/**
 * @file test_nat_engine.c
 * @brief Unit-тесты чистого NAT engine (eim_nat.c/h)
 *
 * Тестирует eim_nat_init_ctx, eim_nat_egress, eim_nat_ingress,
 * eim_nat_add_forward, eim_nat_destroy_ctx.
 *
 * Без TUN, без ETCP, без UTUN_INSTANCE — только crafted IP пакеты.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "../lib/debug_config.h"
#include "../src/eim_nat.h"
#include "etcp.h"
#include "../src/config_parser.h"

#ifdef ENABLE_STATIC_ASSERT
static_assert(sizeof(struct eim_nat_entry) <= 64, "entry size ok");
#endif

static struct {
    int run, passed, failed;
} stats = {0};

#define TEST(name) do { \
    printf("TEST: %-50s ", name); fflush(stdout); \
    stats.run++; \
} while(0)

#define PASS()      do { puts("PASS"); stats.passed++; } while(0)
#define FAIL(msg)   do { printf("FAIL: %s\n", msg); stats.failed++; } while(0)

#define ASSERT(c, m) do { if (!(c)) { FAIL(m); return; } } while(0)
#define ASSERT_EQ(a,b,m) ASSERT((a)==(b),m)

/* ================================================================
 * Helpers: build IP packets (uses htobe32/htobe16 + memcpy: network byte order in wire)
 * ================================================================ */

static void build_ip_hdr(uint8_t* buf, uint8_t proto, uint32_t src_host, uint32_t dst_host, uint16_t total_len) {
    buf[0] = 0x45;
    buf[1] = 0x00;
    uint16_t n = htobe16(total_len);  memcpy(buf + 2, &n, 2);
    buf[4] = 0x12; buf[5] = 0x34;
    memset(buf + 6, 0, 2);
    buf[8] = 64;
    buf[9] = proto;
    memset(buf + 10, 0, 2); // checksum slot = 0 for now
    uint32_t src_net = htobe32(src_host);
    uint32_t dst_net = htobe32(dst_host);
    memcpy(buf + 12, &src_net, 4);
    memcpy(buf + 16, &dst_net, 4);
}

static void compute_ip_checksum(uint8_t* ip) {
    uint32_t sum = 0;
    for (int i = 0; i < 10; i++) { uint16_t w; memcpy(&w, ip + i*2, 2); sum += w; }
    sum = (sum & 0xFFFF) + (sum >> 16);
    sum += (sum >> 16);
    uint16_t c = (uint16_t)(~sum);
    memcpy(ip + 10, &c, 2);
}

static int verify_ip_checksum(uint8_t* ip) {
    uint32_t sum = 0;
    for (int i = 0; i < 10; i++) { uint16_t w; memcpy(&w, ip + i*2, 2); sum += w; }
    sum = (sum & 0xFFFF) + (sum >> 16);
    sum += (sum >> 16);
    return (uint16_t)(~sum) == 0;
}

#define TEST_IP_SRC_HOST  0x0A000002  // 10.0.0.2
#define TEST_IP_DST_HOST  0x08080808  // 8.8.8.8
#define TEST_GW_HOST      0x0A000001  // 10.0.0.1  - gateway NAT IP
#define TEST_PORT_START   10000
#define TEST_PORT_END     20000
#define TEST_SRC_PORT     40000
#define TEST_DST_PORT     53
#define TEST_PAYLOAD_LEN  14

static struct ETCP_CONN mock_conn;

static int mock_conn_initialized = 0;
static struct ETCP_CONN* get_mock_conn(void) {
    if (!mock_conn_initialized) {
        memset(&mock_conn, 0, sizeof(mock_conn));
        mock_conn.peer_node_id = 0xAAAA000000000001ULL;
        mock_conn_initialized = 1;
    }
    return &mock_conn;
}

static struct global_config make_global_config(void) {
    struct global_config g;
    memset(&g, 0, sizeof(g));
    g.nat_enabled = 1;
    g.nat_port_start = TEST_PORT_START;
    g.nat_port_end = TEST_PORT_END;
    g.nat_tun_ip.family = AF_INET;
    g.nat_tun_ip.addr.v4.s_addr = htobe32(TEST_GW_HOST);
    return g;
}

/* ================================================================
 * Test 1: Init / Destroy
 * ================================================================ */
