/**
 * @file nat_detection.c
 * @brief NAT-детекция: STUN-like проверка через третьего пира.
 *
 * Самодостаточный модуль:
 *   - Сам биндится на ETCP_ID_NAT_DETECTION (0x02) и принимает свои пакеты
 *   - Сам подписывается на conn_up/conn_down новых ETCP-соединений
 *   - Ищет третий узел прямо в instance->connections (без привязки к BGP-группе)
 */
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#else
#include <arpa/inet.h>
#endif
#include "../lib/platform_compat.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"
#include "utun_instance.h"
#include "etcp_api.h"
#include "etcp.h"
#include "etcp_connections.h"
#include "config_parser.h"
#include "topo_node.h"
#include "topo_group.h"
#include "route_ping.h"
#include "nat_detection.h"
#include "../lib/u_async.h"


/* ====== Internal NAT check logic ====== */

static struct ETCP_CONN* nat_check_find_third_node(struct NAT_DETECTION* nd, struct ETCP_CONN* exclude) {
    if (!nd || !nd->inst || !nd->inst->connections) return NULL;
    struct ll_entry* e = nd->inst->connections->head;
    while (e) {
        struct conn_queue_entry* ce = (struct conn_queue_entry*)e->data;
        if (ce->conn && ce->conn != exclude && ce->conn->links)
            return ce->conn;
        e = e->next;
    }
    return NULL;
}

static void nat_link_check_cb(int success, uint16_t avg_rtt, uint8_t count_sent, uint8_t count_ok, void* arg) {
    (void)avg_rtt; (void)count_sent; (void)count_ok;
    struct nat_check_arg* na = (struct nat_check_arg*)arg;
    if (!na || !na->link) { u_free(na); return; }
    struct ETCP_LINK* link = na->link;
    if (link->nat_check_status != NAT_CHECK_IN_PROGRESS) {
        DEBUG_DEBUG(DEBUG_CATEGORY_BGP, "nat check cb: link status=%d, skip (conn down?)", link->nat_check_status);
        u_free(na); return;
    }
    link->nat_type = success ? NAT_TYPE_EIM : NAT_TYPE_STRICT;
    link->nat_check_status = success ? NAT_CHECK_EIM : NAT_CHECK_STRICT;
    if (link->etcp)
        nat_detection_send_nat_info(link->etcp->instance->nat_det,
                                     link->etcp, link->remote_socket_id,
                                     na->nat_ip, na->nat_port, link->nat_type);
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "NAT check result: type=%s for conn=%s nat=%s:%u",
               success ? "EIM" : "STRICT", link->etcp->log_name,
               ip_to_str(&na->nat_ip, AF_INET).str, na->nat_port);
    u_free(na);
}

static void nat_detection_start_link_check(struct NAT_DETECTION* nd, struct ETCP_LINK* link) {
    if (!nd || !link || !link->conn || !link->etcp) return;
    if (link->nat_type == NAT_TYPE_DIRECT) return;
    if (link->nat_check_status == NAT_CHECK_IN_PROGRESS) return;
    if (link->conn->type == CFG_SERVER_TYPE_PRIVATE || link->conn->type == CFG_SERVER_TYPE_LOCAL) {
        DEBUG_DEBUG(DEBUG_CATEGORY_BGP, "skip nat check for private/local socket link"); return;
    }
    struct ETCP_CONN* third_conn = nat_check_find_third_node(nd, link->etcp);
    if (!third_conn) { DEBUG_DEBUG(DEBUG_CATEGORY_BGP, "no third node found for link"); return; }

    uint32_t target_ip; uint16_t target_port;
    if (link->nat_ip != 0 && link->nat_port != 0) {
        target_ip = link->nat_ip; target_port = link->nat_port;
    } else {
        struct sockaddr_in* sin = (struct sockaddr_in*)&link->remote_addr;
        target_ip = sin->sin_addr.s_addr; target_port = ntohs(sin->sin_port);
    }
    if (!nd->allow_nat_check_local && is_local_subnet(target_ip)) {
        DEBUG_DEBUG(DEBUG_CATEGORY_BGP, "local connection: %s, skip nat check",
                     ip_to_str(&target_ip, AF_INET).str); return;
    }

    struct TOPO_GROUP* group = nd->inst->topo_groups
        ? topo_groups_get_default(nd->inst->topo_groups) : NULL;
    struct nat_check_arg* arg = u_calloc(1, sizeof(struct nat_check_arg));
    if (!arg) return;
    arg->link = link; arg->nat_ip = target_ip; arg->nat_port = target_port;
    const uint8_t* pubkey = link->etcp->crypto_ctx.peer_key_set ? link->etcp->crypto_ctx.peer_public_key : NULL;
    if (!pubkey) { u_free(arg); DEBUG_DEBUG(DEBUG_CATEGORY_BGP, "no peer pubkey for link=%p", (void*)link); return; }
    int ret = route_ping_send_req_addr(nd, group, third_conn, target_ip, target_port,
                                        3, 500, 1000, 5000, nat_link_check_cb, arg, pubkey);
    if (ret == 0) {
        link->nat_check_status = NAT_CHECK_IN_PROGRESS;
        DEBUG_INFO(DEBUG_CATEGORY_BGP, "ping remote request via=%s (%s:%d)",
                   third_conn->log_name, ip_to_str(&target_ip, AF_INET).str, target_port);
    } else {
        u_free(arg);
        DEBUG_WARN(DEBUG_CATEGORY_BGP, "failed to start for link=%s (%s:%d)",
                   third_conn->log_name, ip_to_str(&target_ip, AF_INET).str, target_port);
    }
}


