// test_routing_mesh.c - Test routing between 3 instances in mesh topology
// Uses programmatic config creation (no temp files)

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#ifdef _WIN32
    #include <winsock2.h>
    #include <ws2tcpip.h>
#else
    #include <arpa/inet.h>
#endif
#include "test_utils.h"
#include "../lib/platform_compat.h"
#include "../src/etcp.h"
#include "../src/etcp_connections.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 "../lib/u_async.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"

#define TEST_TIMEOUT_MS 10000   // 10 seconds timeout

// Keys from test_etcp_two_instances.c (known working keys)
// Instance A keys
#define KEY_A_PRIV "67b705a92b41bcaae105af2d6a17743faa7b26ccebba8b3b9b0af05e9cd1d5fb"
#define KEY_A_PUB  "1c55e4ccae7c4470707759086738b10681bf88b81f198cc2ab54a647d1556e17c65e6b1833e0c771e5a39382c03067c388915a4c732191bc130480f20f8e00b9"

// Instance B keys  
#define KEY_B_PRIV "4813d31d28b7e9829247f488c6be7672f2bdf61b2508333128e386d1759afed2"
#define KEY_B_PUB  "c594f33c91f3a2222795c2c110c527bf214ad1009197ce14556cb13df3c461b3c373bed8f205a8dd1fc0c364f90bf471d7c6f5db49564c33e4235d268569ac71"

static struct UTUN_INSTANCE* inst_a = NULL;
static struct UTUN_INSTANCE* inst_b = NULL;
static struct UTUN_INSTANCE* inst_c = NULL;
static struct UASYNC* ua = NULL;
static int test_phase = 0;
static void* timeout_id = NULL;

// Helper: fill sockaddr_storage from IP:port string
static int make_sockaddr(const char* ip_port, struct sockaddr_storage* ss) {
    char buf[64];
    strncpy(buf, ip_port, sizeof(buf)-1);
    buf[sizeof(buf)-1] = '\0';
    
    char* colon = strrchr(buf, ':');
    if (!colon) return -1;
    *colon = '\0';
    int port = atoi(colon + 1);
    
    struct sockaddr_in* sin = (struct sockaddr_in*)ss;
    sin->sin_family = AF_INET;
    sin->sin_port = htons(port);
    if (inet_pton(AF_INET, buf, &sin->sin_addr) != 1) return -1;
    return 0;
}

// Helper: create CFG_SERVER
static struct CFG_SERVER* create_server(const char* name, const char* ip_port, uint8_t type) {
    struct CFG_SERVER* srv = u_calloc(1, sizeof(struct CFG_SERVER));
    if (!srv) return NULL;
    strncpy(srv->name, name, MAX_CONN_NAME_LEN-1);
    if (make_sockaddr(ip_port, &srv->ip) < 0) {
        u_free(srv);
        return NULL;
    }
    srv->type = type;
    return srv;
}

// Helper: create CFG_CLIENT_LINK
static struct CFG_CLIENT_LINK* create_link(struct CFG_SERVER* local_srv, const char* remote_ip_port) {
    struct CFG_CLIENT_LINK* link = u_calloc(1, sizeof(struct CFG_CLIENT_LINK));
    if (!link) return NULL;
    link->local_srv = local_srv;
    strncpy(link->server_name, local_srv->name, MAX_CONN_NAME_LEN-1);
    if (make_sockaddr(remote_ip_port, &link->remote_addr) < 0) {
        u_free(link);
        return NULL;
    }
    return link;
}

// Helper: create CFG_CLIENT
static struct CFG_CLIENT* create_client(const char* name, const char* peer_key, int keepalive) {
    struct CFG_CLIENT* cli = u_calloc(1, sizeof(struct CFG_CLIENT));
    if (!cli) return NULL;
    strncpy(cli->name, name, MAX_CONN_NAME_LEN-1);
    strncpy(cli->peer_public_key_hex, peer_key, MAX_KEY_LEN-1);
    cli->keepalive = keepalive;
    return cli;
}

