#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 "../src/pkt_normalizer.h"
#include "../lib/u_async.h"
#include "../lib/ll_queue.h"
#include "../lib/debug_config.h"

#define TEST_TIMEOUT_MS 3000  // 3 second timeout
#define TOTAL_PACKETS 100      // Total packets to send
//#define MAX_QUEUE_SIZE 5        // Max packets in input queue
#define MIN_PACKET_SIZE 10      // Minimum packet size
#define MAX_TEST_PACKET_SIZE 2000   // Maximum packet size - must fit in normalizer fragment (mtu-100)

// Packet header size: seq(2) + size(2) + checksum(2) = 6 bytes
#define PACKET_HEADER_SIZE 6

static struct UTUN_INSTANCE* server_instance = NULL;
static struct UTUN_INSTANCE* client_instance = NULL;
static struct PKTNORM* server_pn = NULL;
static struct PKTNORM* client_pn = 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;

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

// Packet sizes for each packet (random)
static int packet_sizes[TOTAL_PACKETS];

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

// Calculate simple checksum (sum of all bytes modulo 65536)
static uint16_t calculate_checksum(const uint8_t* data, int len) {
    uint32_t sum = 0;
    for (int i = 0; i < len; i++) {
        sum += data[i];
    }
    return (uint16_t)(sum & 0xFFFF);
}

// Function to generate packet data with seq, size, checksum and random payload
// Packet format: [seq:2][size:2][checksum:2][payload:N]
static void generate_packet_data(int seq, uint8_t* buffer, int size) {
    // Ensure minimum size for header
    if (size < PACKET_HEADER_SIZE) size = PACKET_HEADER_SIZE;
    
    // Write header
    buffer[0] = (uint8_t)(seq & 0xFF);
    buffer[1] = (uint8_t)((seq >> 8) & 0xFF);
    buffer[2] = (uint8_t)(size & 0xFF);
    buffer[3] = (uint8_t)((size >> 8) & 0xFF);
    
    // Generate random payload (after header)
    int payload_size = size - PACKET_HEADER_SIZE;
    for (int i = 0; i < payload_size; i++) {
        buffer[PACKET_HEADER_SIZE + i] = (uint8_t)(rand() % 256);
    }
    
    // Calculate and write checksum of payload
    uint16_t checksum = calculate_checksum(buffer + PACKET_HEADER_SIZE, payload_size);
    buffer[4] = (uint8_t)(checksum & 0xFF);
    buffer[5] = (uint8_t)((checksum >> 8) & 0xFF);
}

// Verify packet data integrity
// Packet format: [seq:2][size:2][checksum:2][payload:N]
static int verify_packet_data(uint8_t* buffer, int size, int expected_seq) {
    if (size < PACKET_HEADER_SIZE) return 0;
    
    // Parse header
    int seq = buffer[0] | (buffer[1] << 8);
    int pkt_size = buffer[2] | (buffer[3] << 8);
    int stored_checksum = buffer[4] | (buffer[5] << 8);
    
    // Check sequence number (strict order)
    if (seq != expected_seq) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Packet seq mismatch: expected=%d, got=%d", expected_seq, seq);
        return 0;
    }
    
    // Check size
    if (pkt_size != size) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Packet size mismatch: expected=%d, got=%d", size, pkt_size);
        return 0;
    }
    
    // Verify checksum of payload
    int payload_size = size - PACKET_HEADER_SIZE;
    uint16_t calculated_checksum = calculate_checksum(buffer + PACKET_HEADER_SIZE, payload_size);
    if (calculated_checksum != stored_checksum) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Packet checksum mismatch: stored=%d, calculated=%d", 
                    stored_checksum, calculated_checksum);
        return 0;
    }
    
    return 1;
}