static void test_init_destroy(void) {
    TEST("init_destroy_normal");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        int r = eim_nat_init_ctx(&ctx, &g);
        ASSERT_EQ(r, 0, "init returns 0");
        ASSERT(ctx.initialized == 1, "ctx.initialized=1");
        ASSERT(ctx.gateway_ip == TEST_GW_HOST, "gateway_ip correct");
        ASSERT(ctx.port_start == TEST_PORT_START, "port_start correct");
        ASSERT(ctx.port_end == TEST_PORT_END, "port_end correct");
        ASSERT(ctx.table != NULL, "table allocated");
        eim_nat_destroy_ctx(&ctx);
        ASSERT(ctx.initialized == 0, "destroy clears initialized");
        ASSERT(ctx.table == NULL, "destroy frees table");
    }
    PASS();

    TEST("init_null_ctx");
    {
        struct global_config g = make_global_config();
        int r = eim_nat_init_ctx(NULL, &g);
        ASSERT_EQ(r, -1, "returns -1");
    }
    PASS();

    TEST("init_null_config");
    {
        struct eim_nat_ctx ctx;
        int r = eim_nat_init_ctx(&ctx, NULL);
        ASSERT_EQ(r, -1, "returns -1");
    }
    PASS();

    TEST("init_nat_disabled");
    {
        struct global_config g = make_global_config();
        g.nat_enabled = 0;
        struct eim_nat_ctx ctx;
        int r = eim_nat_init_ctx(&ctx, &g);
        ASSERT_EQ(r, 0, "returns 0");
        ASSERT(ctx.initialized == 0, "not initialized");
        ASSERT(ctx.table == NULL, "no table");
    }
    PASS();

    TEST("init_bad_port_range");
    {
        struct global_config g = make_global_config();
        g.nat_port_start = 20000; g.nat_port_end = 10000;
        struct eim_nat_ctx ctx;
        int r = eim_nat_init_ctx(&ctx, &g);
        ASSERT_EQ(r, -1, "returns -1");
    }
    PASS();

    TEST("destroy_twice_safe");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);
        eim_nat_destroy_ctx(&ctx);
        eim_nat_destroy_ctx(&ctx); // second call should be no-op
    }
    PASS();
}

/* ================================================================
 * Test 2: Port Allocation
 * ================================================================ */

static uint8_t* make_udp_pkt(uint32_t src_host, uint16_t src_port_host, uint32_t dst_host, uint16_t dst_port_host) {
    const size_t ip_udp_len = 20 + 8 + TEST_PAYLOAD_LEN;
    uint8_t* pkt = calloc(1, ip_udp_len);
    build_ip_hdr(pkt, IPPROTO_UDP_UINT8, src_host, dst_host, ip_udp_len);
    // UDP header
    uint16_t sp = htobe16(src_port_host), dp = htobe16(dst_port_host);
    memcpy(pkt + 20, &sp, 2);     // src port
    memcpy(pkt + 22, &dp, 2);     // dst port
    uint16_t udp_len = htobe16(8 + TEST_PAYLOAD_LEN);
    memcpy(pkt + 24, &udp_len, 2); // length
    memset(pkt + 26, 0, 2);       // checksum = 0 (no UDP csum)
    memset(pkt + 28, 0xAB, TEST_PAYLOAD_LEN); // payload
    compute_ip_checksum(pkt);
    return pkt;
}

