#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 struct UASYNC* ua = NULL;
static int test_completed = 0;
static void* packet_timeout_id = NULL;

// Test statistics - forward direction (client -> server)
static int packets_sent_fwd = 0;
static int packets_received_fwd = 0;
static int current_packet_seq_fwd = 0;
static uint8_t received_packets_fwd[TOTAL_PACKETS];

// Test statistics - backward direction (server -> client)
static int packets_sent_back = 0;
static int packets_received_back = 0;
static int current_packet_seq_back = 0;
static uint8_t received_packets_back[TOTAL_PACKETS];

static uint8_t packet_buffer[PACKET_SIZE];

// Timing variables
static struct timespec start_time_fwd, end_time_fwd;
static struct timespec start_time_back, end_time_back;
static int phase = 0;  // 0 = connecting, 1 = forward transfer, 2 = backward transfer

// 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 from client to server (forward direction)
static void send_packets_fwd(void) {
    if (!client_instance || packets_sent_fwd >= TOTAL_PACKETS) return;
    
    struct ETCP_CONN* conn = client_instance->connections;
    if (!conn || !conn->input_queue) return;
    
    // Start timing on first packet
    if (packets_sent_fwd == 0) {
        clock_gettime(CLOCK_MONOTONIC, &start_time_fwd);
        phase = 1;
        printf("Starting forward transfer (client -> server)...\n");
    }
    
    // Send while queue has space
    while (packets_sent_fwd < TOTAL_PACKETS) {
        int queue_count = queue_entry_count(conn->input_queue);
        
        if (queue_count >= MAX_QUEUE_SIZE) {
            break;
        }
        
        generate_packet_data(current_packet_seq_fwd, packet_buffer, PACKET_SIZE);
        
        if (etcp_send(conn, packet_buffer, PACKET_SIZE) == 0) {
            packets_sent_fwd++;
            current_packet_seq_fwd++;
        } else {
            break;
        }
    }
    
    if (packets_sent_fwd >= TOTAL_PACKETS) {
        DEBUG_INFO(DEBUG_CATEGORY_ETCP, "All %d forward packets queued", TOTAL_PACKETS);
    }
}

// Send packets from server to client (backward direction)
static void send_packets_back(void) {
    if (!server_instance || packets_sent_back >= TOTAL_PACKETS) return;
    
    struct ETCP_CONN* conn = server_instance->connections;
    if (!conn || !conn->input_queue) return;
    
    // Start timing on first packet
    if (packets_sent_back == 0) {
        clock_gettime(CLOCK_MONOTONIC, &start_time_back);
        phase = 2;
        printf("Starting backward transfer (server -> client)...\n");
    }
    
    // Send while queue has space
    while (packets_sent_back < TOTAL_PACKETS) {
        int queue_count = queue_entry_count(conn->input_queue);
        
        if (queue_count >= MAX_QUEUE_SIZE) {
            break;
        }
        
        generate_packet_data(current_packet_seq_back, packet_buffer, PACKET_SIZE);
        
        if (etcp_send(conn, packet_buffer, PACKET_SIZE) == 0) {
            packets_sent_back++;
            current_packet_seq_back++;
        } else {
            break;
        }
    }
    
    if (packets_sent_back >= TOTAL_PACKETS) {
        DEBUG_INFO(DEBUG_CATEGORY_ETCP, "All %d backward packets queued", TOTAL_PACKETS);
    }
}

// Check packets received by server (forward direction)
static void check_received_packets_fwd(void) {
    if (!server_instance) return;
    
    struct ETCP_CONN* conn = server_instance->connections;
    if (!conn || !conn->output_queue) return;
    
    struct ETCP_FRAGMENT* pkt;
    while ((pkt = (struct ETCP_FRAGMENT*)queue_data_get(conn->output_queue)) != NULL) {
        if (pkt->ll.len >= PACKET_SIZE) {
            int seq = pkt->ll.dgram[0];
            
            uint8_t expected[PACKET_SIZE];
            generate_packet_data(seq, expected, PACKET_SIZE);
            
            if (memcmp(pkt->ll.dgram, expected, PACKET_SIZE) == 0) {
                if (seq >= 0 && seq < TOTAL_PACKETS) {
                    received_packets_fwd[seq] = 1;
                }
                packets_received_fwd++;
            }
        }
        
        if (pkt->ll.dgram) {
            memory_pool_free(conn->instance->data_pool, pkt->ll.dgram);
        }
        queue_entry_free((struct ll_entry*)pkt);
    }
}

// Check packets received by client (backward direction)
static void check_received_packets_back(void) {
    if (!client_instance) return;
    
    struct ETCP_CONN* conn = client_instance->connections;
    if (!conn || !conn->output_queue) return;
    
    struct ETCP_FRAGMENT* pkt;
    while ((pkt = (struct ETCP_FRAGMENT*)queue_data_get(conn->output_queue)) != NULL) {
        if (pkt->ll.len >= PACKET_SIZE) {
            int seq = pkt->ll.dgram[0];

            uint8_t expected[PACKET_SIZE];
            generate_packet_data(seq, expected, PACKET_SIZE);

            if (memcmp(pkt->ll.dgram, expected, PACKET_SIZE) == 0) {
                if (seq >= 0 && seq < TOTAL_PACKETS) {
                    received_packets_back[seq] = 1;
                }
                packets_received_back++;
            }
        }

        if (pkt->ll.dgram) {
            memory_pool_free(conn->instance->data_pool, pkt->ll.dgram);
        }
        queue_entry_free((struct ll_entry*)pkt);
    }
}

