// test_etcp_router.c — Unit test for etcp_router: bind/unbind, send, loopback, forward, reply
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "../lib/platform_compat.h"
#include "test_utils.h"
#ifdef _WIN32
#include <windows.h>
#include <direct.h>
#else
#include <unistd.h>
#endif
#include <time.h>
#include <sys/stat.h>

#include "../src/etcp.h"
#include "../src/etcp_connections.h"
#include "../src/etcp_api.h"
#include "../src/etcp_router.h"
#include "../src/config_parser.h"
#include "../src/utun_instance.h"
#include "../src/routing.h"
#include "../src/tun_if.h"
#include "../src/secure_channel.h"
#include "../src/pkt_normalizer.h"
#include "../src/route_bgp.h"
#include "../lib/u_async.h"
#include "../lib/ll_queue.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"

#define TEST_SVC_ID   0x01
#define TEST_TIMEOUT_MS 5000
#define TOTAL_PACKETS 20
#define MAX_PAYLOAD   200

static char temp_dir[] = "/tmp/utun_test_XXXXXX";
static char server_conf[256], client_conf[256];
static struct UTUN_INSTANCE* srv = NULL;
static struct UTUN_INSTANCE* cli = NULL;
static struct UASYNC* ua = NULL;

static int g_ok = 0;
static int g_test_done = 0;
static int g_phase = 0;  // 0=wait_conn, 1=fwd, 2=reply
static void* g_mon_id = NULL;

// Server config — node 0x1111...
static const char* srv_cfg =
    "[global]\n"
    "my_node_id=0x1111111111111111\n"
    "my_private_key=67b705a92b41bcaae105af2d6a17743faa7b26ccebba8b3b9b0af05e9cd1d5fb\n"
    "my_public_key=1c55e4ccae7c4470707759086738b10681bf88b81f198cc2ab54a647d1556e17c65e6b1833e0c771e5a39382c03067c388915a4c732191bc130480f20f8e00b9\n"
    "tun_ip=10.99.0.1/24\n"
    "tun_ifname=tun99\n"
    "[server: s1]\n"
    "addr=127.0.0.1:9041\n"
    "type=public\n"
    "[allowed_keys]\n"
    "allow_all=1\n";

// Client config — node 0x2222...
static const char* cli_cfg =
    "[global]\n"
    "my_node_id=0x2222222222222222\n"
    "my_private_key=4813d31d28b7e9829247f488c6be7672f2bdf61b2508333128e386d1759afed2\n"
    "my_public_key=c594f33c91f3a2222795c2c110c527bf214ad1009197ce14556cb13df3c461b3c373bed8f205a8dd1fc0c364f90bf471d7c6f5db49564c33e4235d268569ac71\n"
    "tun_ip=10.99.0.2/24\n"
    "tun_ifname=tun98\n"
    "[server: s1]\n"
    "addr=127.0.0.1:9042\n"
    "type=public\n"
    "[client: c1]\n"
    "keepalive=1\n"
    "peer_public_key=1c55e4ccae7c4470707759086738b10681bf88b81f198cc2ab54a647d1556e17c65e6b1833e0c771e5a39382c03067c388915a4c732191bc130480f20f8e00b9\n"
    "link=s1:127.0.0.1:9041\n";

// ======================== Test state ========================
static int       fwd_sent = 0, fwd_rcvd = 0;
static int       reply_sent = 0, reply_rcvd = 0;
static uint8_t   expected_data[MAX_PAYLOAD];
static uint64_t  server_node_id = 0x1111111111111111ULL;
static uint64_t  client_node_id = 0x2222222222222222ULL;

// ======================== Helpers ========================
static int conn_established(struct UTUN_INSTANCE* inst) {
    if (!inst) return 0;
    for (struct ETCP_CONN* c = inst->connections; c; c = c->next) {
        for (struct ETCP_LINK* l = c->links; l; l = l->next) {
            if (l->initialized && c->crypto_ctx.initialized) {
                int ok = 0;
                for (int i = 0; i < SC_SESSION_KEY_SIZE; i++) if (c->crypto_ctx.session_key[i] != 0) { ok = 1; break; }
                if (ok) return 1;
            }
        }
    }
    return 0;
}