static void test_port_alloc(void) {
    TEST("port_alloc_sequential");
    {
        struct global_config g = make_global_config();
        g.nat_port_start = 10000; g.nat_port_end = 10005;
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        // Allocate 3 ports (different internal src ports)
        for (int i = 0; i < 3; i++) {
            uint8_t* pkt = make_udp_pkt(TEST_IP_SRC_HOST, 50000 + i, TEST_IP_DST_HOST, 53);
            int r = eim_nat_egress(&ctx, pkt, 20 + 8 + TEST_PAYLOAD_LEN, 1, get_mock_conn());
            ASSERT_EQ(r, 0, "egress ok");
            // Check port was allocated from table index
            struct eim_nat_entry* e = &ctx.table[10000 + i];
            ASSERT(e->state == EIM_NAT_ENTRY_ACTIVE, "entry active");
            ASSERT_EQ(e->internal_port, htobe16(50000 + i), "internal port stored");
            free(pkt);
        }
        ASSERT(ctx.next_port == 10003, "next_port advanced");
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();

    TEST("port_alloc_wraparound");
    {
        struct global_config g = make_global_config();
        g.nat_port_start = 10000; g.nat_port_end = 10001; // 2 ports
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        // Fill first entry: egress with port not yet seen
        uint8_t* p1 = make_udp_pkt(TEST_IP_SRC_HOST, 50001, TEST_IP_DST_HOST, 53);
        int r = eim_nat_egress(&ctx, p1, 20 + 8 + TEST_PAYLOAD_LEN, 1, get_mock_conn());
        ASSERT_EQ(r, 0, "first egress ok");
        ASSERT(ctx.table[10000].state == EIM_NAT_ENTRY_ACTIVE, "port 10000 used");
        free(p1);

        // Release port 10000 manually
        ctx.table[10000].state = EIM_NAT_ENTRY_FREE;

        // Now alloc should reuse port 10000 (next_port is at 10001, but it wraps around)
        // Actually next_port is 10001 after first alloc. port 10000 is free, but alloc starts from next_port.
        // Let me check the algorithm: it scans from next_port forward. If 10000 is free but 10001 is not current, 
        // it will allocate 10001. But if we free 10000, the scan from 10001 will bypass it.
        // Actually eim_nat_alloc_port starts from ctx->next_port, not from port_start. 
        // After first alloc, next_port=10001. Free 10000. 
        // Now alloc starts from 10001 - it's free (we only allocated 10000, then freed it. next_port was incremented to 10001).
        // Wait, first alloc: port 10000 → next_port becomes 10001 (since 10001 <= port_end).
        // Then we free 10000. 
        // Now alloc starts from 10001. It's free. So it allocates 10001.
        // Then next_port becomes 10002 → > 10001 → wraps to 10000. Now it allocates 10000 since it's free.

        // Allocate second - gets 10001
        uint8_t* p2 = make_udp_pkt(TEST_IP_SRC_HOST, 50002, TEST_IP_DST_HOST, 53);
        r = eim_nat_egress(&ctx, p2, 20 + 8 + TEST_PAYLOAD_LEN, 1, get_mock_conn());
        ASSERT_EQ(r, 0, "second egress ok");
        ASSERT(ctx.table[10001].state == EIM_NAT_ENTRY_ACTIVE, "port 10001 used");
        ASSERT_EQ(ctx.next_port, 10000, "next_port wrapped to 10000");
        free(p2);

        // Third - wraps and gets 10000 (freed)
        uint8_t* p3 = make_udp_pkt(TEST_IP_SRC_HOST, 50003, TEST_IP_DST_HOST, 53);
        r = eim_nat_egress(&ctx, p3, 20 + 8 + TEST_PAYLOAD_LEN, 1, get_mock_conn());
        ASSERT_EQ(r, 0, "third egress ok after wrap");
        ASSERT(ctx.table[10000].state == EIM_NAT_ENTRY_ACTIVE, "port 10000 reused");
        ASSERT_EQ(ctx.table[10000].internal_port, htobe16(50003), "new entry for reused port");
        free(p3);

        eim_nat_destroy_ctx(&ctx);
    }
    PASS();

    TEST("port_exhaustion");
    {
        struct global_config g = make_global_config();
        g.nat_port_start = 10000; g.nat_port_end = 10001; // 2 ports
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        // Fill both ports with different flows
        uint8_t* p1 = make_udp_pkt(TEST_IP_SRC_HOST, 50001, TEST_IP_DST_HOST, TEST_DST_PORT);
        int r = eim_nat_egress(&ctx, p1, 20 + 8 + TEST_PAYLOAD_LEN, 1, get_mock_conn());
        ASSERT_EQ(r, 0, "first ok"); free(p1);

        uint8_t* p2 = make_udp_pkt(0x0A000003, 50002, TEST_IP_DST_HOST, TEST_DST_PORT);
        r = eim_nat_egress(&ctx, p2, 20 + 8 + TEST_PAYLOAD_LEN, 2, get_mock_conn());
        ASSERT_EQ(r, 0, "second ok"); free(p2);

        // Third with different (ip,port) → alloc fails
        uint8_t* p3 = make_udp_pkt(0x0A000004, 50003, TEST_IP_DST_HOST, TEST_DST_PORT);
        r = eim_nat_egress(&ctx, p3, 20 + 8 + TEST_PAYLOAD_LEN, 3, get_mock_conn());
        ASSERT_EQ(r, -1, "fails on exhaustion"); free(p3);

        // Same flow as first → reuses entry
        uint8_t* p4 = make_udp_pkt(TEST_IP_SRC_HOST, 50001, TEST_IP_DST_HOST, TEST_DST_PORT);
        r = eim_nat_egress(&ctx, p4, 20 + 8 + TEST_PAYLOAD_LEN, 1, get_mock_conn());
        ASSERT_EQ(r, 0, "same flow reuses entry"); free(p4);

        eim_nat_destroy_ctx(&ctx);
    }
    PASS();
}

/* ================================================================
 * Test 3: Egress NAT — UDP
 * ================================================================ */
static void test_egress_udp(void) {
    TEST("egress_udp_basic");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        uint8_t* pkt = make_udp_pkt(TEST_IP_SRC_HOST, TEST_SRC_PORT, TEST_IP_DST_HOST, TEST_DST_PORT);
        // Save original dst IP/port (should not be changed by egress)
        uint32_t orig_dst_ip_net; memcpy(&orig_dst_ip_net, pkt + 16, 4);
        uint16_t orig_dst_port_net; memcpy(&orig_dst_port_net, pkt + 22, 2);

        int r = eim_nat_egress(&ctx, pkt, 20 + 8 + TEST_PAYLOAD_LEN, 0x5555, get_mock_conn());
        ASSERT_EQ(r, 0, "egress returns 0");

        // Check src IP = gateway
        uint32_t new_src_ip_net; memcpy(&new_src_ip_net, pkt + 12, 4);
        ASSERT_EQ(be32toh(new_src_ip_net), TEST_GW_HOST, "src IP = gateway");
        // Check dst IP unchanged
        uint32_t dst_ip_net; memcpy(&dst_ip_net, pkt + 16, 4);
        ASSERT_EQ(dst_ip_net, orig_dst_ip_net, "dst IP unchanged");
        // Check src port changed
        uint16_t new_src_port_net; memcpy(&new_src_port_net, pkt + 20, 2);
        ASSERT(be16toh(new_src_port_net) == TEST_PORT_START, "src port = port_start");
        // Check dst port unchanged
        uint16_t dst_port_net; memcpy(&dst_port_net, pkt + 22, 2);
        ASSERT_EQ(dst_port_net, orig_dst_port_net, "dst port unchanged");
        // Check IP checksum valid
        ASSERT(verify_ip_checksum(pkt) == 1, "IP checksum valid");
        // Check NAT entry
        struct eim_nat_entry* e = &ctx.table[TEST_PORT_START];
        ASSERT_EQ(e->state, EIM_NAT_ENTRY_ACTIVE, "entry active");
        ASSERT_EQ(e->internal_ip, TEST_IP_SRC_HOST, "internal_ip stored");
        ASSERT_EQ(e->internal_port, htobe16(TEST_SRC_PORT), "internal_port stored");
        ASSERT_EQ(e->proto, IPPROTO_UDP_UINT8, "proto=UDP");
        ASSERT_EQ(e->src_node_id, 0x5555ULL, "src_node_id stored");
        ASSERT(e->src_conn == get_mock_conn(), "src_conn stored");

        free(pkt);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();
}

/* ================================================================
 * Test 4: Egress NAT — TCP
 * ================================================================ */
static uint8_t* make_tcp_pkt(uint32_t src_host, uint16_t src_port_host, uint32_t dst_host, uint16_t dst_port_host) {
    const size_t ip_tcp_len = 20 + 20 + TEST_PAYLOAD_LEN; // 20 TCP hdr min
    uint8_t* pkt = calloc(1, ip_tcp_len);
    build_ip_hdr(pkt, IPPROTO_TCP_UINT8, src_host, dst_host, ip_tcp_len);
    uint16_t sp = htobe16(src_port_host), dp = htobe16(dst_port_host);
    memcpy(pkt + 20, &sp, 2);
    memcpy(pkt + 22, &dp, 2);
    // seq, ack, offset+flags, window
    pkt[32] = 0x50; // offset=5 (20 bytes), flags=0
    // checksum
    memset(pkt + 36, 0, 2);
    memset(pkt + 40, 0xCC, TEST_PAYLOAD_LEN);
    compute_ip_checksum(pkt);
    return pkt;
}

static void test_egress_tcp(void) {
    TEST("egress_tcp_basic");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        uint8_t* pkt = make_tcp_pkt(TEST_IP_SRC_HOST, TEST_SRC_PORT, TEST_IP_DST_HOST, 80);
        int r = eim_nat_egress(&ctx, pkt, 20 + 20 + TEST_PAYLOAD_LEN, 0x1111, get_mock_conn());
        ASSERT_EQ(r, 0, "egress TCP ok");
        // Check src IP = gateway
        uint32_t new_src; memcpy(&new_src, pkt + 12, 4);
        ASSERT_EQ(be32toh(new_src), TEST_GW_HOST, "src IP=gw");
        // Check src port
        uint16_t new_sport; memcpy(&new_sport, pkt + 20, 2);
        ASSERT_EQ(be16toh(new_sport), TEST_PORT_START, "src port=start");
        // IP checksum valid
        ASSERT(verify_ip_checksum(pkt) == 1, "IP checksum valid");
        // Entry proto = TCP
        ASSERT_EQ(ctx.table[TEST_PORT_START].proto, IPPROTO_TCP_UINT8, "proto=TCP");
        free(pkt);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();
}

/* ================================================================
 * Test 5: Egress ICMP Echo
 * ================================================================ */
static uint8_t* make_icmp_echo_pkt(uint32_t src_host, uint32_t dst_host, uint8_t icmp_type, uint16_t id_host, uint16_t seq_host) {
    const size_t ip_icmp_len = 20 + 8 + TEST_PAYLOAD_LEN;
    uint8_t* pkt = calloc(1, ip_icmp_len);
    build_ip_hdr(pkt, IPPROTO_ICMP_UINT8, src_host, dst_host, ip_icmp_len);
    pkt[20] = icmp_type; // type
    pkt[21] = 0x00;      // code
    // checksum = 0 for now
    memset(pkt + 22, 0, 2);
    uint16_t id_n = htobe16(id_host), seq_n = htobe16(seq_host);
    memcpy(pkt + 24, &id_n, 2);
    memcpy(pkt + 26, &seq_n, 2);
    memset(pkt + 28, 0xDD, TEST_PAYLOAD_LEN);
    compute_ip_checksum(pkt);
    return pkt;
}

static void test_egress_icmp(void) {
    TEST("egress_icmp_echo_request");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        uint16_t icmp_id = 0x1234; // host order
        uint8_t* pkt = make_icmp_echo_pkt(TEST_IP_SRC_HOST, TEST_IP_DST_HOST, 8, icmp_id, 1);
        int r = eim_nat_egress(&ctx, pkt, 20 + 8 + TEST_PAYLOAD_LEN, 0x9999, get_mock_conn());
        ASSERT_EQ(r, 0, "egress ICMP echo ok");
        // ICMP ID should be overwritten with allocated port
        uint16_t new_id_net; memcpy(&new_id_net, pkt + 24, 2);
        ASSERT_EQ(be16toh(new_id_net), TEST_PORT_START, "ICMP ID = allocated port");
        // IP checksum valid
        ASSERT(verify_ip_checksum(pkt) == 1, "IP checksum valid");
        // Entry proto = ICMP
        ASSERT_EQ(ctx.table[TEST_PORT_START].proto, IPPROTO_ICMP_UINT8, "proto=ICMP");
        ASSERT_EQ(ctx.table[TEST_PORT_START].internal_port, htobe16(icmp_id), "internal port = ICMP ID");
        free(pkt);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();

    TEST("egress_icmp_error_no_rewrite");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        // ICMP Dest Unreachable (type 3) — no echo, no id rewrite, no entry
        uint8_t* pkt = make_icmp_echo_pkt(TEST_IP_SRC_HOST, TEST_IP_DST_HOST, 3, 0x1234, 1);
        // Save original data for comparison
        uint8_t backup[20 + 8 + TEST_PAYLOAD_LEN];
        memcpy(backup, pkt, sizeof(backup));
        int r = eim_nat_egress(&ctx, pkt, 20 + 8 + TEST_PAYLOAD_LEN, 0x9999, get_mock_conn());
        ASSERT_EQ(r, 0, "returns 0 (not -1)");
        // Check no entry created
        ASSERT(ctx.table[TEST_PORT_START].state == EIM_NAT_ENTRY_FREE, "no entry created for ICMP error");
        // Packet should be unchanged (non-echo ICMP bypasses)
        ASSERT(memcmp(backup, pkt, sizeof(backup)) == 0, "packet unchanged");
        free(pkt);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();
}

/* ================================================================
 * Test 6: Egress fragments
 * ================================================================ */
static void test_egress_fragments(void) {
    TEST("egress_fragment_mf_no_off");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        uint8_t* pkt = make_udp_pkt(TEST_IP_SRC_HOST, TEST_SRC_PORT, TEST_IP_DST_HOST, TEST_DST_PORT);
        // Set MF=1, offset=0
        pkt[6] = 0x20; pkt[7] = 0x00;
        // Recompute checksum with new frag field
        compute_ip_checksum(pkt);

        int r = eim_nat_egress(&ctx, pkt, 20 + 8 + TEST_PAYLOAD_LEN, 0x1111, get_mock_conn());
        // MF bit set → bypassed (returns 0, no allocation)
        ASSERT_EQ(r, 0, "egress fragment MF=1 bypassed");
        ASSERT(ctx.table[TEST_PORT_START].state == EIM_NAT_ENTRY_FREE, "no entry for fragment");
        free(pkt);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();

    TEST("egress_fragment_nonzero_offset");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        uint8_t* pkt = make_udp_pkt(TEST_IP_SRC_HOST, TEST_SRC_PORT, TEST_IP_DST_HOST, TEST_DST_PORT);
        // offset = 185 (in 8-byte units) → 0x00B9 network order
        pkt[6] = 0x00; pkt[7] = 0xB9;
        compute_ip_checksum(pkt);

        int r = eim_nat_egress(&ctx, pkt, 20 + 8 + TEST_PAYLOAD_LEN, 0x1111, get_mock_conn());
        ASSERT_EQ(r, 0, "egress fragment offset>0 bypassed");
        ASSERT(ctx.table[TEST_PORT_START].state == EIM_NAT_ENTRY_FREE, "no entry");
        free(pkt);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();
}

/* ================================================================
 * Test 7: Egress invalid packets
 * ================================================================ */
static void test_egress_invalid(void) {
    TEST("egress_too_short");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);
        uint8_t buf[10];
        int r = eim_nat_egress(&ctx, buf, 10, 1, get_mock_conn());
        ASSERT_EQ(r, -1, "returns -1");
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();

    TEST("egress_bad_ihl");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);
        uint8_t* pkt = make_udp_pkt(TEST_IP_SRC_HOST, TEST_SRC_PORT, TEST_IP_DST_HOST, TEST_DST_PORT);
        pkt[0] = 0x43; // IHL=3 (invalid, <5)
        int r = eim_nat_egress(&ctx, pkt, 20 + 8 + TEST_PAYLOAD_LEN, 1, get_mock_conn());
        ASSERT_EQ(r, -1, "returns -1 for bad IHL");
        free(pkt);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();

    TEST("egress_not_tcp_udp_icmp");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);
        // Build IP with proto=0x63 (unknown) but pretend it's UDP format
        uint8_t pkt[20 + 8 + TEST_PAYLOAD_LEN];
        build_ip_hdr(pkt, 0x63, TEST_IP_SRC_HOST, TEST_IP_DST_HOST, sizeof(pkt));
        // Fill fake UDP header
        uint16_t sp = htobe16(TEST_SRC_PORT), dp = htobe16(TEST_DST_PORT);
        memcpy(pkt + 20, &sp, 2); memcpy(pkt + 22, &dp, 2);
        compute_ip_checksum(pkt);
        int r = eim_nat_egress(&ctx, pkt, sizeof(pkt), 1, get_mock_conn());
        ASSERT_EQ(r, 0, "unknown proto bypassed (returns 0)");
        ASSERT(ctx.table[TEST_PORT_START].state == EIM_NAT_ENTRY_FREE, "no entry");
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();
}