// Helper: add link to client
static void client_add_link(struct CFG_CLIENT* cli, struct CFG_CLIENT_LINK* link) {
    link->next = cli->links;
    cli->links = link;
}

// Helper: add server to list (at end to preserve order)
static void add_server(struct utun_config* cfg, struct CFG_SERVER* srv) {
    if (!cfg->servers) {
        cfg->servers = srv;
    } else {
        struct CFG_SERVER* tail = cfg->servers;
        while (tail->next) tail = tail->next;
        tail->next = srv;
    }
}

// Helper: find server by name
static struct CFG_SERVER* find_server(struct utun_config* cfg, const char* name) {
    struct CFG_SERVER* srv = cfg->servers;
    while (srv) {
        if (strcmp(srv->name, name) == 0) return srv;
        srv = srv->next;
    }
    return NULL;
}

// Helper: add client to list  
static void add_client(struct utun_config* cfg, struct CFG_CLIENT* cli) {
    cli->next = cfg->clients;
    cfg->clients = cli;
}

// Helper: add subnet
static void add_subnet(struct CFG_ROUTE_ENTRY** list, const char* cidr) {
    struct CFG_ROUTE_ENTRY* entry = u_calloc(1, sizeof(struct CFG_ROUTE_ENTRY));
    if (!entry) return;
    
    char buf[32];
    strncpy(buf, cidr, sizeof(buf)-1);
    char* slash = strchr(buf, '/');
    if (slash) {
        *slash = '\0';
        entry->netmask = atoi(slash + 1);
    }
    entry->ip.family = AF_INET;
    inet_pton(AF_INET, buf, &entry->ip.addr.v4);
    
    entry->next = *list;
    *list = entry;
}

// Create config for instance A
static struct utun_config* create_config_a(const char* key_c_pub) {
    struct utun_config* cfg = u_calloc(1, sizeof(struct utun_config));
    if (!cfg) return NULL;
    
    // Global settings
    strncpy(cfg->global.my_private_key_hex, KEY_A_PRIV, MAX_KEY_LEN-1);
    strncpy(cfg->global.my_public_key_hex, KEY_A_PUB, MAX_KEY_LEN-1);
    cfg->global.my_node_id = 0xAAAAAAAAAAAAAAAAULL;
    strncpy(cfg->global.tun_ifname, "tun_a", 15);
    inet_pton(AF_INET, "10.99.1.1", &cfg->global.tun_ip.addr.v4);
    cfg->global.tun_ip.family = AF_INET;
    cfg->global.tun_test_mode = 1;
    cfg->global.mtu = 1500;
    
    // Servers (for others to connect to)
    add_server(cfg, create_server("b", "127.0.0.1:9101", CFG_SERVER_TYPE_PUBLIC));
    add_server(cfg, create_server("c", "127.0.0.1:9102", CFG_SERVER_TYPE_PUBLIC));
    
    // Clients (to connect to others) - find server by name for correct mapping
    struct CFG_CLIENT* cli_b = create_client("to_b", KEY_B_PUB, 1);
    client_add_link(cli_b, create_link(find_server(cfg, "b"), "127.0.0.1:9201")); // connect to B's server
    add_client(cfg, cli_b);
    
    struct CFG_CLIENT* cli_c = create_client("to_c", key_c_pub, 1);
    client_add_link(cli_c, create_link(find_server(cfg, "c"), "127.0.0.1:9302")); // connect to C's server
    add_client(cfg, cli_c);
    
    // Subnets
    add_subnet(&cfg->my_subnets, "192.168.10.0/24");
    add_subnet(&cfg->my_subnets, "192.168.11.0/24");
    
    return cfg;
}

// Create config for instance B
static struct utun_config* create_config_b(const char* key_c_pub) {
    struct utun_config* cfg = u_calloc(1, sizeof(struct utun_config));
    if (!cfg) return NULL;
    
