#include "etcp_api.h"
#include "etcp.h"
#include "utun_instance.h"
#include "topo_node.h"
#include "stcp_link.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"
#include "../lib/u_async.h"
#include "../lib/socket_compat.h"
#include <string.h>

#define DEBUG_CATEGORY_ETCP_CONNECT DEBUG_CATEGORY_ETCP

struct etcp_connect_cb_node {
    etcp_connect_callback_t cb;
    void* arg;
    uint8_t flags;
    struct etcp_connect_cb_node* next;
};

struct ETCP_CONNECT {
    struct ETCP_CONNECT* next;
    struct UTUN_INSTANCE* instance;
    uint64_t node_id;
    struct ETCP_CONN* conn;
    struct etcp_connect_cb_node* cb_list;
    void* initial_timer;
    void* settle_timer;
    struct stcp_link* tcp_link;
    uint16_t min_rtt;
    uint8_t early_delivered : 1;
    uint8_t done : 1;
    uint8_t tcp_ready : 1;
};
static void tcp_link_ready_cb(struct stcp_link* link, void* arg);
static void tcp_link_close_cb(struct stcp_link* link, int err, void* arg);
static void connect_ready_cb(struct ETCP_CONN* conn, void* arg);

static struct ETCP_CONNECT* connect_find(struct UTUN_INSTANCE* inst, uint64_t node_id) {
    struct ETCP_CONNECT* ctx = inst->pending_connects;
    while (ctx) { if (ctx->node_id == node_id) return ctx; ctx = ctx->next; }
    return NULL;
}

static void connect_deliver(struct ETCP_CONNECT* ctx, int type) {
    struct etcp_connect_cb_node** pp = &ctx->cb_list;
    while (*pp) {
        struct etcp_connect_cb_node* node = *pp;
        if (type == 0 || (node->flags & (uint8_t)type)) {
            *pp = node->next;
            node->cb(node->arg, (type == 0) ? NULL : ctx->conn, type);
            u_free(node);
        } else {
            pp = &node->next;
        }
    }
}

static void connect_cancel(struct ETCP_CONNECT* ctx) {
    if (!ctx) return;
    struct UTUN_INSTANCE* inst = ctx->instance;
    if (inst) {
        struct ETCP_CONNECT** pp = &inst->pending_connects;
        while (*pp && *pp != ctx) pp = &(*pp)->next;
        if (*pp) *pp = ctx->next;
        if (ctx->initial_timer) { uasync_cancel_timeout(inst->ua, ctx->initial_timer); ctx->initial_timer = NULL; }
        if (ctx->settle_timer) { uasync_cancel_timeout(inst->ua, ctx->settle_timer); ctx->settle_timer = NULL; }
    }
    u_free(ctx);
}

static void connect_create_links_v4(struct ETCP_CONNECT* ctx, struct TOPO_NODEQ* node) {
    if (!node || !node->node) return;
    for (const struct TOPO_ADDR4* a = node->node->v4_addrs; a; a = a->next) {
        if (a->port == 0) continue;
        { int zero = 1; for (int j = 0; j < 4; j++) if (a->addr[j] != 0) { zero = 0; break; } if (zero) continue; }
        struct sockaddr_in sin; memset(&sin, 0, sizeof(sin)); sin.sin_family = AF_INET;
        memcpy(&sin.sin_addr.s_addr, a->addr, 4); sin.sin_port = htons(a->port);
        struct sockaddr_storage sa; memcpy(&sa, &sin, sizeof(sin));
        if (!(a->protocol & TOPO_PROTO_TCP)) {
            struct ETCP_SOCKET* s = ctx->instance->etcp_sockets;
            while (s) { if (s->local_addr.ss_family == AF_INET) etcp_link_new(ctx->conn, s, &sa, 0); s = s->next; }
        }
        if ((a->protocol & TOPO_PROTO_TCP) && !ctx->tcp_link) {
            struct stcp_link_config tcp_cfg = {.ua = ctx->instance->ua, .my_keys = &ctx->instance->my_keys, .inst = ctx->instance,
                .peer_pubkey = ctx->conn->crypto_ctx.peer_public_key, .peer_pubkey_mode = 0,
                .remote_addr = &sa, .remote_port = a->port};
            ctx->tcp_link = stcp_link_connect(&tcp_cfg);
            if (ctx->tcp_link) { ctx->conn->transport_link = (void*)ctx->tcp_link; stcp_link_set_on_ready(ctx->tcp_link, tcp_link_ready_cb, ctx); stcp_link_set_on_close(ctx->tcp_link, tcp_link_close_cb, ctx); }
        }
    }
}