/* ================================================================
 * Test 8: Ingress NAT — UDP
 * ================================================================ */
static uint8_t* make_respond_pkt(uint32_t src_host, uint16_t src_port_host, uint32_t dst_host, uint16_t dst_port_host) {
    // Build a packet FROM internet TO gateway (this is the response after egress)
    return make_udp_pkt(src_host, src_port_host, dst_host, dst_port_host);
}

static void test_ingress_udp(void) {
    TEST("ingress_udp_normal");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        // First: egress to create entry
        uint8_t* out = make_udp_pkt(TEST_IP_SRC_HOST, TEST_SRC_PORT, TEST_IP_DST_HOST, TEST_DST_PORT);
        eim_nat_egress(&ctx, out, 20 + 8 + TEST_PAYLOAD_LEN, 0x1111, get_mock_conn());
        free(out);

        // Save internal state
        uint32_t internal_ip = ctx.table[TEST_PORT_START].internal_ip;
        uint16_t internal_port_net = ctx.table[TEST_PORT_START].internal_port;

        // Build response: FROM 8.8.8.8:53 TO gateway:port_start
        uint8_t* resp = make_respond_pkt(TEST_IP_DST_HOST, TEST_DST_PORT, TEST_GW_HOST, TEST_PORT_START);
        struct eim_nat_entry* entry = NULL;
        int r = eim_nat_ingress(&ctx, resp, 20 + 8 + TEST_PAYLOAD_LEN, &entry);
        ASSERT_EQ(r, 0, "ingress ok");
        ASSERT(entry != NULL, "entry returned");
        ASSERT(entry == &ctx.table[TEST_PORT_START], "correct entry");
        // Check dst IP → internal IP
        uint32_t new_dst_net; memcpy(&new_dst_net, resp + 16, 4);
        ASSERT_EQ(be32toh(new_dst_net), internal_ip, "dst IP = internal IP");
        // Check dst port → internal port
        uint16_t new_dst_port_net; memcpy(&new_dst_port_net, resp + 22, 2);
        ASSERT_EQ(new_dst_port_net, internal_port_net, "dst port = internal port");
        // Check src IP unchanged
        uint32_t src_net; memcpy(&src_net, resp + 12, 4);
        ASSERT_EQ(be32toh(src_net), TEST_IP_DST_HOST, "src IP unchanged");
        ASSERT(verify_ip_checksum(resp) == 1, "IP checksum valid");

        free(resp);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();

    TEST("ingress_no_entry");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);
        uint8_t* resp = make_respond_pkt(TEST_IP_DST_HOST, 53, TEST_GW_HOST, TEST_PORT_START + 500);
        int r = eim_nat_ingress(&ctx, resp, 20 + 8 + TEST_PAYLOAD_LEN, NULL);
        ASSERT_EQ(r, 0, "returns 0 (no entry → bypass)");
        // Packet unchanged (port not in range anyway? Actually it IS in range but entry is FREE)
        // Wait: port_start+500 = 10500, which IS in range [10000,20000]. Entry state = FREE.
        // Code checks: if entry->state == FREE return 0
        ASSERT(ctx.table[10500].state == EIM_NAT_ENTRY_FREE, "entry is free");
        free(resp);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();

    TEST("ingress_not_for_gateway");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);
        // Packet dst = 8.8.8.8, not gateway → bypass
        uint8_t* resp = make_respond_pkt(TEST_IP_DST_HOST, 53, TEST_IP_DST_HOST, TEST_PORT_START);
        uint8_t backup[20 + 8 + TEST_PAYLOAD_LEN];
        memcpy(backup, resp, sizeof(backup));
        int r = eim_nat_ingress(&ctx, resp, 20 + 8 + TEST_PAYLOAD_LEN, NULL);
        ASSERT_EQ(r, 0, "bypass (dst not gateway)");
        ASSERT(memcmp(backup, resp, sizeof(backup)) == 0, "packet unchanged");
        free(resp);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();
}

