#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <time.h>

#include "../lib/memory_pool.h"
#include "../lib/u_async.h"
#include "../lib/ll_queue.h"
#include "../tinycrypt/lib/include/tinycrypt/aes.h"
#include "../tinycrypt/lib/include/tinycrypt/ccm_mode.h"
#include "../tinycrypt/lib/include/tinycrypt/ecc.h"
#include "../tinycrypt/lib/include/tinycrypt/ecc_dh.h"
#include "../tinycrypt/lib/include/tinycrypt/constants.h"
#include "../tinycrypt/lib/include/tinycrypt/sha256.h"

#include "../src/utun_instance.h"
#include "../src/config_parser.h"
#include "../src/routing.h"
#include "../src/etcp.h"
#include "../src/etcp_connections.h"
#include "../src/secure_channel.h"
#include "../src/crc32.h"

struct test_context {
    struct UTUN_INSTANCE* server;
    struct UTUN_INSTANCE* client;
    int init_received_server;
    int init_received_client;
    int response_received_server;
    int response_received_client;
    int test_completed;
};

static void print_hex(const char* label, const uint8_t* data, size_t len) {
    printf("%s: ", label);
    for (size_t i = 0; i < len && i < 32; i++) {
        printf("%02x", data[i]);
    }
    printf("\n");
}

static int wait_for_init_completion(struct test_context* ctx, int timeout_ms) {
    int elapsed = 0;
    while (elapsed < timeout_ms && !ctx->test_completed) {
        uasync_poll(ctx->client->ua, 10);
        uasync_poll(ctx->server->ua, 10);
        usleep(10000);
        elapsed += 10;
    }
    return ctx->test_completed;
}

static void check_connections(struct test_context* ctx) {
    struct ETCP_CONN* conn = ctx->client->connections;
    while (conn) {
        struct ETCP_LINK* link = conn->links;
        while (link) {
            if (link->is_server == 0) {
                printf("Client link: initialized=%d, retry_count=%d, timer=%s\n",
                       link->initialized, link->init_retry_count, 
                       link->init_timer ? "active" : "null");
                if (link->initialized) {
                    printf("SUCCESS: Client connection established!\n");
                    ctx->test_completed = 1;
                }
            }
            link = link->next;
        }
        conn = conn->next;
    }

    conn = ctx->server->connections;
    while (conn) {
        struct ETCP_LINK* link = conn->links;
        while (link) {
            if (link->is_server == 1) {
                printf("Server link: initialized=%d\n", link->initialized);
            }
            link = link->next;
        }
        conn = conn->next;
    }
}

int main() {
    printf("=== ETCP Connection Establishment Test ===\n\n");
    
    struct test_context ctx = {0};
    
    ctx.server = calloc(1, sizeof(struct UTUN_INSTANCE));
    ctx.client = calloc(1, sizeof(struct UTUN_INSTANCE));
    
    if (!ctx.server || !ctx.client) {
        printf("Failed to allocate instances\n");
        return 1;
    }
    
    ctx.server->ua = uasync_create();
    ctx.client->ua = uasync_create();
    
    if (!ctx.server->ua || !ctx.client->ua) {
        printf("Failed to create uasync instances\n");
        return 1;
    }
    
    ctx.server->node_id = 0x1111111111111111ULL;
    ctx.client->node_id = 0x2222222222222222ULL;
    
    // Generate keypairs
    sc_generate_keypair(&ctx.server->my_keys);
    sc_generate_keypair(&ctx.client->my_keys);
    
    printf("Server node_id: 0x%llx\n", (unsigned long long)ctx.server->node_id);
    printf("Client node_id: 0x%llx\n", (unsigned long long)ctx.client->node_id);
    print_hex("Server public key", ctx.server->my_keys.public_key, SC_PUBKEY_SIZE);
    print_hex("Client public key", ctx.client->my_keys.public_key, SC_PUBKEY_SIZE);
    
    // Configure server socket (listen on localhost:9001)
    struct ETCP_SOCKET* server_sock = etcp_socket_add(ctx.server, NULL, 0, 0, 0);
    if (!server_sock) {
        printf("Failed to create server socket\n");
        return 1;
    }
    
    struct sockaddr_in server_addr = {0};
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr("127.0.0.1");
    server_addr.sin_port = htons(9001);
    
    if (bind(server_sock->fd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        printf("Failed to bind server socket\n");
        return 1;
    }
    
    // Create server ETCP connection
    struct ETCP_CONN* server_conn = etcp_connection_create(ctx.server);
    if (!server_conn) {
        printf("Failed to create server ETCP connection\n");
        return 1;
    }
    sc_init_ctx(&server_conn->crypto_ctx, &ctx.server->my_keys);
    
    printf("\n=== Starting connection test ===\n");
    
    // Create client socket
    struct ETCP_SOCKET* client_sock = etcp_socket_add(ctx.client, NULL, 0, 0, 0);
    if (!client_sock) {
        printf("Failed to create client socket\n");
        return 1;
    }
    
    // Create client ETCP connection
    struct ETCP_CONN* client_conn = etcp_connection_create(ctx.client);
    if (!client_conn) {
        printf("Failed to create client ETCP connection\n");
        return 1;
    }
    sc_init_ctx(&client_conn->crypto_ctx, &ctx.client->my_keys);
    
    // Set peer public key for both sides
    sc_set_peer_public_key(&server_conn->crypto_ctx, ctx.server->my_keys.public_key, 0);
    sc_set_peer_public_key(&client_conn->crypto_ctx, ctx.client->my_keys.public_key, 0);
    
    // Create client link (this will auto-start init)
    struct sockaddr_storage server_remote_addr = {0};
    memcpy(&server_remote_addr, &server_addr, sizeof(server_addr));
    server_remote_addr.ss_family = AF_INET;
    
    struct ETCP_LINK* client_link = etcp_link_new(client_conn, client_sock, &server_remote_addr, 0);
    if (!client_link) {
        printf("Failed to create client link\n");
        return 1;
    }
    
    printf("\nClient link created, is_server=%d, init_timer=%s\n", 
           client_link->is_server, client_link->init_timer ? "active" : "null");
    
    // Poll for 5 seconds to allow handshake
    printf("\nWaiting for connection establishment (5 seconds)...\n");
    int completed = wait_for_init_completion(&ctx, 5000);
    
    check_connections(&ctx);
    
    if (completed) {
        printf("\n=== TEST PASSED ===\n");
        printf("Connection established successfully!\n");
        return 0;
    } else {
        printf("\n=== TEST FAILED ===\n");
        printf("Connection not established within timeout\n");
        return 1;
    }
}