static void connect_create_links_v6(struct ETCP_CONNECT* ctx, struct TOPO_NODEQ* node) {
    if (!node || !node->node) return;
    for (const struct TOPO_ADDR6* a = node->node->v6_addrs; a; a = a->next) {
        if (a->port == 0) continue;
        int zero = 1; for (int j = 0; j < 16; j++) if (a->addr[j] != 0) { zero = 0; break; } if (zero) continue;
        struct sockaddr_in6 sin6; memset(&sin6, 0, sizeof(sin6)); sin6.sin6_family = AF_INET6;
        memcpy(&sin6.sin6_addr, a->addr, 16); sin6.sin6_port = htons(a->port);
        struct sockaddr_storage sa; memcpy(&sa, &sin6, sizeof(sin6));
        if (!(a->protocol & TOPO_PROTO_TCP)) {
            struct ETCP_SOCKET* s = ctx->instance->etcp_sockets;
            while (s) { if (s->local_addr.ss_family == AF_INET6) etcp_link_new(ctx->conn, s, &sa, 0); s = s->next; }
        }
        if ((a->protocol & TOPO_PROTO_TCP) && !ctx->tcp_link) {
            struct stcp_link_config tcp_cfg = {.ua = ctx->instance->ua, .my_keys = &ctx->instance->my_keys, .inst = ctx->instance,
                .peer_pubkey = ctx->conn->crypto_ctx.peer_public_key, .peer_pubkey_mode = 0,
                .remote_addr = &sa, .remote_port = a->port};
            ctx->tcp_link = stcp_link_connect(&tcp_cfg);
            if (ctx->tcp_link) {
                ctx->conn->transport_link = (void*)ctx->tcp_link;
                stcp_link_set_on_ready(ctx->tcp_link, tcp_link_ready_cb, ctx);
                stcp_link_set_on_close(ctx->tcp_link, tcp_link_close_cb, ctx);
            }
        }
    }
}

static void tcp_link_ready_cb(struct stcp_link* link, void* arg) {
    struct ETCP_CONNECT* ctx = (struct ETCP_CONNECT*)arg;
    if (!ctx || ctx->done || ctx->tcp_ready) return;
    ctx->tcp_ready = 1;
    DEBUG_INFO(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] TCP link ready for node 0x%016llx", (unsigned long long)ctx->node_id);
    connect_ready_cb(ctx->conn, ctx);
}

static void tcp_link_close_cb(struct stcp_link* link, int err, void* arg) {
    struct ETCP_CONNECT* ctx = (struct ETCP_CONNECT*)arg;
    if (!ctx || ctx->done) return;
    DEBUG_ERROR(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] TCP link closed node 0x%016llx err=%d",
                (unsigned long long)ctx->node_id, err);
    ctx->tcp_link = NULL;
}

static void connect_bgp_ready_cb(struct ETCP_CONN* conn) {
    struct ETCP_CONNECT* ctx = connect_find(conn->instance, conn->peer_node_id);
    if (!ctx || ctx->done) return;
    connect_deliver(ctx, ETCP_CONNECT_BGP_READY);
}