/* ================================================================
 * Test 9: Ingress ICMP Echo Reply
 * ================================================================ */
static void test_ingress_icmp(void) {
    TEST("ingress_icmp_echo_reply");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        uint16_t icmp_id = 0x4321;
        // 1. Egress ICMP Echo Request → rewrites ID to allocated port
        uint8_t* req = make_icmp_echo_pkt(TEST_IP_SRC_HOST, TEST_IP_DST_HOST, 8, icmp_id, 1);
        eim_nat_egress(&ctx, req, 20 + 8 + TEST_PAYLOAD_LEN, 0x2222, get_mock_conn());
        free(req);

        // 2. Build ICMP Echo Reply FROM internet TO gateway:port_start
        const size_t len = 20 + 8 + TEST_PAYLOAD_LEN;
        uint8_t* reply = calloc(1, len);
        build_ip_hdr(reply, IPPROTO_ICMP_UINT8, TEST_IP_DST_HOST, TEST_GW_HOST, len);
        reply[20] = 0; // Echo Reply
        reply[21] = 0;
        memset(reply + 22, 0, 2); // checksum
        uint16_t alloc_id_net = htobe16(TEST_PORT_START); // the port allocated by egress
        uint16_t seq = htobe16(1);
        memcpy(reply + 24, &alloc_id_net, 2);
        memcpy(reply + 26, &seq, 2);
        memset(reply + 28, 0xDD, TEST_PAYLOAD_LEN);
        compute_ip_checksum(reply);

        // 3. Ingress → should rewrite ID back to icmp_id
        struct eim_nat_entry* entry = NULL;
        int r = eim_nat_ingress(&ctx, reply, len, &entry);
        ASSERT_EQ(r, 0, "ingress icmp reply ok");
        uint16_t new_id_net; memcpy(&new_id_net, reply + 24, 2);
        ASSERT_EQ(be16toh(new_id_net), icmp_id, "ICMP ID restored to original");
        ASSERT(entry != NULL, "entry returned");
        ASSERT(verify_ip_checksum(reply) == 1, "IP checksum valid");
        free(reply);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();
}