/* ====== Handle incoming NAT messages (static, dispatched from receive_cbk) ====== */

static void nat_detection_handle_nat_info(struct NAT_DETECTION* nd, struct TOPO_GROUP* group,
                                           struct ETCP_CONN* from_conn,
                                           const uint8_t* data, size_t len) {
    if (!nd || !from_conn || !data || len < sizeof(struct NATDET_MSG_NAT_INFO)) return;

    const struct NATDET_MSG_NAT_INFO* info = (const struct NATDET_MSG_NAT_INFO*)data;
    uint32_t nat_ip = (info->nat_ip[0] << 24) | (info->nat_ip[1] << 16)
                    | (info->nat_ip[2] << 8) | info->nat_ip[3];
    uint16_t nat_port = ntohs(info->nat_port);
    uint8_t socket_id = info->socket_id;

    uint8_t verified_type;
    if (info->nat_type == NAT_TYPE_EIM) verified_type = NAT_VERIFIED_EIM;
    else if (info->nat_type == NAT_TYPE_STRICT) verified_type = NAT_VERIFIED_STRICT;
    else if (info->nat_type == NAT_TYPE_DIRECT) verified_type = NAT_VERIFIED_DIRECT;
    else verified_type = NAT_VERIFIED_UNKNOWN;

    if (group && group->local_node) {
        struct TOPO_NODE* ni = topo_node_registry_find(group->instance->topo_groups, group->local_node->node_id);
        if (!ni) return;
        int data_changed = 0;
        struct TOPO_SOCKMETA4* sm = ni->v4_sock_meta;
        while (sm) {
            if (sm->id == socket_id) {
                if (sm->nat_type != verified_type) { sm->nat_type = verified_type; data_changed = 1; }
                break;
            }
            sm = sm->next;
        }

        struct ETCP_SOCKET* es = group->instance->etcp_sockets;
        while (es) {
            if (es->sock_id == socket_id) {
                if (verified_type == NAT_VERIFIED_STRICT) {
                    struct sockaddr_in* nat_sin = (struct sockaddr_in*)&es->nat_addr;
                    nat_sin->sin_family = AF_INET;
                    nat_sin->sin_addr.s_addr = nat_ip;
                    nat_sin->sin_port = htons(nat_port);
                    es->nat_type = verified_type;
                } else {
                    struct sockaddr_in* nat_sin = (struct sockaddr_in*)&es->nat_addr;
                    uint32_t old_nat_ip = nat_sin->sin_addr.s_addr;
                    uint16_t old_nat_port = ntohs(nat_sin->sin_port);
                    nat_sin->sin_family = AF_INET;
                    nat_sin->sin_addr.s_addr = nat_ip;
                    nat_sin->sin_port = htons(nat_port);
                    if (old_nat_ip != nat_ip || old_nat_port != nat_port) data_changed = 1;
                    es->nat_type = verified_type;
                }
                break;
            }
            es = es->next;
        }

        if (data_changed) {
            int prev_v4a = topo_list_count((struct _topo_head*)ni->v4_addrs);
            topo_group_update_my_nodeinfo(group->instance, group);
            if (topo_list_count((struct _topo_head*)ni->v4_addrs) != prev_v4a) {
                group->local_node->dirty = 1;
                ni->ver = (ni->ver % 255) + 1;
                group->local_node->last_ver = ni->ver;
            }
            if (topo_list_count((struct _topo_head*)ni->v4_addrs) == prev_v4a) {
                struct TOPO_ADDR4* a = ni->v4_addrs;
                while (a) {
                    if (a->type == TOPO_ADDR_INTERFACE && a->socket_id == socket_id) {
                        uint32_t a_ip; memcpy(&a_ip, a->addr, 4);
                        if (a_ip == nat_ip && a->port == nat_port) {
                            a->type = TOPO_ADDR_NAT; a->socket_id = socket_id | 1;
                            group->local_node->dirty = 1;
                            ni->ver = (ni->ver % 255) + 1;
                            group->local_node->last_ver = ni->ver;
                            DEBUG_INFO(DEBUG_CATEGORY_BGP, "NAT_INFO matched interface addr, updated entry type to NAT sock=%d", socket_id);
                        }
                        break;
                    }
                    a = a->next;
                }
            }
            if (group->local_node->dirty && group->senders_list) {
                struct ll_entry* se = group->senders_list->head;
                while (se) {
                    struct TOPO_GROUP_CONN_ITEM* item = (struct TOPO_GROUP_CONN_ITEM*)se->data;
                    if (item && item->conn)
                        topo_group_send_nodeinfo(group, group->local_node, item->conn, 0);
                    se = se->next;
                }
            }
        }
    } else {
        struct ETCP_SOCKET* es = nd->inst->etcp_sockets;
        while (es) {
            if (es->sock_id == socket_id) {
                struct sockaddr_in* nat_sin = (struct sockaddr_in*)&es->nat_addr;
                nat_sin->sin_family = AF_INET;
                nat_sin->sin_addr.s_addr = nat_ip;
                nat_sin->sin_port = htons(nat_port);
                es->nat_type = verified_type;
                DEBUG_INFO(DEBUG_CATEGORY_BGP, "NAT_INFO socket=%s id=%d type=%d addr=%s:%u (no topo_group)",
                           es->name, socket_id, verified_type,
                           ip_to_str(&nat_ip, AF_INET).str, nat_port);
                break;
            }
            es = es->next;
        }
    }

    struct ETCP_LINK* l = from_conn->links;
    while (l) { l->nat_type = info->nat_type; l = l->next; }
}

