#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.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/ll_queue.h"
#include "../lib/debug_config.h"

#define TEST_TIMEOUT_MS 5000   // 5 seconds for packet transmission
#define PACKET_SIZE 100        // Test packet size

static struct UTUN_INSTANCE* server_instance = NULL;
static struct UTUN_INSTANCE* client_instance = NULL;
static int test_completed = 0;  // 0 = running, 1 = success, 2 = timeout/failure
static void* packet_timeout_id = NULL;

// Test packet data
static uint8_t test_packet_data[PACKET_SIZE];
static int packet_sent = 0;
static int packet_received = 0;

// Function to check if connection is established - UPDATED WITH DEBUG
static int is_connection_established(struct UTUN_INSTANCE* inst) {
    printf("[DEBUG] is_connection_established: Checking instance %p\n", inst);
    fflush(stdout);
    DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "is_connection_established: Checking instance %p", inst);
    
    if (!inst) {
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "is_connection_established: Instance is NULL");
        return 0;
    }
    
    printf("[DEBUG] is_connection_established: Instance has %d connections\n", inst->connections_count);
    fflush(stdout);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "is_connection_established: Instance has %d connections", inst->connections_count);
    
    struct ETCP_CONN* conn = inst->connections;
    while (conn) {
        printf("[DEBUG] is_connection_established: Checking connection %p\n", conn);
        fflush(stdout);
        DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "is_connection_established: Checking connection %p", conn);
        struct ETCP_LINK* link = conn->links;
        while (link) {
            printf("[DEBUG] is_connection_established: Link %p - initialized=%d, is_server=%d\n", 
                   link, link->initialized, link->is_server);
            fflush(stdout);
            DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "is_connection_established: Link %p - initialized=%d, is_server=%d", 
                       link, link->initialized, link->is_server);
            if (link->initialized) {
                printf("[DEBUG] is_connection_established: FOUND initialized link!\n");
                fflush(stdout);
                DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "is_connection_established: FOUND initialized link!");
                return 1;
            }
            link = link->next;
        }
        conn = conn->next;
    }
    
    DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "is_connection_established: No initialized links found");
    return 0;
}

// Function to send test packet to client's input queue
static void send_test_packet(void) {
    printf("[DEBUG] send_test_packet: ENTERING\n");
    fflush(stdout);
    DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "send_test_packet: ENTERING");
    
    if (!client_instance || packet_sent) {
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "send_test_packet: SKIPPING (client_instance=%p, packet_sent=%d)", client_instance, packet_sent);
        return;
    }
    
    struct ETCP_CONN* conn = client_instance->connections;
    if (!conn) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "send_test_packet: No ETCP connection found on client");
        return;
    }
    
    printf("[DEBUG] send_test_packet: Found connection %p\n", conn);
    fflush(stdout);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "send_test_packet: Found connection %p", conn);
    printf("[DEBUG] send_test_packet: connection input_queue=%p\n", conn->input_queue);
    fflush(stdout);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "send_test_packet: connection input_queue=%p", conn->input_queue);
    printf("[DEBUG] send_test_packet: connection output_queue=%p\n", conn->output_queue);
    fflush(stdout);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "send_test_packet: connection output_queue=%p", conn->output_queue);
    printf("[DEBUG] send_test_packet: connection next_tx_id=%u\n", conn->next_tx_id);
    fflush(stdout);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "send_test_packet: connection next_tx_id=%u", conn->next_tx_id);
    
    if (!conn->input_queue) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "send_test_packet: Client input queue is NULL");
        return;
    }
    
    // Create test packet data
    for (int i = 0; i < PACKET_SIZE; i++) {
        test_packet_data[i] = (uint8_t)(i % 256);
    }
    
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "send_test_packet: Created test data, first byte=%02X, last byte=%02X", 
                test_packet_data[0], test_packet_data[PACKET_SIZE-1]);
    
    // Use new etcp_send function - much simpler API
    printf("[DEBUG] send_test_packet: Calling etcp_send with len=%d\n", PACKET_SIZE);
    fflush(stdout);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "send_test_packet: Calling etcp_send with len=%d", PACKET_SIZE);
    int result = etcp_send(conn, test_packet_data, PACKET_SIZE);
    printf("[DEBUG] send_test_packet: etcp_send returned %d\n", result);
    fflush(stdout);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "send_test_packet: etcp_send returned %d", result);
    
    if (result != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "send_test_packet: Failed to send packet via etcp_send");
        return;
    }
    
    packet_sent = 1;
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "send_test_packet: SUCCESS - Test packet sent via etcp_send (%d bytes)", PACKET_SIZE);
}