/* ================================================================
 * Test 10: Port Forwarding (static entries)
 * ================================================================ */
static void test_port_forward(void) {
    TEST("add_forward_basic");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        int r = eim_nat_add_forward(&ctx, IPPROTO_TCP_UINT8,
                                    0x0A0000FE, htobe16(8080), // internal 10.0.0.254:8080
                                    10050);
        ASSERT_EQ(r, 0, "add_forward ok");
        ASSERT(ctx.table[10050].state == EIM_NAT_ENTRY_STATIC, "entry static");
        ASSERT_EQ(ctx.table[10050].internal_ip, 0x0A0000FE, "internal_ip");
        ASSERT_EQ(ctx.table[10050].internal_port, htobe16(8080), "internal_port");
        ASSERT_EQ(ctx.table[10050].proto, IPPROTO_TCP_UINT8, "proto=TCP");
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();

    TEST("add_forward_duplicate");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        eim_nat_add_forward(&ctx, IPPROTO_TCP_UINT8, 0x0A0000FE, htobe16(8080), 10050);
        int r = eim_nat_add_forward(&ctx, IPPROTO_UDP_UINT8, 0x0A0000FF, htobe16(9090), 10050);
        ASSERT_EQ(r, -1, "duplicate rejects");
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();

    TEST("ingress_hits_static_entry");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        // Static forward: external port 10050 → internal 10.0.0.254:8080 TCP
        eim_nat_add_forward(&ctx, IPPROTO_TCP_UINT8, 0x0A0000FE, htobe16(8080), 10050);

        // Ingress TCP packet to gateway:10050
        uint8_t* resp = make_tcp_pkt(TEST_IP_DST_HOST, 443, TEST_GW_HOST, 10050);
        struct eim_nat_entry* entry = NULL;
        int r = eim_nat_ingress(&ctx, resp, 20 + 20 + TEST_PAYLOAD_LEN, &entry);
        ASSERT_EQ(r, 0, "ingress static ok");
        ASSERT(entry != NULL, "entry returned");
        ASSERT_EQ(entry->state, EIM_NAT_ENTRY_STATIC, "static entry");
        // Check dst IP → 10.0.0.254
        uint32_t dst_net; memcpy(&dst_net, resp + 16, 4);
        ASSERT_EQ(be32toh(dst_net), 0x0A0000FE, "dst IP = 10.0.0.254");
        // Check dst port → 8080
        uint16_t dst_port; memcpy(&dst_port, resp + 22, 2);
        ASSERT_EQ(dst_port, htobe16(8080), "dst port = 8080");
        ASSERT(verify_ip_checksum(resp) == 1, "IP checksum valid");
        free(resp);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();
}