    // Global settings
    strncpy(cfg->global.my_private_key_hex, KEY_B_PRIV, MAX_KEY_LEN-1);
    strncpy(cfg->global.my_public_key_hex, KEY_B_PUB, MAX_KEY_LEN-1);
    cfg->global.my_node_id = 0xBBBBBBBBBBBBBBBBULL;
    strncpy(cfg->global.tun_ifname, "tun_b", 15);
    inet_pton(AF_INET, "10.99.2.1", &cfg->global.tun_ip.addr.v4);
    cfg->global.tun_ip.family = AF_INET;
    cfg->global.tun_test_mode = 1;
    cfg->global.mtu = 1500;
    
    // Servers
    add_server(cfg, create_server("a", "127.0.0.1:9201", CFG_SERVER_TYPE_PUBLIC));
    add_server(cfg, create_server("c", "127.0.0.1:9202", CFG_SERVER_TYPE_PUBLIC));
    
    // Clients - find server by name for correct mapping
    struct CFG_CLIENT* cli_a = create_client("to_a", KEY_A_PUB, 1);
    client_add_link(cli_a, create_link(find_server(cfg, "a"), "127.0.0.1:9101"));
    add_client(cfg, cli_a);
    
    struct CFG_CLIENT* cli_c = create_client("to_c", key_c_pub, 1);
    client_add_link(cli_c, create_link(find_server(cfg, "c"), "127.0.0.1:9303"));
    add_client(cfg, cli_c);
    
    // Subnets
    add_subnet(&cfg->my_subnets, "192.168.20.0/24");
    add_subnet(&cfg->my_subnets, "192.168.21.0/24");
    
    return cfg;
}

// Create config for instance C
static struct utun_config* create_config_c(const char* key_c_priv, const char* key_c_pub) {
    struct utun_config* cfg = u_calloc(1, sizeof(struct utun_config));
    if (!cfg) return NULL;
    
    // Global settings
    strncpy(cfg->global.my_private_key_hex, key_c_priv, MAX_KEY_LEN-1);
    strncpy(cfg->global.my_public_key_hex, key_c_pub, MAX_KEY_LEN-1);
    cfg->global.my_node_id = 0xCCCCCCCCCCCCCCCCULL;
    strncpy(cfg->global.tun_ifname, "tun_c", 15);
    inet_pton(AF_INET, "10.99.3.1", &cfg->global.tun_ip.addr.v4);
    cfg->global.tun_ip.family = AF_INET;
    cfg->global.tun_test_mode = 1;
    cfg->global.mtu = 1500;
    
    // Servers
    add_server(cfg, create_server("a", "127.0.0.1:9302", CFG_SERVER_TYPE_PUBLIC));
    add_server(cfg, create_server("b", "127.0.0.1:9303", CFG_SERVER_TYPE_PUBLIC));
    
    // Clients - find server by name for correct mapping
    struct CFG_CLIENT* cli_a = create_client("to_a", KEY_A_PUB, 1);
    client_add_link(cli_a, create_link(find_server(cfg, "a"), "127.0.0.1:9102"));
    add_client(cfg, cli_a);
    
    struct CFG_CLIENT* cli_b = create_client("to_b", KEY_B_PUB, 1);
    client_add_link(cli_b, create_link(find_server(cfg, "b"), "127.0.0.1:9202"));
    add_client(cfg, cli_b);
    
    // Subnets
    add_subnet(&cfg->my_subnets, "192.168.30.0/24");
    add_subnet(&cfg->my_subnets, "192.168.31.0/24");
    
    return cfg;
}

// Free config (simplified - just free top level, rest will be cleaned on exit)
static void free_config_manual(struct utun_config* cfg) {
    (void)cfg; // Config structures cleaned up by utun_instance_destroy
    // Note: utun_instance_create_from_config makes copies of config data,
    // so we don't need to keep the config structures after instance creation
}

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

static void check_connections(void* arg) {
    (void)arg;
    
    int a = count_initialized_links(inst_a);
    int b = count_initialized_links(inst_b);
    int c = count_initialized_links(inst_c);
    
    printf("[CONNECT] Links ready: A=%d, B=%d, C=%d (need 2 each)\n", a, b, c);
    
    if (a >= 2 && b >= 2 && c >= 2) {
        printf("\n[SUCCESS] All connections established!\n");
        test_phase = 1;
        return;
    }
    
    timeout_id = uasync_set_timeout(ua, 100, NULL, check_connections);
}

