#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include "../lib/platform_compat.h"
#include "test_utils.h"
#ifndef _WIN32
#include <unistd.h>
#endif

#include "etcp.h"
#include "etcp_connections.h"
#include "etcp_api.h"
#include "../src/config_parser.h"
#include "../src/config_updater.h"
#include "../src/utun_instance.h"
#include "topo_group.h"
#include "topo_node.h"
#include "secure_channel.h"
#include "../lib/u_async.h"
#include "../lib/ll_queue.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"

#define TIMEOUT_TB 300000
#define POLL_MS 5
static uint64_t nid_a = 0;
static uint64_t nid_b = 0;

static struct UTUN_INSTANCE* g_a = NULL, *g_b = NULL;
static struct UASYNC* ua = NULL;
static volatile int result = 0;
static void* ttimer = NULL;
static char tdir[] = "/tmp/utun_ec_XXXXXX";
static char ca[256], cb[256];
static int pa = 0, pb = 0;

static int wf(const char* p, const char* f, ...) {
    va_list ap; FILE* fp = fopen(p, "w"); if (!fp) return -1;
    va_start(ap, f); vfprintf(fp, f, ap); va_end(ap); fclose(fp); return 0;
}
static char* gv(const char* p, const char* k) {
    struct utun_config* c = parse_config(p); if (!c) return NULL;
    char* r = (strcmp(k, "pub") == 0) ? u_strdup(c->global.my_public_key_hex) : u_strdup(c->global.my_private_key_hex);
    free_config(c); return r;
}
static int lks(struct UTUN_INSTANCE* i) {
    int n = 0; struct ll_entry* entry = i->connections->head;
    while (entry) { struct conn_queue_entry* ce = (struct conn_queue_entry*)entry->data;
        struct ETCP_CONN* c = ce->conn;
        struct ETCP_LINK* l = c->links; while (l) { if (l->initialized && l->link_status) n++; l = l->next; }
        entry = entry->next; }
    return n;
}
static void fail(const char* msg) { fprintf(stderr, "FAIL: %s\n", msg); fflush(stderr); result = 1; }
static void to_cb(void* arg) { (void)arg; fprintf(stderr, "TIMEOUT\n"); result = 1; }

static void test1(void* arg); static void test2(void* arg);
static void test3(void* arg); static void test4(void* arg);
static void test5(void* arg); static void test6(void* arg);
static void test7(void* arg); static void test8(void* arg);
static void test9(void* arg); static void test10(void* arg);

/* ----- callback state per test ----- */
static volatile int cb_type = 0, cb_conn_ok = 0, cb_count = 0;

static void connect_ccb(void* a, struct ETCP_CONN* conn, int type) { cb_type = type; cb_conn_ok = (conn != NULL); cb_count++; }

/* ----- helper: create minimal TOPO_NODEQ with 1 v4 addr ----- */
static struct TOPO_NODEQ* mknode(uint64_t nid, const uint8_t pubkey[32],
                                  uint8_t a, uint8_t b, uint8_t c, uint8_t d, uint16_t port) {
    struct TOPO_NODEQ* nq = u_calloc(1, sizeof(struct TOPO_NODEQ));
    struct TOPO_NODE* ni = u_calloc(1, sizeof(struct TOPO_NODE));
    if (!nq || !ni) { u_free(nq); u_free(ni); return NULL; }
    ni->ref_count = 1; ni->node_id = nid; ni->ver = 0;
    memcpy(ni->public_key, pubkey, SC_PUBKEY_SIZE);
    nq->ll.size = sizeof(struct TOPO_NODEQ) - sizeof(struct ll_entry);
    nq->node = ni; nq->hash_node_id = nid;
    struct TOPO_SOCKMETA4* sm = u_calloc(1, sizeof(struct TOPO_SOCKMETA4));
    if (sm) { sm->id = 0; sm->config_type = CFG_SERVER_TYPE_PUBLIC; sm->nat_type = NAT_TYPE_UNKNOWN; sm->next = ni->v4_sock_meta; ni->v4_sock_meta = sm; }
    struct TOPO_ADDR4* addr = u_calloc(1, sizeof(struct TOPO_ADDR4));
    if (addr) { addr->addr[0] = a; addr->addr[1] = b; addr->addr[2] = c; addr->addr[3] = d; addr->port = port; addr->type = TOPO_ADDR_INTERFACE; addr->socket_id = 0; addr->next = ni->v4_addrs; ni->v4_addrs = addr; }
    return nq;
}