/* ================================================================
 * Test 11: ICMP non-echo egress (bypass, no modification)
 * ================================================================ */
static void test_bypass_flows(void) {
    TEST("egress_flow_reuse");
    {
        struct global_config g = make_global_config();
        struct eim_nat_ctx ctx;
        eim_nat_init_ctx(&ctx, &g);

        // Same flow (ip:port:proto) twice → same port
        uint8_t* p1 = make_udp_pkt(TEST_IP_SRC_HOST, TEST_SRC_PORT, TEST_IP_DST_HOST, TEST_DST_PORT);
        eim_nat_egress(&ctx, p1, 20 + 8 + TEST_PAYLOAD_LEN, 0x1111, get_mock_conn());
        free(p1);

        uint8_t* p2 = make_udp_pkt(TEST_IP_SRC_HOST, TEST_SRC_PORT, 0x08080404, TEST_DST_PORT);
        eim_nat_egress(&ctx, p2, 20 + 8 + TEST_PAYLOAD_LEN, 0x1111, get_mock_conn());
        // Should reuse same port (matching internal_ip:port:proto)
        uint16_t sp; memcpy(&sp, p2 + 20, 2);
        ASSERT_EQ(be16toh(sp), TEST_PORT_START, "same port reused");
        free(p2);
        eim_nat_destroy_ctx(&ctx);
    }
    PASS();
}

/* ================================================================
 * Main
 * ================================================================ */
int main(void) {
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_ERROR);

    test_init_destroy();
    test_port_alloc();
    test_egress_udp();
    test_egress_tcp();
    test_egress_icmp();
    test_egress_fragments();
    test_egress_invalid();
    test_ingress_udp();
    test_ingress_icmp();
    test_port_forward();
    test_bypass_flows();

    printf("\n=== Results: %d run, %d passed, %d failed ===\n",
           stats.run, stats.passed, stats.failed);
    return stats.failed ? 1 : 0;
}