static void nat_detection_handle_nat_check_req(struct NAT_DETECTION* nd, struct TOPO_GROUP* group,
                                                struct ETCP_CONN* from_conn,
                                                const uint8_t* data, size_t len) {
    if (!nd || !from_conn || !data || len < sizeof(struct NATDET_MSG_NAT_CHECK_REQ)) return;
    (void)group;
    const struct NATDET_MSG_NAT_CHECK_REQ* req = (const struct NATDET_MSG_NAT_CHECK_REQ*)data;
    uint8_t socket_id = req->socket_id;
    struct ETCP_LINK* target_link = NULL;
    struct ETCP_LINK* l = from_conn->links;
    while (l) {
        if (l->remote_socket_id == socket_id) { target_link = l; break; }
        l = l->next;
    }
    if (!target_link) return;
    if (target_link->nat_check_status == NAT_CHECK_IN_PROGRESS) return;
    nat_detection_start_link_check(nd, target_link);
}


/* ====== ETCP callbacks (self-registering) ====== */

static void nat_detection_trigger_checks(struct NAT_DETECTION* nd);

static void nat_detection_on_conn_up(struct ETCP_CONN* conn, int event, void* arg) { (void)event;
    struct NAT_DETECTION* nd = (struct NAT_DETECTION*)arg;
    if (!nd) return;
    nat_detection_trigger_checks(nd);
}

static void nat_detection_on_conn_down(struct ETCP_CONN* conn, int event, void* arg) { (void)event;
    struct NAT_DETECTION* nd = (struct NAT_DETECTION*)arg;
    if (!nd || !conn) return;
    route_ping_cancel_for_conn(nd, conn);
    struct ETCP_LINK* lnk = conn->links;
    while (lnk) {
        if (lnk->nat_check_status == NAT_CHECK_IN_PROGRESS) lnk->nat_check_status = NAT_CHECK_NONE;
        lnk = lnk->next;
    }
}

