#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
#define MAX_QUEUE_ENTRIES 100  // Max entries in queue before considering it full

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
static int is_connection_established(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;
}

// Function to send test packet via normalizer packer input queue
static void send_test_packet(void) {
    if (!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;
    }
    
    if (!conn->normalizer || !conn->normalizer->packer || !conn->normalizer->packer->input) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "send_test_packet: Client normalizer/packer/input not initialized");
        return;
    }
    
    // Check queue fullness before writing
    size_t queue_count = queue_entry_count(conn->normalizer->packer->input);
    if (queue_count >= MAX_QUEUE_ENTRIES) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "send_test_packet: Packer input queue is full (%zu entries), waiting...", queue_count);
        return;
    }
    
    // Create test packet data
    for (int i = 0; i < PACKET_SIZE; i++) {
        test_packet_data[i] = (uint8_t)(i % 256);
    }
    
    // Allocate memory for packet using queue API
    void* entry_data = queue_data_new(PACKET_SIZE);
    if (!entry_data) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "send_test_packet: Failed to allocate queue data");
        return;
    }
    
    memcpy(entry_data, test_packet_data, PACKET_SIZE);
    
    // Set the length in ll_entry
    struct ll_entry* entry = (struct ll_entry *)((char *)entry_data - offsetof(struct ll_entry, data));
    entry->len = PACKET_SIZE;
    
    // Write to packer input queue
    if (queue_data_put(conn->normalizer->packer->input, entry_data, 0) < 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "send_test_packet: Failed to put data in packer input queue");
        queue_data_free(entry_data);
        return;
    }
    
    packet_sent = 1;
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "send_test_packet: Test packet sent to packer input queue (%d bytes)", PACKET_SIZE);
}

// Function to check if packet received in server's unpacker output queue
static void check_packet_received(void) {
    if (!server_instance || packet_received) {
        return;
    }
    
    struct ETCP_CONN* conn = server_instance->connections;
    if (!conn) {
        return;
    }
    
    if (!conn->normalizer || !conn->normalizer->unpacker || !conn->normalizer->unpacker->output) {
        return;
    }
    
    // Read from unpacker output queue
    void* data = queue_data_get(conn->normalizer->unpacker->output);
    if (!data) {
        return;
    }
    
    // Get the length from ll_entry
    struct ll_entry* entry = (struct ll_entry *)((char *)data - offsetof(struct ll_entry, data));
    size_t size = entry->len;
    
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "check_packet_received: Found packet in unpacker output queue (size=%zu)", size);
    
    // Verify packet data
    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 and data matches (%zu bytes)", size);
    } else {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "check_packet_received: Packet received but data differs (expected %d bytes, got %zu)", PACKET_SIZE, size);
    }
    
    // Free the data
    queue_data_free(data);
}

// Monitor connection and send packet when ready
static void monitor_and_send(void* 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);
        
        if (client_ready) {
            DEBUG_INFO(DEBUG_CATEGORY_ETCP, "monitor_and_send: Client link initialized, ready to send packet");
            connection_checked = 1;
        }
    }
    
    // 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 (via normalizer queues) ===\n\n");
    
    // Enable debug output
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_INFO);
    debug_set_categories(DEBUG_CATEGORY_ETCP | DEBUG_CATEGORY_CONNECTION);
    
    // 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 ETCP connections
    if (init_connections(server_instance) < 0) {
        printf("Failed to initialize server connections\n");
        utun_instance_destroy(server_instance);
        return 1;
    }
    
    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 ETCP connections
    if (init_connections(client_instance) < 0) {
        printf("Failed to initialize client connections\n");
        utun_instance_destroy(server_instance);
        utun_instance_destroy(client_instance);
        return 1;
    }
    
    printf("Client instance ready (node_id=%llx)\n\n", (unsigned long long)client_instance->node_id);
    
    // 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 packer input to server unpacker output\n");
        printf("ETCP + normalizer integration 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");
        return 1;
    } else {
        printf("\n=== TEST FAILED: Unknown error ===\n");
        return 1;
    }
}