// Calculate time difference in milliseconds
static double time_diff_ms(struct timespec* start, struct timespec* end) {
    double seconds = end->tv_sec - start->tv_sec;
    double nanoseconds = end->tv_nsec - start->tv_nsec;
    return (seconds * 1000.0) + (nanoseconds / 1000000.0);
}

// 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) {
        // Phase 1: Forward transfer (client -> server)
        if (packets_sent_fwd < TOTAL_PACKETS || packets_received_fwd < TOTAL_PACKETS) {
            send_packets_fwd();
            check_received_packets_fwd();
            
            // Check if forward phase completed
            if (packets_sent_fwd >= TOTAL_PACKETS && packets_received_fwd >= TOTAL_PACKETS) {
                if (end_time_fwd.tv_sec == 0) {
                    clock_gettime(CLOCK_MONOTONIC, &end_time_fwd);
                    double duration = time_diff_ms(&start_time_fwd, &end_time_fwd);
                    printf("✅ Forward transfer completed: %d/%d packets in %.2f ms\n", 
                           packets_received_fwd, TOTAL_PACKETS, duration);
                }
            }
        }
        // Phase 2: Backward transfer (server -> client)
        else if (packets_sent_back < TOTAL_PACKETS || packets_received_back < TOTAL_PACKETS) {
            send_packets_back();
            check_received_packets_back();
        }
        // Check completion
        else {
            clock_gettime(CLOCK_MONOTONIC, &end_time_back);
            double duration_back = time_diff_ms(&start_time_back, &end_time_back);
            double duration_total = time_diff_ms(&start_time_fwd, &end_time_back);
            
            printf("✅ Backward transfer completed: %d/%d packets in %.2f ms\n", 
                   packets_received_back, TOTAL_PACKETS, duration_back);
            
            test_completed = 1;
            printf("\n=== SUCCESS: Bidirectional transfer completed! ===\n");
            printf("Forward (client->server):  %d/%d packets in %.2f ms\n", 
                   packets_received_fwd, TOTAL_PACKETS, time_diff_ms(&start_time_fwd, &end_time_fwd));
            printf("Backward (server->client): %d/%d packets in %.2f ms\n", 
                   packets_received_back, TOTAL_PACKETS, duration_back);
            printf("Total time: %.2f ms\n", duration_total);
            
            if (packet_timeout_id) {
                uasync_cancel_timeout(ua, packet_timeout_id);
                packet_timeout_id = NULL;
            }
            return;
        }
    }
    
    if (!test_completed) {
        packet_timeout_id = uasync_set_timeout(ua, 10, NULL, monitor_and_send);
    }
}

// Timeout handler
static void test_timeout(void* arg) {
    (void)arg;
    if (!test_completed) {
        printf("\n=== TIMEOUT ===\n");
        printf("Forward:  Sent: %d/%d, Received: %d/%d\n", 
               packets_sent_fwd, TOTAL_PACKETS, packets_received_fwd, TOTAL_PACKETS);
        printf("Backward: Sent: %d/%d, Received: %d/%d\n", 
               packets_sent_back, TOTAL_PACKETS, packets_received_back, TOTAL_PACKETS);
        test_completed = 2;
        if (packet_timeout_id) {
            uasync_cancel_timeout(ua, packet_timeout_id);
            packet_timeout_id = NULL;
        }
    }
}

int main() {
    printf("=== ETCP 100 Packets Bidirectional Test ===\n\n");
    
    memset(received_packets_fwd, 0, sizeof(received_packets_fwd));
    memset(received_packets_back, 0, sizeof(received_packets_back));
    
    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");
    ua = uasync_create();
    server_instance = utun_instance_create(ua, "test_etcp_100_packets_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");
    client_instance = utun_instance_create(ua, "test_etcp_100_packets_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 in each direction (max queue size: %d)...\n", TOTAL_PACKETS, MAX_QUEUE_SIZE);
    packet_timeout_id = uasync_set_timeout(ua, 500, NULL, monitor_and_send);
    void* global_timeout_id = uasync_set_timeout(ua, TEST_TIMEOUT_MS, NULL, test_timeout);

    while (!test_completed) {
        uasync_poll(ua, 100);
    }

    printf("\nCleaning up...\n");

    if (packet_timeout_id) uasync_cancel_timeout(ua, packet_timeout_id);
    if (global_timeout_id) uasync_cancel_timeout(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);
    }

    // Destroy shared uasync instance after both instances are destroyed
    if (ua) {
        uasync_destroy(ua, 0);
        ua = NULL;
    }

    if (test_completed == 1) {
        printf("\n=== TEST PASSED ===\n");
        printf("✅ All %d packets transmitted in each direction\n", TOTAL_PACKETS);
        return 0;
    } else {
        printf("\n=== TEST FAILED ===\n");
        printf("❌ Forward:  Sent: %d/%d, Received: %d/%d\n", 
               packets_sent_fwd, TOTAL_PACKETS, packets_received_fwd, TOTAL_PACKETS);
        printf("❌ Backward: Sent: %d/%d, Received: %d/%d\n", 
               packets_sent_back, TOTAL_PACKETS, packets_received_back, TOTAL_PACKETS);
        return 1;
    }
}