static void nat_detection_new_conn_cbk(struct ETCP_CONN* conn, void* arg) {
    if (!conn || !conn->instance) return;
    struct NAT_DETECTION* nd = (struct NAT_DETECTION*)arg;
    DEBUG_DEBUG(DEBUG_CATEGORY_NAT, "NAT detection set callbacks: %s", conn->log_name);
    etcp_conn_add_cbk(conn, nat_detection_on_conn_up, nd, ETCP_CBK_EVENT_UP);
    etcp_conn_add_cbk(conn, nat_detection_on_conn_down, nd, ETCP_CBK_EVENT_DOWN);
}


/* ====== Receive callback (self-binded on ETCP_ID_NAT_DETECTION) ====== */

static const char* nat_subcmd_name(uint8_t subcmd) {
    switch (subcmd) {
        case NATDET_SUBCMD_PING_REQ:      return "PING_REQ";
        case NATDET_SUBCMD_PING_RESP:     return "PING_RESP";
        case NATDET_SUBCMD_NAT_INFO:      return "NAT_INFO";
        case NATDET_SUBCMD_NAT_CHECK_REQ: return "NAT_CHECK_REQ";
        default: return "?";
    }
}

static void nat_detection_receive_cbk(struct ETCP_CONN* from_conn, struct ll_entry* entry) {
    if (!from_conn || !entry || entry->len < 2) {
        if (entry) { queue_dgram_free(entry); queue_entry_free(entry); } return;
    }
    struct UTUN_INSTANCE* instance = from_conn->instance;
    if (!instance || !instance->nat_det) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "invalid instance/nat_det");
        queue_dgram_free(entry); queue_entry_free(entry); return;
    }
    struct TOPO_GROUP* group = instance->topo_groups
        ? topo_groups_get_default(instance->topo_groups) : NULL;
    struct NAT_DETECTION* nd = instance->nat_det;
    uint8_t* data = entry->dgram; uint8_t subcmd = data[1];

    DEBUG_INFO(DEBUG_CATEGORY_BGP, "NAT recv %s from %s len=%zu group=%s",
               nat_subcmd_name(subcmd), from_conn->log_name, entry->len, group ? "yes" : "no");

    if (subcmd == NATDET_SUBCMD_PING_REQ)
        route_ping_handle_req(nd, group, from_conn, data, entry->len);
    else if (subcmd == NATDET_SUBCMD_PING_RESP)
        route_ping_handle_resp(nd, from_conn, data, entry->len);
    else if (subcmd == NATDET_SUBCMD_NAT_INFO)
        nat_detection_handle_nat_info(nd, group, from_conn, data, entry->len);
    else if (subcmd == NATDET_SUBCMD_NAT_CHECK_REQ)
        nat_detection_handle_nat_check_req(nd, group, from_conn, data, entry->len);

    queue_dgram_free(entry); queue_entry_free(entry);
}


/* ====== NAT check trigger ====== */

static void nat_detection_trigger_checks(struct NAT_DETECTION* nd) {
    if (!nd || !nd->inst) return;
    struct ll_entry* entry = nd->inst->connections->head;
    while (entry) {
        struct conn_queue_entry* ce = (struct conn_queue_entry*)entry->data;
        struct ETCP_LINK* l = ce->conn->links;
        while (l) {
            if (l->initialized && l->conn && l->nat_check_status < NAT_CHECK_IN_PROGRESS)
                nat_detection_start_link_check(nd, l);
            l = l->next;
        }
        entry = entry->next;
    }
}


/* ====== Lifecycle ====== */

struct NAT_DETECTION* nat_detection_create(struct UTUN_INSTANCE* inst) {
    if (!inst) { DEBUG_ERROR(DEBUG_CATEGORY_BGP, "nat_detection_create: inst is NULL"); return NULL; }
    struct NAT_DETECTION* nd = u_calloc(1, sizeof(struct NAT_DETECTION));
    if (!nd) { DEBUG_ERROR(DEBUG_CATEGORY_BGP, "nat_detection_create: alloc failed"); return NULL; }
    nd->inst = inst;
    nd->next_ping_req_id = 1;
    nd->allow_nat_check_local = 0;
    etcp_bind(inst, ETCP_ID_NAT_DETECTION, nat_detection_receive_cbk);
    etcp_add_new_conn_cbk(inst, nat_detection_new_conn_cbk, nd);
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "NAT detection initialized");
    return nd;
}