static struct ETCP_CONN* first_conn(struct UTUN_INSTANCE* inst) {
    return inst ? inst->connections : NULL;
}

// ======================== Server handler ========================
static void srv_handler(struct ETCP_CONN* conn, struct ll_entry* entry) {
    (void)conn;
    if (!entry || !entry->dgram || entry->len < 10) { // svc_id(1) + subcmd(1) + seq(4) + data_len(4)
        if (entry) { queue_entry_free(entry); queue_dgram_free(entry); }
        return;
    }
    uint8_t  subcmd   = entry->dgram[1];
    uint32_t seq      = 0; memcpy(&seq, entry->dgram + 2, 4);
    uint32_t data_len = 0; memcpy(&data_len, entry->dgram + 6, 4);
    uint8_t* payload  = entry->dgram + 10;

    if (subcmd == 0x01) { // DATA
        if (seq != (uint32_t)fwd_rcvd) {
            printf("[FAIL] server: seq mismatch expected=%u got=%u\n", (uint32_t)fwd_rcvd, seq);
            g_test_done = -1;
        } else if (data_len > 0 && memcmp(payload, expected_data, data_len) != 0) {
            printf("[FAIL] server: data mismatch at seq=%u\n", seq);
            g_test_done = -1;
        } else {
            fwd_rcvd++;
        }
        queue_entry_free(entry); queue_dgram_free(entry);

        // Send reply back
        if (!g_test_done && reply_sent < TOTAL_PACKETS) {
            uint8_t buf[10 + MAX_PAYLOAD];
            buf[0] = TEST_SVC_ID;
            buf[1] = 0x02; // REPLY subcmd
            uint32_t rseq = reply_sent;
            memcpy(buf + 2, &rseq, 4);
            memcpy(buf + 6, &data_len, 4);
            memcpy(buf + 10, expected_data, data_len);
            struct ll_entry* re = queue_entry_new(0);
            if (re) { re->dgram = u_malloc(10 + data_len); memcpy(re->dgram, buf, 10 + data_len); re->len = 10 + data_len;
                etcp_route_send(srv, client_node_id, re, 0); reply_sent++; }
        }
    } else {
        queue_entry_free(entry); queue_dgram_free(entry);
    }
}

// ======================== Client handler ========================
static void cli_handler(struct ETCP_CONN* conn, struct ll_entry* entry) {
    (void)conn;
    if (!entry || !entry->dgram || entry->len < 10) {
        if (entry) { queue_entry_free(entry); queue_dgram_free(entry); }
        return;
    }
    uint8_t  subcmd   = entry->dgram[1];
    uint32_t seq      = 0; memcpy(&seq, entry->dgram + 2, 4);
    uint32_t data_len = 0; memcpy(&data_len, entry->dgram + 6, 4);
    uint8_t* payload  = entry->dgram + 10;

    if (subcmd == 0x02) { // REPLY
        if (seq != (uint32_t)reply_rcvd) {
            printf("[FAIL] client: reply seq mismatch expected=%u got=%u\n", (uint32_t)reply_rcvd, seq);
            g_test_done = -1;
        } else if (data_len > 0 && memcmp(payload, expected_data, data_len) != 0) {
            printf("[FAIL] client: reply data mismatch at seq=%u\n", seq);
            g_test_done = -1;
        } else {
            reply_rcvd++;
        }
    }
    queue_entry_free(entry); queue_dgram_free(entry);
}

// ======================== Loopback test (no ETCP) ========================
static int   loop_rcvd = 0;
static int   loop_ok   = 0;
static void  loop_handler(struct ETCP_CONN* conn, struct ll_entry* entry) {
    (void)conn;
    if (entry && entry->dgram && entry->len >= 6) {
        loop_rcvd++;
        if (entry->dgram[1] == 0xAA && entry->dgram[2] == 0xBB && entry->dgram[3] == 0xCC) loop_ok = 1;
    }
    if (entry) { queue_entry_free(entry); queue_dgram_free(entry); }
}

