/**
 * @file test_nat_detection.c
 * @brief Интеграционный тест детекции типа NAT через три узла
 *
 * Топология: C1 -> S <- C2
 * - S получает NODEINFO от C1 (direct peer)
 * - S использует C2 как третий узел для NAT-пинга C1
 * - C2 пингует C1, результат возвращается S
 * - S отправляет NAT_INFO C1
 * - Проверяем корректность заполнения всех NAT-полей
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include "test_utils.h"
#include "../src/etcp.h"
#include "../src/etcp_connections.h"
#include "../src/config_parser.h"
#include "../src/config_updater.h"
#include "../src/utun_instance.h"
#include "../src/routing.h"
#include "../src/route_bgp.h"
#include "../src/route_ping.h"
#include "../src/route_node.h"
#include "../src/tun_if.h"
#include "../src/secure_channel.h"
#include "../lib/u_async.h"
#include "../lib/debug_config.h"

#define TEST_TIMEOUT_MS 5000
#define NODE_ID_S  0x1111111111111111ULL
#define NODE_ID_C1 0x2222222222222222ULL
#define NODE_ID_C2 0x3333333333333333ULL

static struct UTUN_INSTANCE* inst_s = NULL;
static struct UTUN_INSTANCE* inst_c1 = NULL;
static struct UTUN_INSTANCE* inst_c2 = NULL;
static struct UASYNC* ua = NULL;
static int test_timed_out = 0;
static void* test_timeout_id = NULL;

static char temp_dir[] = "/tmp/utun_nat_test_XXXXXX";
static char config_s[256];
static char config_c1[256];
static char config_c2[256];

static struct {
    int done;
    int success;
    uint16_t avg_rtt;
    uint8_t count_sent;
    uint8_t count_ok;
    uint32_t recv_ip;
    uint16_t recv_port;
} nat_ping_result;

static int write_config(const char* path, const char* content) {
    FILE* f = fopen(path, "w");
    if (!f) return -1;
    fprintf(f, "%s", content);
    fclose(f);
    return 0;
}

static char* get_pubkey_from_config(const char* path) {
    struct utun_config* cfg = parse_config(path);
    if (!cfg) return NULL;
    char* pub = strdup(cfg->global.my_public_key_hex);
    free_config(cfg);
    return pub;
}

static int create_temp_configs(void) {
    if (test_mkdtemp(temp_dir) != 0) {
        fprintf(stderr, "Failed to create temp directory\n");
        return -1;
    }
    snprintf(config_s, sizeof(config_s), "%s/s.conf", temp_dir);
    snprintf(config_c1, sizeof(config_c1), "%s/c1.conf", temp_dir);
    snprintf(config_c2, sizeof(config_c2), "%s/c2.conf", temp_dir);

    const char* tpl_s =
        "[global]\n"
        "my_node_id=0x1111111111111111\n"
        "tun_ip=10.100.0.1/24\n"
        "tun_ifname=tun100\n"
        "\n"
        "[server:test_s]\n"
        "addr=127.0.0.1:39011\n"
        "type=public\n";
    if (write_config(config_s, tpl_s) != 0) return -1;
    if (config_ensure_keys_and_node_id(config_s) != 0) return -1;
    char* pub_s = get_pubkey_from_config(config_s);
    if (!pub_s) return -1;

    const char* tpl_c2 =
        "[global]\n"
        "my_node_id=0x3333333333333333\n"
        "tun_ip=10.100.0.3/24\n"
        "tun_ifname=tun103\n"
        "\n"
        "[server:test_c2]\n"
        "addr=127.0.0.1:39013\n"
        "type=public\n"
        "\n"
        "[client:to_s]\n"
        "keepalive=1\n"
        "peer_public_key=%s\n"
        "link=test_c2:127.0.0.1:39011\n";
    char tpl_c2_full[4096];
    snprintf(tpl_c2_full, sizeof(tpl_c2_full), tpl_c2, pub_s);
    if (write_config(config_c2, tpl_c2_full) != 0) { free(pub_s); return -1; }
    if (config_ensure_keys_and_node_id(config_c2) != 0) { free(pub_s); return -1; }
    char* pub_c2 = get_pubkey_from_config(config_c2);
    if (!pub_c2) { free(pub_s); return -1; }

    const char* tpl_c1 =
        "[global]\n"
        "my_node_id=0x2222222222222222\n"
        "tun_ip=10.100.0.2/24\n"
        "tun_ifname=tun102\n"
        "\n"
        "[server:test_c1]\n"
        "addr=127.0.0.1:39012\n"
        "type=public\n"
        "\n"
        "[client:to_s]\n"
        "keepalive=1\n"
        "peer_public_key=%s\n"
        "link=test_c1:127.0.0.1:39011\n";
    char tpl_c1_full[4096];
    snprintf(tpl_c1_full, sizeof(tpl_c1_full), tpl_c1, pub_s);
    free(pub_s);
    if (write_config(config_c1, tpl_c1_full) != 0) { free(pub_c2); return -1; }
    if (config_ensure_keys_and_node_id(config_c1) != 0) { free(pub_c2); return -1; }
    free(pub_c2);

    return 0;
}

static void cleanup_temp_configs(void) {
    test_unlink(config_s);
    test_unlink(config_c1);
    test_unlink(config_c2);
    test_rmdir(temp_dir);
}

static int has_initialized_link(struct UTUN_INSTANCE* inst) {
    if (!inst) return 0;
    struct ETCP_CONN* conn = inst->connections;
    while (conn) {
        struct ETCP_LINK* link = conn->links;
        while (link) {
            if (link->initialized) return 1;
            link = link->next;
        }
        conn = conn->next;
    }
    return 0;
}

static void test_timeout_cb(void* arg) {
    (void)arg;
    test_timed_out = 1;
    DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "test_nat_detection: overall test timeout");
}

static void nat_ping_resp_cb(int success, uint16_t avg_rtt, uint8_t count_sent, uint8_t count_ok,
                               uint32_t recv_ip, uint16_t recv_port, void* arg) {
    (void)arg;
    nat_ping_result.done = 1;
    nat_ping_result.success = success;
    nat_ping_result.avg_rtt = avg_rtt;
    nat_ping_result.count_sent = count_sent;
    nat_ping_result.count_ok = count_ok;
    nat_ping_result.recv_ip = recv_ip;
    nat_ping_result.recv_port = recv_port;
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "nat_ping_resp_cb: success=%d avg_rtt=%u sent=%u ok=%u recv_ip=%08x recv_port=%u",
               success, (unsigned)avg_rtt, (unsigned)count_sent, (unsigned)count_ok,
               (unsigned)recv_ip, (unsigned)recv_port);
}

int main(void) {
    int test_result = 1;
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_INFO);
    debug_set_categories(DEBUG_CATEGORY_ETCP | DEBUG_CATEGORY_BGP | DEBUG_CATEGORY_ROUTING);
    utun_instance_set_tun_init_enabled(0);

    if (create_temp_configs() != 0) {
        fprintf(stderr, "Failed to create temp configs\n");
        return 1;
    }

    ua = uasync_create();
    if (!ua) {
        cleanup_temp_configs();
        return 1;
    }

    inst_s = utun_instance_create(ua, config_s);
    inst_c1 = utun_instance_create(ua, config_c1);
    inst_c2 = utun_instance_create(ua, config_c2);
    if (!inst_s || !inst_c1 || !inst_c2) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create instances");
        goto cleanup;
    }

    if (init_connections(inst_s) != 0 || init_connections(inst_c1) != 0 || init_connections(inst_c2) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to init connections");
        goto cleanup;
    }

    test_timeout_id = uasync_set_timeout(ua, TEST_TIMEOUT_MS * 10, NULL, test_timeout_cb);

    // 1. Wait for links C1->S and C2->S to initialize
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Waiting for ETCP links to initialize...");
    // Give some time for connections to establish before checking
    for (int i = 0; i < 100; i++) {
        uasync_poll(ua, 10);
    }
    while (!test_timed_out) {
        if (has_initialized_link(inst_c1) && has_initialized_link(inst_c2)) {
            DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Links C1->S and C2->S initialized");
            break;
        }
        uasync_poll(ua, 10);
    }
    if (test_timed_out) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Timeout waiting for links");
        goto cleanup;
    }

    // 2. Wait for BGP exchange so S learns about C1 and C2
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "Waiting for BGP exchange (S learns C1 and C2)...");
    int bgp_wait_cycles = 0;
    while (!test_timed_out && bgp_wait_cycles < 500) {
        if (inst_s->bgp && route_bgp_get_node(inst_s->bgp, NODE_ID_C1) != NULL &&
            route_bgp_get_node(inst_s->bgp, NODE_ID_C2) != NULL) {
            DEBUG_INFO(DEBUG_CATEGORY_BGP, "S learned about C1 and C2");
            break;
        }
        uasync_poll(ua, 10);
        bgp_wait_cycles++;
    }
    if (!inst_s->bgp || route_bgp_get_node(inst_s->bgp, NODE_ID_C1) == NULL ||
        route_bgp_get_node(inst_s->bgp, NODE_ID_C2) == NULL) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "S did not learn about C1/C2 in time");
        goto cleanup;
    }

    // 3. Wait for NAT detection to complete for C1 on server
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "Waiting for NAT detection to complete for C1...");
    bgp_wait_cycles = 0;
    while (!test_timed_out && bgp_wait_cycles < 1500) {
        struct NODEINFO_Q* node_c1 = route_bgp_get_node(inst_s->bgp, NODE_ID_C1);
        if (node_c1 && node_c1->nat_check_status == NAT_CHECK_OPEN) {
            DEBUG_INFO(DEBUG_CATEGORY_BGP, "NAT detection completed for C1: OPEN");
            break;
        }
        uasync_poll(ua, 10);
        bgp_wait_cycles++;
    }
    struct NODEINFO_Q* node_c1 = route_bgp_get_node(inst_s->bgp, NODE_ID_C1);
    if (!node_c1 || node_c1->nat_check_status != NAT_CHECK_OPEN) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "NAT detection did not complete for C1");
        goto cleanup;
    }

    // 4. Wait for C1 to receive NAT_INFO
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "Waiting for C1 to receive NAT_INFO...");
    bgp_wait_cycles = 0;
    while (!test_timed_out && bgp_wait_cycles < 500) {
        struct ETCP_SOCKET* sock = inst_c1->etcp_sockets;
        int found = 0;
        while (sock) {
            if (sock->nat_type == NAT_TYPE_OPEN) { found = 1; break; }
            sock = sock->next;
        }
        if (found) {
            DEBUG_INFO(DEBUG_CATEGORY_BGP, "C1 received NAT_INFO");
            break;
        }
        uasync_poll(ua, 10);
        bgp_wait_cycles++;
    }

    // 5. Verify all NAT fields on server
    node_c1 = route_bgp_get_node(inst_s->bgp, NODE_ID_C1);
    if (!node_c1) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "NODEINFO_Q for C1 disappeared");
        goto cleanup;
    }
    if (node_c1->nat_check_status != NAT_CHECK_OPEN) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "FAIL: nat_check_status=%d, expected OPEN(%d)",
                    (int)node_c1->nat_check_status, NAT_CHECK_OPEN);
        goto cleanup;
    }
    if (node_c1->nat_type != NAT_TYPE_OPEN) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "FAIL: nat_type=%d, expected OPEN(%d)",
                    (int)node_c1->nat_type, NAT_TYPE_OPEN);
        goto cleanup;
    }
    if (node_c1->nat_ip == 0 || node_c1->nat_port == 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "FAIL: nat_ip/port not set (ip=%08x port=%u)",
                    (unsigned)node_c1->nat_ip, (unsigned)node_c1->nat_port);
        goto cleanup;
    }

    // Verify link nat_type on server
    struct ETCP_CONN* conn_sc1 = NULL;
    struct ETCP_LINK* link_sc1 = NULL;
    struct ETCP_CONN* conn = inst_s->connections;
    while (conn) {
        if (conn->peer_node_id == NODE_ID_C1) {
            conn_sc1 = conn;
            link_sc1 = conn->links;
            break;
        }
        conn = conn->next;
    }
    if (!link_sc1 || link_sc1->nat_type != NAT_TYPE_OPEN) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "FAIL: link nat_type not OPEN on server");
        goto cleanup;
    }

    // Verify socket nat_type on C1
    struct ETCP_SOCKET* sock_c1 = NULL;
    struct ETCP_SOCKET* sock = inst_c1->etcp_sockets;
    while (sock) {
        if (sock->nat_type == NAT_TYPE_OPEN) {
            sock_c1 = sock;
            break;
        }
        sock = sock->next;
    }
    if (!sock_c1) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "FAIL: no socket with nat_type OPEN on C1");
        goto cleanup;
    }

    // 6. Separate STUN check via explicit route_ping_send_req_addr
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "Performing explicit STUN ping from S via C2 to C1...");
    struct ETCP_CONN* conn_sc2 = NULL;
    conn = inst_s->connections;
    while (conn) {
        if (conn->peer_node_id == NODE_ID_C2) { conn_sc2 = conn; break; }
        conn = conn->next;
    }
    if (!conn_sc2) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "FAIL: S has no connection to C2");
        goto cleanup;
    }

    memset(&nat_ping_result, 0, sizeof(nat_ping_result));
    int ret = route_ping_send_req_addr(inst_s->bgp, conn_sc2, NODE_ID_C1,
                                        node_c1->nat_ip, node_c1->nat_port,
                                        3, 10, 200, 3000, nat_ping_resp_cb, NULL,
                                        node_c1->node.public_key);
    if (ret != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "route_ping_send_req_addr failed: %d", ret);
        goto cleanup;
    }

    bgp_wait_cycles = 0;
    while (!test_timed_out && bgp_wait_cycles < 500 && !nat_ping_result.done) {
        uasync_poll(ua, 10);
        bgp_wait_cycles++;
    }
    if (!nat_ping_result.done) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "FAIL: STUN ping timeout");
        goto cleanup;
    }
    if (!nat_ping_result.success || nat_ping_result.count_ok == 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "FAIL: STUN ping failed success=%d ok=%u",
                    nat_ping_result.success, (unsigned)nat_ping_result.count_ok);
        goto cleanup;
    }
    if (nat_ping_result.recv_ip == 0 || nat_ping_result.recv_port == 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_BGP, "FAIL: STUN recv_ip/port zero (ip=%08x port=%u)",
                    (unsigned)nat_ping_result.recv_ip, (unsigned)nat_ping_result.recv_port);
        goto cleanup;
    }
    // recv_ip should match C2's socket address (127.0.0.1)
    if (nat_ping_result.recv_ip != 0x7F000001) {
        DEBUG_WARN(DEBUG_CATEGORY_BGP, "STUN recv_ip=%08x (expected 127.0.0.1), tolerating",
                   (unsigned)nat_ping_result.recv_ip);
    }

    DEBUG_INFO(DEBUG_CATEGORY_BGP, "STUN check PASSED: recv_ip=%08x recv_port=%u",
               (unsigned)nat_ping_result.recv_ip, (unsigned)nat_ping_result.recv_port);

    DEBUG_INFO(DEBUG_CATEGORY_BGP, "NAT detection test PASSED");
    test_result = 0;

cleanup:
    if (test_timeout_id) uasync_cancel_timeout(ua, test_timeout_id);
    if (inst_s) utun_instance_destroy(inst_s);
    if (inst_c1) utun_instance_destroy(inst_c1);
    if (inst_c2) utun_instance_destroy(inst_c2);
    if (ua) uasync_destroy(ua, 0);
    cleanup_temp_configs();
    return test_result;
}