static void connect_initial_timeout_cb(void* arg) {
    struct ETCP_CONNECT* ctx = (struct ETCP_CONNECT*)arg;
    if (!ctx || ctx->done) return;
    ctx->done = 1;
    ctx->initial_timer = NULL;
    DEBUG_ERROR(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] timeout for node 0x%016llx", (unsigned long long)ctx->node_id);
    struct ETCP_CONN* conn = ctx->conn; ctx->conn = NULL;
    if (ctx->tcp_link) { stcp_link_close(ctx->tcp_link); ctx->tcp_link = NULL; }
    connect_deliver(ctx, 0);
    if (conn) etcp_connection_close(conn);
    connect_cancel(ctx);
}

static void connect_settle_timeout_cb(void* arg) {
    struct ETCP_CONNECT* ctx = (struct ETCP_CONNECT*)arg;
    if (!ctx || ctx->done) return;
    ctx->done = 1;
    ctx->settle_timer = NULL;
    DEBUG_INFO(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] settle for node 0x%016llx, min_rtt=%u",
               (unsigned long long)ctx->node_id, ctx->min_rtt);
    struct ETCP_LINK* link = ctx->conn->links;
    while (link) {
        struct ETCP_LINK* next = link->next;
        if (!link->initialized) {
            DEBUG_INFO(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] removing failed link=%p for node 0x%016llx",
                       (void*)link, (unsigned long long)ctx->node_id);
            etcp_link_close(link);
        }
        link = next;
    }
    if (ctx->tcp_link && !ctx->tcp_ready) {
        DEBUG_INFO(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] removing failed TCP link for node 0x%016llx",
                   (unsigned long long)ctx->node_id);
        stcp_link_close(ctx->tcp_link); ctx->tcp_link = NULL;
    }
    etcp_conn_remove_ready_cbk(ctx->conn, connect_ready_cb, ctx);
    connect_deliver(ctx, ETCP_CONNECT_LATE);
    connect_cancel(ctx);
}

static void connect_ready_cb(struct ETCP_CONN* conn, void* arg) {
    struct ETCP_CONNECT* ctx = (struct ETCP_CONNECT*)arg;
    if (!ctx || ctx->done) return;
    struct ETCP_LINK* link = conn->links;
    while (link) {
        if (link->initialized && link->rtt_last && link->rtt_last < ctx->min_rtt)
            ctx->min_rtt = link->rtt_last;
        link = link->next;
    }
    if (ctx->early_delivered) return;
    ctx->early_delivered = 1;
    DEBUG_INFO(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] early ready for node 0x%016llx, min_rtt=%u",
               (unsigned long long)ctx->node_id, ctx->min_rtt);
    connect_deliver(ctx, ETCP_CONNECT_EARLY);
    if (ctx->initial_timer) { uasync_cancel_timeout(ctx->instance->ua, ctx->initial_timer); ctx->initial_timer = NULL; }
    uint32_t settle_tb;
    if (ctx->min_rtt == 0xFFFF) settle_tb = 20000;
    else {
        settle_tb = (uint32_t)ctx->min_rtt * 8;
        if (settle_tb < 5000) settle_tb = 5000;
        if (settle_tb > 20000) settle_tb = 20000;
    }
    ctx->settle_timer = uasync_set_timeout(ctx->instance->ua, (int)settle_tb, ctx,
                        connect_settle_timeout_cb, "etcp_connect_settle");
}