// Function to check if packet received in server's output queue
static void check_packet_received(void) {
    DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "check_packet_received: ENTERING");
    
    if (!server_instance || packet_received) {
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "check_packet_received: SKIPPING (server_instance=%p, packet_received=%d)", server_instance, packet_received);
        return;
    }
    
    struct ETCP_CONN* conn = server_instance->connections;
    if (!conn) {
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "check_packet_received: No connection on server");
        return;
    }
    
    if (!conn->output_queue) {
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "check_packet_received: No output_queue on server connection");
        return;
    }
    
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "check_packet_received: Checking server connection %p", conn);
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "check_packet_received: Server output_queue count: %d", queue_entry_count(conn->output_queue));
    
    // Check if there's any packet in output queue
    struct ETCP_FRAGMENT* pkt = (struct ETCP_FRAGMENT*)queue_data_get(conn->output_queue);
    if (pkt) {
        DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "check_packet_received: Found packet in output queue");

        // ETCP_FRAGMENT содержит ll_entry в начале, данные в pkt_data
        size_t size = pkt->ll.len;
        uint8_t* data = pkt->ll.dgram;

        DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "check_packet_received: Packet size=%zu, expected=%d", size, PACKET_SIZE);
        DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "check_packet_received: First byte received=%02X, expected=%02X",
                    data[0], test_packet_data[0]);

        // Проверяем только данные (первые PACKET_SIZE байт), размер может включать ETCP заголовки
        if (size >= PACKET_SIZE && memcmp(data, test_packet_data, PACKET_SIZE) == 0) {
            packet_received = 1;
            DEBUG_INFO(DEBUG_CATEGORY_ETCP, "check_packet_received: SUCCESS - Packet received in server output queue (%zu bytes), data matches", size);
        } else {
            DEBUG_WARN(DEBUG_CATEGORY_ETCP, "check_packet_received: Packet received but data differs (expected %d bytes, got %zu)", PACKET_SIZE, size);
        }

        // Освобождаем ETCP_FRAGMENT и данные
        if (pkt->ll.dgram) {
            memory_pool_free(conn->instance->data_pool, pkt->ll.dgram);
        }
        queue_entry_free((struct ll_entry*)pkt);
    } else {
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "check_packet_received: No packet found in output queue");
    }
}

// Monitor connection and send packet when ready
static void monitor_and_send(void* arg) {
    printf("[DEBUG] monitor_and_send: ENTERING (arg=%p)\n", arg);
    fflush(stdout);
    DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "monitor_and_send: ENTERING (arg=%p)", arg);
    (void)arg;
    
    if (test_completed) {
        packet_timeout_id = NULL;
        return;
    }
    
    static int connection_checked = 0;
    static int packet_sent_flag = 0;
    
    // Check connection
    if (!connection_checked) {
        int server_ready = is_connection_established(server_instance);
        int client_ready = is_connection_established(client_instance);
        
        printf("[DEBUG] monitor_and_send: Connection check - server=%d, client=%d\n", server_ready, client_ready);
        fflush(stdout);
        DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "monitor_and_send: Connection check - server=%d, client=%d", server_ready, client_ready);
        
        // We only need client link initialized to send packet
        // Server will accept packets via its socket
        if (client_ready) {
            DEBUG_INFO(DEBUG_CATEGORY_ETCP, "monitor_and_send: Client link initialized, ready to send packet (server=%d, client=%d)", server_ready, client_ready);
            connection_checked = 1;
        } else {
            DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "monitor_and_send: Waiting for connection... (server=%d, client=%d)", server_ready, client_ready);
        }
    }
    
    // Send packet once connection is ready
    if (connection_checked && !packet_sent_flag) {
        send_test_packet();
        packet_sent_flag = 1;
    }
    
    // Check for received packet
    if (packet_sent_flag) {
        check_packet_received();
        if (packet_received) {
            test_completed = 1;  // Success
            printf("\n=== SUCCESS: Packet transmitted successfully! ===\n");
            if (packet_timeout_id) {
                uasync_cancel_timeout(server_instance->ua, packet_timeout_id);
                packet_timeout_id = NULL;
            }
            return;
        }
    }
    
    // Schedule next check
    if (!test_completed) {
        packet_timeout_id = uasync_set_timeout(server_instance->ua, 100, NULL, monitor_and_send);
    }
}