void nat_detection_destroy(struct NAT_DETECTION* nd) {
    if (!nd) return;
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "NAT detection destroy");
    etcp_remove_new_conn_cbk(nd->inst, nat_detection_new_conn_cbk, nd);
    etcp_unbind(nd->inst, ETCP_ID_NAT_DETECTION);
    route_ping_destroy_pending(nd);
    u_free(nd);
}


/* ====== Link lifecycle ====== */

void nat_detection_link_ready(struct NAT_DETECTION* nd, struct ETCP_LINK* link) {
    if (!nd || !link) return;
    if (link->nat_check_status < NAT_CHECK_IN_PROGRESS)
        nat_detection_start_link_check(nd, link);
}

void nat_detection_cancel_for_conn(struct NAT_DETECTION* nd, struct ETCP_CONN* conn) {
    if (!nd || !conn) return;
    route_ping_cancel_for_conn(nd, conn);
}


/* ====== Send NAT messages ====== */

void nat_detection_send_nat_info(struct NAT_DETECTION* nd, struct ETCP_CONN* to_conn,
                                  uint8_t socket_id, uint32_t nat_ip, uint16_t nat_port,
                                  uint8_t nat_type) {
    (void)nd;
    if (!to_conn) return;
    struct NATDET_MSG_NAT_INFO* pkt = u_calloc(1, sizeof(struct NATDET_MSG_NAT_INFO));
    if (!pkt) return;
    pkt->cmd = ETCP_ID_NAT_DETECTION; pkt->subcmd = NATDET_SUBCMD_NAT_INFO;
    pkt->socket_id = socket_id;
    pkt->nat_ip[0] = (nat_ip >> 24) & 0xFF; pkt->nat_ip[1] = (nat_ip >> 16) & 0xFF;
    pkt->nat_ip[2] = (nat_ip >> 8) & 0xFF; pkt->nat_ip[3] = nat_ip & 0xFF;
    pkt->nat_port = htons(nat_port); pkt->nat_type = nat_type;
    struct ll_entry* e = queue_entry_new(0);
    if (!e) { u_free(pkt); return; }
    e->dgram = (uint8_t*)pkt; e->len = sizeof(struct NATDET_MSG_NAT_INFO);
    if (etcp_send(to_conn, e) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "etcp_send NAT_INFO failed");
        u_free(pkt); queue_entry_free(e);
    }
}

static void nat_detection_send_nat_check_req(struct NAT_DETECTION* nd, struct ETCP_CONN* to_conn,
                                              uint8_t socket_id) {
    (void)nd;
    if (!to_conn) return;
    struct NATDET_MSG_NAT_CHECK_REQ* pkt = u_calloc(1, sizeof(struct NATDET_MSG_NAT_CHECK_REQ));
    if (!pkt) return;
    pkt->cmd = ETCP_ID_NAT_DETECTION; pkt->subcmd = NATDET_SUBCMD_NAT_CHECK_REQ;
    pkt->socket_id = socket_id;
    struct ETCP_LINK* l = to_conn->links;
    while (l) {
        if (l->remote_socket_id == socket_id) {
            struct sockaddr_in* sin = (struct sockaddr_in*)&l->conn->interface_addr;
            if (sin->sin_family == AF_INET) {
                pkt->interface_ip = sin->sin_addr.s_addr;
                pkt->interface_port = sin->sin_port;
            } break;
        }
        l = l->next;
    }
    struct ll_entry* e = queue_entry_new(0);
    if (!e) { u_free(pkt); return; }
    e->dgram = (uint8_t*)pkt; e->len = sizeof(struct NATDET_MSG_NAT_CHECK_REQ);
    if (etcp_send(to_conn, e) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "etcp_send NAT_CHECK_REQ failed");
        u_free(pkt); queue_entry_free(e);
    }
}


/* ====== Control API ====== */

void nat_detection_request_check_all(struct NAT_DETECTION* nd) {
    if (!nd || !nd->inst) return;
    struct ll_entry* entry = nd->inst->connections->head;
    while (entry) {
        struct conn_queue_entry* ce = (struct conn_queue_entry*)entry->data;
        struct ETCP_LINK* link = ce->conn->links;
        while (link) {
            if (link->nat_check_status < NAT_CHECK_IN_PROGRESS)
                nat_detection_start_link_check(nd, link);
            link = link->next;
        }
        entry = entry->next;
    }
}


/* ====== Test helper ====== */

void nat_detection_set_allow_local(struct NAT_DETECTION* nd, int allow) {
    if (!nd) return;
    nd->allow_nat_check_local = allow ? 1 : 0;
}