// 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) via normalizer
static void send_packets_fwd(void) {
    printf("DEBUG send_packets_fwd: client_instance=%p, client_pn=%p, packets_sent_fwd=%d/%d\n",
           (void*)client_instance, (void*)client_pn, packets_sent_fwd, TOTAL_PACKETS);
    fflush(stdout);
    
    if (!client_instance || !client_pn || packets_sent_fwd >= TOTAL_PACKETS) {
        printf("DEBUG send_packets_fwd: early return (instance=%p, pn=%p, sent=%d)\n",
               (void*)client_instance, (void*)client_pn, packets_sent_fwd);
        fflush(stdout);
        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) via normalizer...\n");
        fflush(stdout);
    }
    
    // Send while we have packets
    while (packets_sent_fwd < TOTAL_PACKETS) {
        int size = packet_sizes[packets_sent_fwd];
        uint8_t* buffer = malloc(size);
        if (!buffer) break;
        
        generate_packet_data(current_packet_seq_fwd, buffer, size);
        
        pn_packer_send(client_pn, buffer, size);
        
        free(buffer);
        packets_sent_fwd++;
        current_packet_seq_fwd++;
    }
    
    if (packets_sent_fwd >= TOTAL_PACKETS) {
        DEBUG_INFO(DEBUG_CATEGORY_ETCP, "All %d forward packets queued to normalizer", TOTAL_PACKETS);
    }
}

// Send packets from server to client (backward direction) via normalizer
static void send_packets_back(void) {
    if (!server_instance || !server_pn || packets_sent_back >= TOTAL_PACKETS) 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) via normalizer...\n");
    }
    
    // Send while we have packets
    while (packets_sent_back < TOTAL_PACKETS) {
        int size = packet_sizes[packets_sent_back];
        uint8_t* buffer = malloc(size);
        if (!buffer) break;
        
        generate_packet_data(current_packet_seq_back, buffer, size);
        
        pn_packer_send(server_pn, buffer, size);
        
        free(buffer);
        packets_sent_back++;
        current_packet_seq_back++;
    }
    
    if (packets_sent_back >= TOTAL_PACKETS) {
        DEBUG_INFO(DEBUG_CATEGORY_ETCP, "All %d backward packets queued to normalizer", TOTAL_PACKETS);
    }
}

// Check packets received by server (forward direction) via normalizer output
static void check_received_packets_fwd(void) {
    if (!server_instance || !server_pn) return;

    // Debug: check output queue count
    int output_count = queue_entry_count(server_pn->output);
    if (output_count > 0) {
        DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Server output queue has %d entries", output_count);
    }

    void* data;
    while ((data = queue_data_get(server_pn->output)) != NULL) {
        struct ll_entry* entry = (struct ll_entry*)data;

        if (entry->len >= PACKET_HEADER_SIZE) {
            // Verify packet with strict sequence checking
            // packets_received_fwd is the next expected sequence number
            if (verify_packet_data(entry->dgram, entry->len, packets_received_fwd)) {
                packets_received_fwd++;
            } else {
                int seq = entry->dgram[0] | (entry->dgram[1] << 8);
                DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Packet verification failed, seq=%d, expected_seq=%d, len=%d",
                           seq, packets_received_fwd, entry->len);
            }
        } else {
            DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Packet too small: len=%d, expected at least %d",
                       entry->len, PACKET_HEADER_SIZE);
        }

        queue_dgram_free(entry);
        queue_entry_free(data);
    }
}

// Check packets received by client (backward direction) via normalizer output
static void check_received_packets_back(void) {
    printf("DEBUG check_received_packets_back: client_instance=%p, client_pn=%p\n", 
           (void*)client_instance, (void*)client_pn);
    fflush(stdout);
    if (!client_instance || !client_pn) return;

    void* data;
    while ((data = queue_data_get(client_pn->output)) != NULL) {
        struct ll_entry* entry = (struct ll_entry*)data;

        if (entry->len >= PACKET_HEADER_SIZE) {
            // Verify packet with strict sequence checking
            // packets_received_back is the next expected sequence number
            if (verify_packet_data(entry->dgram, entry->len, packets_received_back)) {
                packets_received_back++;
            } else {
                int seq = entry->dgram[0] | (entry->dgram[1] << 8);
                DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Packet verification failed, seq=%d, expected_seq=%d, len=%d",
                           seq, packets_received_back, entry->len);
            }
        } else {
            DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Packet too small: len=%d, expected at least %d",
                       entry->len, PACKET_HEADER_SIZE);
        }

        queue_dgram_free(entry);
        queue_entry_free(data);
    }
}