static void timeout_handler(void* arg) {
    (void)arg;
    printf("\n[TIMEOUT] Test failed\n");
    test_phase = -1;
}

int main(void) {
    printf("========================================\n");
    printf("  Routing Mesh Test (3 instances)\n");
    printf("  (programmatic config, no temp files)\n");
    printf("========================================\n\n");
    
    // Generate keys for C
    struct SC_MYKEYS keys_c;
    if (sc_generate_keypair(&keys_c) != SC_OK) {
        fprintf(stderr, "Failed to generate keys for C\n");
        return 1;
    }
    
    // Convert to hex
    char key_c_priv[65], key_c_pub[129];
    const char* hex_chars = "0123456789abcdef";
    for (int i = 0; i < 32; i++) {
        key_c_priv[i*2] = hex_chars[(keys_c.private_key[i] >> 4) & 0xF];
        key_c_priv[i*2+1] = hex_chars[keys_c.private_key[i] & 0xF];
    }
    key_c_priv[64] = '\0';
    
    for (int i = 0; i < 64; i++) {
        key_c_pub[i*2] = hex_chars[(keys_c.public_key[i] >> 4) & 0xF];
        key_c_pub[i*2+1] = hex_chars[keys_c.public_key[i] & 0xF];
    }
    key_c_pub[128] = '\0';
    
    printf("Generated C keys:\n");
    printf("  priv: %.16s...\n", key_c_priv);
    printf("  pub:  %.16s...\n\n", key_c_pub);
    
    // Create configs programmatically
    struct utun_config* cfg_a = create_config_a(key_c_pub);
    struct utun_config* cfg_b = create_config_b(key_c_pub);
    struct utun_config* cfg_c = create_config_c(key_c_priv, key_c_pub);
    
    if (!cfg_a || !cfg_b || !cfg_c) {
        fprintf(stderr, "Failed to create configs\n");
        return 1;
    }
    
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_WARN);
    
    ua = uasync_create();
    if (!ua) {
        fprintf(stderr, "Failed to create uasync\n");
        return 1;
    }
    
    printf("[INIT] Creating instances from programmatic config...\n");
    
    inst_a = utun_instance_create_from_config(ua, cfg_a);
    inst_b = utun_instance_create_from_config(ua, cfg_b);
    inst_c = utun_instance_create_from_config(ua, cfg_c);
    
    if (!inst_a || !inst_b || !inst_c) {
        fprintf(stderr, "Failed to create instances\n");
        return 1;
    }
    
    printf("  Instance A: node_id=%016llX\n", (unsigned long long)inst_a->node_id);
    printf("  Instance B: node_id=%016llX\n", (unsigned long long)inst_b->node_id);
    printf("  Instance C: node_id=%016llX\n", (unsigned long long)inst_c->node_id);
    
    printf("\n[INIT] Initializing connections...\n");
    
    if (init_connections(inst_a) < 0 ||
        init_connections(inst_b) < 0 ||
        init_connections(inst_c) < 0) {
        fprintf(stderr, "Failed to initialize connections\n");
        return 1;
    }
    
    printf("  Connections initialized\n");
    
    printf("\n[PHASE] Waiting for connections to establish...\n");
    timeout_id = uasync_set_timeout(ua, TEST_TIMEOUT_MS, NULL, timeout_handler);
    check_connections(NULL);
    
    while (test_phase == 0) {
        uasync_poll(ua, 10);
    }
    
    // Cleanup
    if (inst_a) utun_instance_destroy(inst_a);
    if (inst_b) utun_instance_destroy(inst_b);
    if (inst_c) utun_instance_destroy(inst_c);
    if (ua) uasync_destroy(ua, 0);
    
    // Free configs
    free_config_manual(cfg_a);
    free_config_manual(cfg_b);
    free_config_manual(cfg_c);
    
    if (test_phase == 1) {
        printf("\n✓ TEST PASSED\n");
        return 0;
    } else {
        printf("\n✗ TEST FAILED\n");
        return 1;
    }
}