int etcp_connect(struct UTUN_INSTANCE* inst, struct TOPO_NODEQ* node,
                 etcp_connect_callback_t cb, void* arg, uint8_t flags) {
    if (!inst || !node || !cb) return -1;
    uint64_t node_id = node->node->node_id;

    struct ETCP_CONN* existing = inst->connections;
    while (existing) {
        if (existing->peer_node_id == node_id) {
            struct ETCP_LINK* l = existing->links;
            while (l) { if (l->link_state == 3) break; l = l->next; }
            if (l) {
                DEBUG_INFO(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] node 0x%016llx already connected, delivering immediately",
                           (unsigned long long)node_id);
                if (flags & ETCP_CONNECT_EARLY) cb(arg, existing, ETCP_CONNECT_EARLY);
                if (flags & ETCP_CONNECT_LATE)  cb(arg, existing, ETCP_CONNECT_LATE);
                if ((flags & ETCP_CONNECT_BGP_READY) && existing->routing_exchange_active >= 3)
                    cb(arg, existing, ETCP_CONNECT_BGP_READY);
                return 0;
            }
        }
        existing = existing->next;
    }

    struct ETCP_CONNECT* ctx = connect_find(inst, node_id);
    if (ctx) {
        struct etcp_connect_cb_node* cn = u_calloc(1, sizeof(struct etcp_connect_cb_node));
        if (!cn) return -1;
        cn->cb = cb; cn->arg = arg; cn->flags = flags; cn->next = ctx->cb_list; ctx->cb_list = cn;
        DEBUG_INFO(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] node 0x%016llx already connecting, added cb",
                   (unsigned long long)node_id);
        return 0;
    }

    struct ETCP_CONN* conn = etcp_connection_create(inst, NULL);
    if (!conn) { DEBUG_ERROR(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] failed to create ETCP_CONN for node 0x%016llx",
                              (unsigned long long)node_id); cb(arg, NULL, 0); return -1; }
    conn->next = inst->connections; inst->connections = conn; inst->connections_count++;
    if (sc_init_ctx(&conn->crypto_ctx, &inst->my_keys) != SC_OK) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] sc_init_ctx failed for node 0x%016llx", (unsigned long long)node_id);
        etcp_connection_close(conn); cb(arg, NULL, 0); return -1;
    }
    if (sc_set_peer_public_key(&conn->crypto_ctx, node->node->public_key, 0) != SC_OK) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] sc_set_peer_public_key failed for node 0x%016llx", (unsigned long long)node_id);
        etcp_connection_close(conn); cb(arg, NULL, 0); return -1;
    }

    ctx = u_calloc(1, sizeof(struct ETCP_CONNECT));
    if (!ctx) { DEBUG_ERROR(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] alloc failed");
                 etcp_connection_close(conn); cb(arg, NULL, 0); return -1; }
    ctx->instance = inst; ctx->node_id = node_id; ctx->conn = conn;
    ctx->min_rtt = 0xFFFF;
    {
        struct etcp_connect_cb_node* cn = u_calloc(1, sizeof(struct etcp_connect_cb_node));
        if (!cn) { u_free(ctx); etcp_connection_close(conn); cb(arg, NULL, 0); return -1; }
        cn->cb = cb; cn->arg = arg; cn->flags = flags; ctx->cb_list = cn;
    }

    etcp_conn_add_ready_cbk(conn, connect_ready_cb, ctx);
    conn->bgp_ready_cbk = connect_bgp_ready_cb;

    connect_create_links_v4(ctx, node);
    connect_create_links_v6(ctx, node);

    if (!conn->links && !conn->transport_link) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] no links created for node 0x%016llx",
                    (unsigned long long)node_id);
        if (ctx->tcp_link) { stcp_link_close(ctx->tcp_link); ctx->tcp_link = NULL; }
        connect_deliver(ctx, 0);
        etcp_connection_close(conn); u_free(ctx); return -1;
    }

    ctx->initial_timer = uasync_set_timeout(inst->ua, (int)inst->etcp_connect_timeout_tb, ctx, connect_initial_timeout_cb, "etcp_connect_init");
    ctx->next = inst->pending_connects; inst->pending_connects = ctx;
    DEBUG_INFO(DEBUG_CATEGORY_ETCP_CONNECT, "[etcp_connect] started for node 0x%016llx", (unsigned long long)node_id);
    return 0;
}