// 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) {
        int client_established = is_connection_established(client_instance);
        int server_established = is_connection_established(server_instance);
        printf("Connection check: client=%d, server=%d\n", client_established, server_established);
        printf("Connections: client=%p, server=%p\n", 
               (void*)client_instance->connections, (void*)server_instance->connections);
        if (client_established) {
            DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Connection established, starting transmission via normalizer");
            connection_checked = 1;
            
            // Initialize normalizers after connection is established
            printf("DEBUG: client_pn=%p, client_instance->connections=%p\n", 
                   (void*)client_pn, (void*)client_instance->connections);
            fflush(stdout);
            if (!client_pn && client_instance->connections) {
                printf("Creating client normalizer...\n");
                fflush(stdout);
                client_pn = pn_init(client_instance->connections);
                if (!client_pn) {
                    printf("Failed to create client normalizer (returned NULL)\n");
                    fflush(stdout);
                    test_completed = 2;
                    return;
                }
                printf("Client normalizer created (frag_size=%d)\n", client_pn->frag_size);
                fflush(stdout);
            } else if (!client_pn) {
                printf("No client connections available\n");
                fflush(stdout);
            } else {
                printf("Client normalizer already exists\n");
                fflush(stdout);
            }
            
            printf("DEBUG: server_pn=%p, server_instance->connections=%p\n",
                   (void*)server_pn, (void*)server_instance->connections);
            fflush(stdout);
            if (!server_pn && server_instance->connections) {
                printf("Creating server normalizer...\n");
                fflush(stdout);
                server_pn = pn_init(server_instance->connections);
                if (!server_pn) {
                    printf("Failed to create server normalizer (returned NULL)\n");
                    fflush(stdout);
                    test_completed = 2;
                    return;
                }
                printf("Server normalizer created (frag_size=%d)\n\n", server_pn->frag_size);
                fflush(stdout);
            } else if (!server_pn) {
                printf("No server connections available\n");
                fflush(stdout);
            } else {
                printf("Server normalizer already exists\n");
                fflush(stdout);
            }
        }
    }
    
    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 via normalizer 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(server_instance->ua, packet_timeout_id);
            packet_timeout_id = NULL;
        }
    }
}

int main() {
    printf("=== PKT Normalizer + ETCP Test ===\n");
    printf("Testing with %d packets of random sizes (%d-%d bytes)\n\n", 
           TOTAL_PACKETS, MIN_PACKET_SIZE, MAX_TEST_PACKET_SIZE);
    
    // Generate random packet sizes
    srand((unsigned)time(NULL));
    int total_bytes = 0;
    for (int i = 0; i < TOTAL_PACKETS; i++) {
        packet_sizes[i] = MIN_PACKET_SIZE + rand() % (MAX_TEST_PACKET_SIZE - MIN_PACKET_SIZE + 1);
        total_bytes += packet_sizes[i];
    }
    printf("Total data to transfer: %d bytes (%.2f KB average per packet)\n\n", 
           total_bytes, (float)total_bytes / TOTAL_PACKETS / 1024);
    
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_WARN);
//    debug_set_level(DEBUG_LEVEL_DEBUG);
//    debug_set_level(DEBUG_LEVEL_TRACE);
    debug_set_categories(DEBUG_CATEGORY_ALL);
    
    utun_instance_set_tun_init_enabled(0);

    printf("Creating server...\n");
    ua = uasync_create();
    server_instance = utun_instance_create(ua, "test_pkt_normalizer_etcp_server.conf");
    if (!server_instance || init_connections(server_instance) < 0) {
        printf("Failed to create server\n");
        return 1;
    }
    printf("Server created, waiting for connection...\n\n");

    printf("Creating client...\n");
    client_instance = utun_instance_create(ua, "test_pkt_normalizer_etcp_client.conf");
    if (!client_instance || init_connections(client_instance) < 0) {
        printf("Failed to create client\n");
        return 1;
    }
    printf("Client created\n\n");

    printf("Sending %d packets in each direction via normalizer...\n", TOTAL_PACKETS);
    packet_timeout_id = uasync_set_timeout(ua, 500, NULL, monitor_and_send);
    void* global_timeout_id = uasync_set_timeout(ua, TEST_TIMEOUT_MS*10, NULL, test_timeout);

    while (!test_completed) {
        uasync_poll(ua, 100);  // 100ms timeout to allow timer processing
    }
    
    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_pn) {
        pn_pair_deinit(server_pn);
    }
    if (client_pn) {
        pn_pair_deinit(client_pn);
    }
    
    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 via normalizer\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;
    }
}