static int test_loopback(void) {
    loop_rcvd = loop_ok = 0;
    etcp_router_bind(srv, 0xF0, loop_handler);
    uint8_t data[6] = { 0xF0, 0xAA, 0xBB, 0xCC, 0x00, 0x00 };
    struct ll_entry* e = queue_entry_new(0);
    e->dgram = u_malloc(6); memcpy(e->dgram, data, 6); e->len = 6;
    etcp_route_send(srv, srv->node_id, e, 0); // loopback
    etcp_router_unbind(srv, 0xF0);
    if (loop_rcvd != 1 || !loop_ok) {
        printf("[FAIL] loopback: rcvd=%d ok=%d\n", loop_rcvd, loop_ok);
        return 1;
    }
    printf("  loopback: OK\n");
    return 0;
}

// ======================== Test API bind/unbind ========================
static int test_api(void) {
    if (etcp_router_bind(srv, 0xEE, loop_handler) != 0) { printf("[FAIL] bind\n"); return 1; }
    if (etcp_router_bind(srv, 0xEE, loop_handler) != 0) { printf("[FAIL] rebind (overwrite)\n"); etcp_router_unbind(srv, 0xEE); return 1; }
    if (etcp_router_unbind(srv, 0xEE) != 0) { printf("[FAIL] unbind\n"); return 1; }
    if (etcp_router_unbind(srv, 0xEE) == 0) { printf("[FAIL] double unbind should fail\n"); return 1; }
    // Invalid args:
    if (etcp_router_bind(NULL, 0, NULL) == 0) { printf("[FAIL] bind null\n"); return 1; }
    if (etcp_router_bind(srv, 0xEE, NULL) == 0) { printf("[FAIL] bind null cb\n"); return 1; }
    printf("  api bind/unbind: OK\n");
    return 0;
}

// ======================== Monitor & sender ========================
static uint64_t dedup_bgp_log = 0;
static void monitor(void* arg) {
    (void)arg;
    if (g_test_done) { g_mon_id = NULL; return; }

    static int conn_ok = 0, conn_delay = 0;
    if (!conn_ok) {
        if (conn_established(srv) && conn_established(cli)) {
            conn_delay++;
            if (conn_delay < 40) { g_mon_id = uasync_set_timeout(ua, 10, NULL, monitor, "mon"); return; }
            conn_ok = 1;
            printf("  connections established\n");
            // Run API tests + loopback
            if (test_api() != 0) { g_test_done = -1; return; }
            if (test_loopback() != 0) { g_test_done = -1; return; }
            // Bind handlers
            etcp_router_bind(srv, TEST_SVC_ID, srv_handler);
            etcp_router_bind(cli, TEST_SVC_ID, cli_handler);
            // Generate shared random data
            for (int i = 0; i < MAX_PAYLOAD; i++) expected_data[i] = (uint8_t)(rand() & 0xFF);
            g_phase = 1;
            printf("  sending %d packets client→server...\n", TOTAL_PACKETS);
            fflush(stdout);
        }
    }

    // Phase 1: send forward
    if (g_phase == 1 && fwd_sent < TOTAL_PACKETS) {
        int data_len = 16 + (fwd_sent % 32);
        uint8_t buf[10 + MAX_PAYLOAD];
        buf[0] = TEST_SVC_ID;
        buf[1] = 0x01; // DATA subcmd
        uint32_t seq = fwd_sent;
        memcpy(buf + 2, &seq, 4);
        memcpy(buf + 6, &data_len, 4);
        memcpy(buf + 10, expected_data, data_len);
        struct ll_entry* e = queue_entry_new(0);
        if (e) { e->dgram = u_malloc(10 + data_len); memcpy(e->dgram, buf, 10 + data_len); e->len = 10 + data_len;
            if (etcp_route_send(cli, server_node_id, e, 0) == 0) fwd_sent++;
            else { queue_entry_free(e); queue_dgram_free(e); }
        }
    }

    if (g_phase == 1 && fwd_sent >= TOTAL_PACKETS && fwd_rcvd >= TOTAL_PACKETS && reply_sent >= TOTAL_PACKETS) {
        g_phase = 2;
        printf("  forward phase done: sent=%d rcvd=%d reply_sent=%d\n", fwd_sent, fwd_rcvd, reply_sent);
    }

    if (g_phase == 2 && reply_rcvd >= TOTAL_PACKETS) {
        printf("  reply phase done: rcvd=%d\n", reply_rcvd);
        g_test_done = 1;
        return;
    }

    g_mon_id = uasync_set_timeout(ua, 10, NULL, monitor, "mon");
}

