#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 60000   // 60 seconds for 100 packets
#define PACKET_SIZE 100         // Test packet size
#define TOTAL_PACKETS 100       // Total packets to send
#define MAX_QUEUE_SIZE 5        // Max packets in input queue

static struct UTUN_INSTANCE* server_instance = NULL;
static struct UTUN_INSTANCE* client_instance = NULL;
static int test_completed = 0;
static void* packet_timeout_id = NULL;

// Test statistics
static int packets_sent = 0;
static int packets_received = 0;
static uint8_t packet_buffer[PACKET_SIZE];
static int current_packet_seq = 0;
static uint8_t received_packets[TOTAL_PACKETS];

// Function to generate packet data
static void generate_packet_data(int seq, uint8_t* buffer, int size) {
    buffer[0] = (uint8_t)(seq & 0xFF);
    for (int i = 1; i < size; i++) {
        buffer[i] = (uint8_t)((seq + i) % 256);
    }
}

// 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;
}

// Send packets while queue has space
static void send_packets(void) {
    if (!client_instance || packets_sent >= TOTAL_PACKETS) return;
    
    struct ETCP_CONN* conn = client_instance->connections;
    if (!conn || !conn->input_queue) return;
    
    // Send while queue has space
    while (packets_sent < TOTAL_PACKETS) {
        int queue_count = queue_entry_count(conn->input_queue);
        
        if (queue_count >= MAX_QUEUE_SIZE) {
            // Queue full, stop sending
            break;
        }
        
        generate_packet_data(current_packet_seq, packet_buffer, PACKET_SIZE);
        
        if (etcp_send(conn, packet_buffer, PACKET_SIZE) == 0) {
            packets_sent++;
            current_packet_seq++;
        } else {
            break;  // Send failed
        }
    }
    
    if (packets_sent >= TOTAL_PACKETS) {
        DEBUG_INFO(DEBUG_CATEGORY_ETCP, "All %d packets queued for sending", TOTAL_PACKETS);
    }
}

// Check received packets
static void check_received_packets(void) {
    if (!server_instance) return;
    
    struct ETCP_CONN* conn = server_instance->connections;
    if (!conn || !conn->output_queue) return;
    
    struct ETCP_FRAGMENT* pkt;
    while ((pkt = queue_data_get(conn->output_queue)) != NULL) {
        if (pkt->ll.size >= PACKET_SIZE) {
            int seq = pkt->pkt_data[0];
            
            uint8_t expected[PACKET_SIZE];
            generate_packet_data(seq, expected, PACKET_SIZE);
            
            if (memcmp(pkt->pkt_data, expected, PACKET_SIZE) == 0) {
                if (seq >= 0 && seq < TOTAL_PACKETS) {
                    received_packets[seq] = 1;
                }
                packets_received++;
            }
        }
        
        if (pkt->pkt_data) {
            memory_pool_free(conn->instance->data_pool, pkt->pkt_data);
        }
        queue_data_free(pkt);
    }
}

// Monitor function
static void monitor_and_send(void* arg) {
    (void)arg;
    
    if (test_completed) {
        packet_timeout_id = NULL;
        return;
    }
    
    static int connection_checked = 0;
    
    if (!connection_checked) {
        if (is_connection_established(client_instance)) {
            DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Connection established, starting transmission");
            connection_checked = 1;
        }
    }
    
    if (connection_checked) {
        // Send packets if queue has space
        send_packets();
        
        // Check received packets
        check_received_packets();
        
        // Check completion
        if (packets_sent >= TOTAL_PACKETS && packets_received >= TOTAL_PACKETS) {
            test_completed = 1;
            printf("\n=== SUCCESS: All %d packets transmitted! ===\n", TOTAL_PACKETS);
            if (packet_timeout_id) {
                uasync_cancel_timeout(server_instance->ua, packet_timeout_id);
                packet_timeout_id = NULL;
            }
            return;
        }
    }
    
    if (!test_completed) {
        packet_timeout_id = uasync_set_timeout(server_instance->ua, 10, NULL, monitor_and_send);
    }
}

// Timeout handler
static void test_timeout(void* arg) {
    (void)arg;
    if (!test_completed) {
        printf("\n=== TIMEOUT ===\n");
        printf("Sent: %d/%d, Received: %d/%d\n", packets_sent, TOTAL_PACKETS, packets_received, TOTAL_PACKETS);
        test_completed = 2;
        if (packet_timeout_id) {
            uasync_cancel_timeout(server_instance->ua, packet_timeout_id);
            packet_timeout_id = NULL;
        }
    }
}

int main() {
    printf("=== ETCP 100 Packets Test ===\n\n");
    
    memset(received_packets, 0, sizeof(received_packets));
    
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_DEBUG);
    debug_set_categories(DEBUG_CATEGORY_ETCP);
    
    utun_instance_set_tun_init_enabled(0);

    printf("Creating server...\n");
    struct UASYNC* server_ua = uasync_create();
    server_instance = utun_instance_create(server_ua, "test_server.conf");
    if (!server_instance || init_connections(server_instance) < 0) {
        printf("Failed to create server\n");
        return 1;
    }
    printf("✅ Server ready\n\n");
    
    printf("Creating client...\n");
    struct UASYNC* client_ua = uasync_create();
    client_instance = utun_instance_create(client_ua, "test_client.conf");
    if (!client_instance || init_connections(client_instance) < 0) {
        printf("Failed to create client\n");
        return 1;
    }
    printf("✅ Client ready\n\n");
    
    printf("Sending %d packets (max queue size: %d)...\n", TOTAL_PACKETS, MAX_QUEUE_SIZE);
    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);
    
    int elapsed = 0;
    while (!test_completed && elapsed < TEST_TIMEOUT_MS + 1000) {
        if (server_ua) uasync_poll(server_ua, 5);
        if (client_ua) uasync_poll(client_ua, 5);
        usleep(5000);
        elapsed += 5;
    }
    
    printf("\nCleaning up...\n");
    
    if (packet_timeout_id) uasync_cancel_timeout(server_ua, packet_timeout_id);
    if (global_timeout_id) uasync_cancel_timeout(server_ua, global_timeout_id);
    
    if (server_instance) {
        server_instance->running = 0;
        utun_instance_destroy(server_instance);
    }
    if (client_instance) {
        client_instance->running = 0;
        utun_instance_destroy(client_instance);
    }
    
    if (test_completed == 1) {
        printf("\n=== TEST PASSED ===\n");
        printf("✅ All %d packets transmitted\n", TOTAL_PACKETS);
        return 0;
    } else {
        printf("\n=== TEST FAILED ===\n");
        printf("❌ Sent: %d/%d, Received: %d/%d\n", packets_sent, TOTAL_PACKETS, packets_received, TOTAL_PACKETS);
        return 1;
    }
}