// Timeout handler
static void test_timeout(void* arg) {
    (void)arg;
    if (!test_completed) {
        printf("\n=== TIMEOUT: Packet not received within %d seconds ===\n", TEST_TIMEOUT_MS/1000);
        test_completed = 2;  // Timeout/failure
        if (packet_timeout_id) {
            uasync_cancel_timeout(server_instance->ua, packet_timeout_id);
            packet_timeout_id = NULL;
        }
    }
}

int main() {
    printf("=== ETCP Simple Traffic Test (Queue-based) ===\n\n");
    
    // Enable debug output - MAXIMUM DEBUGGING
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_TRACE);
    debug_set_categories(DEBUG_CATEGORY_ALL);
    debug_enable_function_name(1);

    DEBUG_TRACE(DEBUG_CATEGORY_MEMORY, "*************1");

    
    printf("Function names enabled: YES\n");
    printf("===========================\n\n");
    
    // Explicitly disable TUN initialization for this test
    utun_instance_set_tun_init_enabled(0);

    // Create server instance
    printf("Creating server instance...\n");
    struct UASYNC* server_ua = uasync_create();
    server_instance = utun_instance_create(server_ua, "test_server.conf");
    if (!server_instance) {
        printf("Failed to create server instance\n");
        return 1;
    }

    // Initialize connections and register sockets regardless of TUN state
    if (server_instance->tun.fd < 0) {
        printf("ℹ️  Server TUN disabled - initializing connections only\n");
    }

    // Initialize ETCP connections (creates sockets and links)
    if (init_connections(server_instance) < 0) {
        printf("Failed to initialize server connections\n");
        utun_instance_destroy(server_instance);
        return 1;
    } else {
        printf("Server connections initialized: sockets=%p, connections=%p\n", 
               server_instance->etcp_sockets, server_instance->connections);
    }
    

    
    printf("✅ Server instance initialized successfully (node_id=%llx)\n", (unsigned long long)server_instance->node_id);
    printf("Server instance ready (node_id=%llx)\n\n", (unsigned long long)server_instance->node_id);
    
    // Create client instance
    printf("Creating client instance...\n");
    struct UASYNC* client_ua = uasync_create();
    client_instance = utun_instance_create(client_ua, "test_client.conf");
    if (!client_instance) {
        printf("Failed to create client instance\n");
        utun_instance_destroy(server_instance);
        return 1;
    }
    
    // Initialize connections and register sockets regardless of TUN state
    if (client_instance->tun.fd < 0) {
        printf("ℹ️  Client TUN disabled - initializing connections only\n");
    }
    
    // Initialize ETCP connections (creates sockets and links)
    printf("About to call init_connections() for client instance\n");
    fflush(stdout);
    
    int conn_result = init_connections(client_instance);
    printf("init_connections() returned: %d\n", conn_result);
    fflush(stdout);
    
    if (conn_result < 0) {
        printf("Failed to initialize client connections (result=%d)\n", conn_result);
        printf("But continuing test to analyze the issue...\n");
        fflush(stdout);
        // Don't return error - continue to analyze
    } else {
        printf("Client connections initialized: sockets=%p, connections=%p, count=%d\n", 
               client_instance->etcp_sockets, client_instance->connections,
               client_instance ? client_instance->connections_count : -1);
        fflush(stdout);
    }
    

    
    printf("✅ Client instance initialized successfully (node_id=%llx)\n", (unsigned long long)client_instance->node_id);
    printf("Client instance ready (node_id=%llx)\n\n", (unsigned long long)client_instance->node_id);
    
    // Debug: print connection and link info
    printf("\n=== Connection Debug ===\n");
    struct ETCP_CONN* conn = client_instance->connections;
    while (conn) {
        printf("Client connection %p: peer_node_id=%llx\n", conn, (unsigned long long)conn->peer_node_id);
        struct ETCP_LINK* link = conn->links;
        while (link) {
            printf("  Link %p: initialized=%d, is_server=%d, remote_addr family=%d, conn->fd=%d\n", 
                   link, link->initialized, link->is_server, link->remote_addr.ss_family, link->conn ? link->conn->fd : -1);
            link = link->next;
        }
        conn = conn->next;
    }
    conn = server_instance->connections;
    while (conn) {
        printf("Server connection %p: peer_node_id=%llx\n", conn, (unsigned long long)conn->peer_node_id);
        struct ETCP_LINK* link = conn->links;
        while (link) {
            printf("  Link %p: initialized=%d, is_server=%d\n", link, link->initialized, link->is_server);
            link = link->next;
        }
        conn = conn->next;
    }
    printf("=== End Debug ===\n\n");
    
    // Start monitoring and packet transmission
    printf("Starting packet transmission test...\n");
    packet_timeout_id = uasync_set_timeout(server_ua, 500, NULL, monitor_and_send);
    void* global_timeout_id = uasync_set_timeout(server_ua, TEST_TIMEOUT_MS, NULL, test_timeout);
    
    // Main event loop
    printf("Running event loop...\n\n");
    
    // Give instances time to initialize
    printf("Waiting 0.5 seconds for instances initialization...\n");
    for (int i = 0; i < 50; i++) {
        if (server_ua) uasync_poll(server_ua, 10);
        if (client_ua) uasync_poll(client_ua, 10);
        usleep(10000);
    }
    printf("Starting connection and packet transmission...\n");
    
    int elapsed = 0;
    int poll_interval = 5;
    while (!test_completed && elapsed < TEST_TIMEOUT_MS + 1000) {
        if (server_ua) uasync_poll(server_ua, poll_interval);
        if (client_ua) uasync_poll(client_ua, poll_interval);
        
        usleep(poll_interval * 1000);
        elapsed += poll_interval;
        
        // Quick exit if packet received
        if (test_completed == 1) {
            printf("[TEST] Packet received, exiting early after %d ms\n", elapsed);
            break;
        }
    }
    
    // Cleanup
    printf("\nCleaning up...\n");
    
    // Cancel timeouts
    if (packet_timeout_id) {
        uasync_cancel_timeout(server_ua, packet_timeout_id);
        packet_timeout_id = NULL;
    }
    if (global_timeout_id) {
        uasync_cancel_timeout(server_ua, global_timeout_id);
    }
    
    // Destroy instances
    if (server_instance) {
        server_instance->running = 0;
        utun_instance_destroy(server_instance);
    }
    if (client_instance) {
        client_instance->running = 0;
        utun_instance_destroy(client_instance);
    }
    
    // Evaluate test result
    if (test_completed == 1) {
        printf("\n=== TEST PASSED ===\n");
        printf("✅ Packet successfully transmitted from client input queue to server output queue\n");
        printf("✅ ETCP connection and queue mechanisms verified\n");
        return 0;
    } else if (test_completed == 2) {
        printf("\n=== TEST FAILED: Packet not received within timeout ===\n");
        printf("❌ Connection may not have been established\n");
        printf("❌ Check if UDP sockets were created and bound correctly\n");
        return 1;
    } else {
        printf("\n=== TEST FAILED: Unknown error ===\n");
        return 1;
    }
}