static void timeout(void* arg) {
    (void)arg;
    if (!g_test_done) { printf("[FAIL] timeout: sent=%d rcvd=%d reply_sent=%d reply_rcvd=%d\n",
                                fwd_sent, fwd_rcvd, reply_sent, reply_rcvd);
        g_test_done = -1; }
    if (g_mon_id) { uasync_cancel_timeout(ua, g_mon_id); g_mon_id = NULL; }
}

// ======================== Main ========================
int main(void) {
    if (test_mkdtemp(temp_dir) != 0) { fprintf(stderr, "mkdtemp fail\n"); return 1; }
    snprintf(server_conf, sizeof(server_conf), "%s/server.conf", temp_dir);
    snprintf(client_conf, sizeof(client_conf), "%s/client.conf", temp_dir);

    FILE* f = fopen(server_conf, "w");
    if (!f) { fprintf(stderr, "fopen fail\n"); test_rmdir(temp_dir); return 1; }
    fprintf(f, "%s", srv_cfg); fclose(f);
    f = fopen(client_conf, "w");
    if (!f) { fprintf(stderr, "fopen fail\n"); test_unlink(server_conf); test_rmdir(temp_dir); return 1; }
    fprintf(f, "%s", cli_cfg); fclose(f);

    printf("=== test_etcp_router ===\n");

    debug_config_init();
    debug_set_level(DEBUG_LEVEL_ERROR);
    debug_set_categories(DEBUG_CATEGORY_ALL);
    utun_instance_set_tun_init_enabled(0);
    srand((unsigned)time(NULL));

    ua = uasync_create();
    if (!ua) { printf("[FAIL] uasync_create\n"); goto done; }

    srv = utun_instance_create(ua, server_conf);
    if (!srv || utun_instance_init(srv) < 0) { printf("[FAIL] server create\n"); goto done; }

    cli = utun_instance_create(ua, client_conf);
    if (!cli || utun_instance_init(cli) < 0) { printf("[FAIL] client create\n"); goto done; }

    g_mon_id = uasync_set_timeout(ua, 100, NULL, monitor, "mon");
    void* to_id = uasync_set_timeout(ua, TEST_TIMEOUT_MS * 10, NULL, timeout, "to");

    while (!g_test_done) uasync_poll(ua, 100);

    if (to_id) uasync_cancel_timeout(ua, to_id);
    if (g_mon_id) uasync_cancel_timeout(ua, g_mon_id);

    if (g_test_done == 1) {
        // Check all statistics
        if (fwd_sent == TOTAL_PACKETS && fwd_rcvd == TOTAL_PACKETS &&
            reply_sent == TOTAL_PACKETS && reply_rcvd == TOTAL_PACKETS) {
            printf("[PASS] test_etcp_router — %d packets forwarded, %d replied\n", fwd_rcvd, reply_rcvd);
            g_ok = 1;
        } else {
            printf("[FAIL] incomplete: fwd_sent=%d fwd_rcvd=%d reply_sent=%d reply_rcvd=%d\n",
                   fwd_sent, fwd_rcvd, reply_sent, reply_rcvd);
        }
    } else {
        printf("[FAIL] test did not complete\n");
    }

    etcp_router_unbind(srv, TEST_SVC_ID);
    etcp_router_unbind(cli, TEST_SVC_ID);

done:
    if (srv) { srv->running = 0; utun_instance_destroy(srv); }
    if (cli) { cli->running = 0; utun_instance_destroy(cli); }
    if (ua) uasync_destroy(ua, 0);
    test_unlink(server_conf); test_unlink(client_conf); test_rmdir(temp_dir);
    return g_ok ? 0 : 1;
}