#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "../lib/platform_compat.h"
#include "test_utils.h"
#ifdef _WIN32
#include <windows.h>
#include <direct.h>
#else
#include <unistd.h>
#endif
#include <time.h>
#include <sys/stat.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"
#include "../lib/mem.h"

extern int pn_packer_send(void* pn, const uint8_t* buffer, size_t size);

#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

// Temp config file paths
static char temp_dir[] = "/tmp/utun_test_XXXXXX";
static char server_config_path[256];
static char client_config_path[256];

// Server config content (uses ports 9041 - unique to avoid conflicts)
static const char* server_config_content =
    "[global]\n"
    "my_node_id=0x1111111111111111\n"
    "my_private_key=67b705a92b41bcaae105af2d6a17743faa7b26ccebba8b3b9b0af05e9cd1d5fb\n"
    "my_public_key=1c55e4ccae7c4470707759086738b10681bf88b81f198cc2ab54a647d1556e17c65e6b1833e0c771e5a39382c03067c388915a4c732191bc130480f20f8e00b9\n"
    "tun_ip=10.99.0.1/24\n"
    "tun_ifname=tun99\n"
    "\n"
    "[server: test]\n"
    "addr=127.0.0.1:9041\n"
    "type=public\n";

// Client config content (uses ports 9042 - unique to avoid conflicts)
static const char* client_config_content =
    "[global]\n"
    "my_node_id=0x2222222222222222\n"
    "my_private_key=4813d31d28b7e9829247f488c6be7672f2bdf61b2508333128e386d1759afed2\n"
    "my_public_key=c594f33c91f3a2222795c2c110c527bf214ad1009197ce14556cb13df3c461b3c373bed8f205a8dd1fc0c364f90bf471d7c6f5db49564c33e4235d268569ac71\n"
    "tun_ip=10.99.0.2/24\n"
    "tun_ifname=tun98\n"
    "\n"
    "[server: test]\n"
    "addr=127.0.0.1:9042\n"
    "type=public\n"
    "\n"
    "[client: test_client]\n"
    "keepalive=1\n"
    "peer_public_key=1c55e4ccae7c4470707759086738b10681bf88b81f198cc2ab54a647d1556e17c65e6b1833e0c771e5a39382c03067c388915a4c732191bc130480f20f8e00b9\n"
    "link=test:127.0.0.1:9041\n";

// Create temp config files
static int create_temp_configs(void) {
    if (test_mkdtemp(temp_dir) != 0) {
        fprintf(stderr, "Failed to create temp directory\n");
        return -1;
    }

    snprintf(server_config_path, sizeof(server_config_path), "%s/server.conf", temp_dir);
    snprintf(client_config_path, sizeof(client_config_path), "%s/client.conf", temp_dir);

    FILE* f = fopen(server_config_path, "w");
    if (!f) {
        fprintf(stderr, "Failed to create server config file\n");
        return -1;
    }
    fprintf(f, "%s", server_config_content);
    fclose(f);

    f = fopen(client_config_path, "w");
    if (!f) {
        fprintf(stderr, "Failed to create client config file\n");
        test_unlink(server_config_path);
        return -1;
    }
    fprintf(f, "%s", client_config_content);
    fclose(f);

    return 0;
}

// Cleanup temp config files
static void cleanup_temp_configs(void) {
    if (server_config_path[0]) test_unlink(server_config_path);
    if (client_config_path[0]) test_unlink(client_config_path);
    if (temp_dir[0]) test_rmdir(temp_dir);
}

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 = u_malloc(size);
        if (!buffer) break;
        
        generate_packet_data(current_packet_seq_fwd, buffer, size);
        
        pn_packer_send(client_pn, buffer, size);
        
        u_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 = u_malloc(size);
        if (!buffer) break;
        
        generate_packet_data(current_packet_seq_back, buffer, size);
        
        pn_packer_send(server_pn, buffer, size);
        
        u_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;
                }
                // Reset callback for test manual polling
                queue_set_callback(client_pn->output, NULL, NULL);
                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;
                }
                // Reset callback for test manual polling
                queue_set_callback(server_pn->output, NULL, NULL);
                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() {
    // Create temp config files first
    if (create_temp_configs() != 0) {
        fprintf(stderr, "Failed to create temporary config files\n");
        return 1;
    }

    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, server_config_path);
    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, client_config_path);
    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_deinit(server_pn);
    }
    if (client_pn) {
        pn_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;
    }

    // Cleanup temp config files
    cleanup_temp_configs();

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