/* ======================== Test 1: already connected → immediate EARLY+LATE ======================== */
static void test1(void* arg) {
    (void)arg; if (result) return;
    if (lks(g_a) < 1 || lks(g_b) < 1) { uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test1, "t1"); return; }
    struct TOPO_NODEQ* nb = topo_node_find_by_id(topo_groups_get_default(g_a->topo_groups), nid_b);
    if (!nb) { uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test1, "t1b"); return; }
    fprintf(stderr, "Test 1: already connected → EARLY|LATE\n"); fflush(stderr);
    cb_type = 0; cb_conn_ok = 0; cb_count = 0;
    int r = etcp_connect(g_a, nb, connect_ccb, NULL, ETCP_CONNECT_EARLY | ETCP_CONNECT_LATE);
    if (r != 0) { fail("etcp_connect returned error"); return; }
    uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test2, "t2");
}
static void test2(void* arg) {
    (void)arg; if (result) return;
    if (!cb_count) { uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test2, "t2"); return; }
    if (cb_count != 2) { fail("test1: expected 2 callbacks"); return; }
    if (!cb_conn_ok) { fail("test1: conn=NULL in callback"); return; }
    fprintf(stderr, "  OK: %d callbacks, conn=OK\n", cb_count); fflush(stderr);
    /* advance to test3 */
    uasync_call_soon(ua, NULL, test3);
}

/* ======================== Test 3: unreachable node → timeout error ======================== */
static void test3(void* arg) {
    (void)arg; if (result) return;
    fprintf(stderr, "Test 3: unreachable node → error\n"); fflush(stderr);
    struct SC_MYKEYS fk; sc_generate_keypair(&fk);
    uint64_t fnid = 0xF000000000000001ULL;
    struct TOPO_NODEQ* fn = mknode(fnid, fk.public_key, 127,0,0,1, 49999);
    if (!fn) { fail("mknode failed"); return; }
    queue_data_put_with_index(topo_groups_get_default(g_a->topo_groups)->nodes, &fn->ll);
    g_a->etcp_connect_timeout_tb = 2000;
    cb_type = -1; cb_conn_ok = -1; cb_count = 0;
    int r = etcp_connect(g_a, fn, connect_ccb, NULL, ETCP_CONNECT_EARLY | ETCP_CONNECT_LATE);
    if (r != 0) { fail("etcp_connect(unreachable) returned error"); return; }
    uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test4, "t4");
}
static void test4(void* arg) {
    (void)arg; if (result) return;
    if (!cb_count) { uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test4, "t4"); return; }
    if (cb_count != 1) { fail("test3: expected 1 callback"); return; }
    if (cb_conn_ok != 0) { fail("test3: expected conn=NULL"); return; }
    if (cb_type != 0) { fail("test3: expected type=0 (error)"); return; }
    fprintf(stderr, "  OK: error callback delivered (NULL, type=0)\n"); fflush(stderr);
    uasync_call_soon(ua, NULL, test5);
}

/* ======================== Test 5: double connect to unreachable → both callbacks ======================== */
static int t5_cb_count = 0;
static void t5_ccb(void* a, struct ETCP_CONN* conn, int type) {
    (void)a; (void)conn; if (type == 0) t5_cb_count++;
}

static void test5(void* arg) {
    (void)arg; if (result) return;
    fprintf(stderr, "Test 5: double connect to unreachable → both callbacks\n"); fflush(stderr);
    struct SC_MYKEYS fk; sc_generate_keypair(&fk);
    uint64_t fnid = 0xF000000000000002ULL;
    struct TOPO_NODEQ* fn = mknode(fnid, fk.public_key, 127,0,0,1, 49998);
    if (!fn) { fail("mknode failed"); return; }
    queue_data_put_with_index(topo_groups_get_default(g_a->topo_groups)->nodes, &fn->ll);
    g_a->etcp_connect_timeout_tb = 2000;
    t5_cb_count = 0;
    int r1 = etcp_connect(g_a, fn, t5_ccb, NULL, ETCP_CONNECT_EARLY | ETCP_CONNECT_LATE);
    int r2 = etcp_connect(g_a, fn, t5_ccb, NULL, ETCP_CONNECT_EARLY);
    if (r1 != 0 || r2 != 0) { fail("etcp_connect double call returned error"); return; }
    uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test6, "t6");
}
static void test6(void* arg) {
    (void)arg; if (result) return;
    if (t5_cb_count < 2) { uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test6, "t6"); return; }
    if (t5_cb_count != 2) { fail("test5: expected 2 error callbacks"); return; }
    fprintf(stderr, "  OK: both callbacks fired\n"); fflush(stderr);
    uasync_call_soon(ua, NULL, test7);
}

/* ======================== Test 7: EARLY-only flag ======================== */
static void test7(void* arg) {
    (void)arg; if (result) return;
    struct TOPO_NODEQ* nb = topo_node_find_by_id(topo_groups_get_default(g_a->topo_groups), nid_b);
    if (!nb || lks(g_a) < 1) { uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test7, "t7"); return; }
    fprintf(stderr, "Test 7: EARLY-only flag\n"); fflush(stderr);
    cb_type = 0; cb_conn_ok = 0; cb_count = 0;
    int r = etcp_connect(g_a, nb, connect_ccb, NULL, ETCP_CONNECT_EARLY);
    if (r != 0) { fail("etcp_connect(EARLY) returned error"); return; }
    uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test8, "t8");
}
static void test8(void* arg) {
    (void)arg; if (result) return;
    if (!cb_count) { uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test8, "t8"); return; }
    if (cb_count != 1) { fail("test7: expected exactly 1 callback"); return; }
    if (cb_type != ETCP_CONNECT_EARLY) { fail("test7: expected EARLY type"); return; }
    fprintf(stderr, "  OK: only EARLY fired\n"); fflush(stderr);
    uasync_call_soon(ua, NULL, test9);
}

/* ======================== Test 9: LATE-only flag ======================== */
static void test9(void* arg) {
    (void)arg; if (result) return;
    struct TOPO_NODEQ* nb = topo_node_find_by_id(topo_groups_get_default(g_a->topo_groups), nid_b);
    if (!nb || lks(g_a) < 1) { uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test9, "t9"); return; }
    fprintf(stderr, "Test 9: LATE-only flag\n"); fflush(stderr);
    cb_type = 0; cb_conn_ok = 0; cb_count = 0;
    int r = etcp_connect(g_a, nb, connect_ccb, NULL, ETCP_CONNECT_LATE);
    if (r != 0) { fail("etcp_connect(LATE) returned error"); return; }
    uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test10, "t10");
}
static void test10(void* arg) {
    (void)arg; if (result) return;
    if (!cb_count) { uasync_set_timeout(ua, 50, NULL, (timeout_callback_t)test10, "t10"); return; }
    if (cb_count != 1) { fail("test9: expected exactly 1 callback"); return; }
    if (cb_type != ETCP_CONNECT_LATE) { fail("test9: expected LATE type"); return; }
    fprintf(stderr, "  OK: only LATE fired\n"); fflush(stderr);
    fprintf(stderr, "=== ALL PASSED ===\n"); fflush(stderr);
    result = 2;
}

/* ======================== setup / cleanup / main ======================== */

static void setup(void) {
    test_mkdtemp(tdir);
    int base = 48000 + (getpid() % 10000); pa = base; pb = base + 1;
    snprintf(ca, sizeof(ca), "%s/a.conf", tdir); snprintf(cb, sizeof(cb), "%s/b.conf", tdir);
    wf(ca, "[global]\ntun_ip=10.99.0.1/24\ntun_ifname=tun99\n[server: s1]\naddr=127.0.0.1:%d\ntype=public\n[allowed_keys]\nallow_all=1\n", pa);
    wf(cb, "[global]\ntun_ip=10.99.0.2/24\ntun_ifname=tun98\n[server: s1]\naddr=127.0.0.1:%d\ntype=public\n[allowed_keys]\nallow_all=1\n", pb);
    config_ensure_keys_and_node_id(ca); config_ensure_keys_and_node_id(cb);
    { struct utun_config* cfa = parse_config(ca); struct utun_config* cfb = parse_config(cb);
      nid_a = cfa->global.my_node_id; nid_b = cfb->global.my_node_id;
      free_config(cfa); free_config(cfb); }
    char *p0 = gv(ca,"pub"), *r0 = gv(ca,"priv"), *p1 = gv(cb,"pub"), *r1 = gv(cb,"priv");
    wf(ca, "[global]\nmy_private_key=%s\nmy_public_key=%s\ntun_ip=10.99.0.1/24\ntun_ifname=tun99\n[server: s1]\naddr=127.0.0.1:%d\ntype=public\n[client: to_b]\nkeepalive=1\npeer_public_key=%s\nlink=s1:127.0.0.1:%d\n[allowed_keys]\nallow_all=1\n", r0, p0, pa, p1, pb);
    wf(cb, "[global]\nmy_private_key=%s\nmy_public_key=%s\ntun_ip=10.99.0.2/24\ntun_ifname=tun98\n[server: s1]\naddr=127.0.0.1:%d\ntype=public\n[allowed_keys]\nallow_all=1\n", r1, p1, pb);
    u_free(p0); u_free(r0); u_free(p1); u_free(r1);
}
static void cleanup(void) { test_unlink(ca); test_unlink(cb); test_rmdir(tdir); }

int main(void) {
    debug_config_init(); debug_set_level(DEBUG_LEVEL_WARN);
    utun_instance_set_tun_init_enabled(0); setup();
    ua = uasync_create();
    g_a = utun_instance_create(ua, ca); g_b = utun_instance_create(ua, cb);
    if (!g_a || !g_b) goto done;
    utun_instance_init(g_a); utun_instance_init(g_b);
    uasync_call_soon(ua, NULL, test1);
    ttimer = uasync_set_timeout(ua, TIMEOUT_TB, NULL, to_cb, "to");
    { uint64_t start = get_time_tb();
      while (!result && (int)(get_time_tb() - start) < TIMEOUT_TB + 50000) { uasync_poll(ua, POLL_MS); } }
    fprintf(stderr, "final result=%d\n", result); fflush(stderr);
done:
    if (ttimer) uasync_cancel_timeout(ua, ttimer);
    if (g_a) { g_a->running = 0; utun_instance_destroy(g_a); }
    if (g_b) { g_b->running = 0; utun_instance_destroy(g_b); }
    if (ua) { uasync_destroy(ua, 0); ua = NULL; }
    cleanup();
    return (result == 2) ? 0 : 1;
}