// test_etcp_traffic_flow.c - TUN-less ETCP traffic flow debugging test
// This test bypasses TUN devices to focus on ETCP protocol traffic analysis

#include "../src/etcp.h"
#include "../src/etcp_connections.h"
#include "../src/etcp_loadbalancer.h"
#include "../lib/u_async.h"
#include "../lib/debug_config.h"
#include "../lib/memory_pool.h"
#include "../lib/ll_queue.h"
#include "../src/secure_channel.h"
#include "../src/crc32.h"
#include "../src/config_parser.h"
#include "../src/utun_instance.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <time.h>
#include <errno.h>

// Traffic analysis structures
typedef struct {
    uint64_t timestamp;
    uint16_t packet_id;
    uint8_t packet_type;
    uint16_t src_port;
    uint16_t dst_port;
    uint16_t data_len;
    uint8_t* data_copy;  // Copy of packet data for analysis
} packet_capture_t;

typedef struct {
    packet_capture_t* captures;
    size_t max_captures;
    size_t capture_count;
    pthread_mutex_t lock;
} traffic_analyzer_t;

// Test instance configuration
typedef struct {
    struct UTUN_INSTANCE* instance;
    struct ETCP_CONN* etcp_conn;
    struct ETCP_SOCKET* socket;
    traffic_analyzer_t* analyzer;
    uint16_t listen_port;
    uint16_t peer_port;
    uint64_t node_id;
    int is_server;
    pthread_t thread;
    int running;
} test_instance_t;

// Global test state
static test_instance_t* server_instance = NULL;
static test_instance_t* client_instance = NULL;
static int test_completed = 0;
static int connection_established = 0;

// Traffic capture function
static void capture_packet(test_instance_t* inst, struct ETCP_DGRAM* dgram, const char* direction) {
    if (!inst->analyzer || !dgram) return;
    
    pthread_mutex_lock(&inst->analyzer->lock);
    
    if (inst->analyzer->capture_count < inst->analyzer->max_captures) {
        packet_capture_t* cap = &inst->analyzer->captures[inst->analyzer->capture_count++];
        cap->timestamp = get_current_timestamp() * 10; // Convert to 0.1ms units
        cap->data_len = dgram->data_len;
        cap->src_port = inst->listen_port;
        cap->dst_port = inst->peer_port;
        
        // Copy packet data for analysis
        if (dgram->data_len > 0 && dgram->data_len <= 1500) {
            cap->data_copy = malloc(dgram->data_len);
            if (cap->data_copy) {
                memcpy(cap->data_copy, dgram->data, dgram->data_len);
            }
        }
        
        // Analyze packet type from first byte
        if (dgram->data_len > 0) {
            cap->packet_type = dgram->data[0];
            
            // Try to extract packet ID if it's a payload packet
            if (dgram->data_len >= 5 && dgram->data[0] == 0x00) { // ETCP_SECTION_PAYLOAD
                cap->packet_id = (dgram->data[3] << 8) | dgram->data[4];
            } else {
                cap->packet_id = 0xFFFF; // Not a payload packet
            }
        }
        
        DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "CAPTURE %s: inst=%p type=0x%02x id=%u len=%u ts=%llu",
                    direction, inst, cap->packet_type, cap->packet_id, cap->data_len,
                    (unsigned long long)cap->timestamp);
    }
    
    pthread_mutex_unlock(&inst->analyzer->lock);
}

// Enhanced packet input function with traffic analysis
static void test_etcp_input_with_capture(struct ETCP_DGRAM* dgram) {
    if (!dgram || !dgram->link || !dgram->link->etcp) return;
    
    test_instance_t* inst = NULL;
    
    // Determine which instance this packet belongs to
    if (server_instance && dgram->link->etcp == server_instance->etcp_conn) {
        inst = server_instance;
    } else if (client_instance && dgram->link->etcp == client_instance->etcp_conn) {
        inst = client_instance;
    }
    
    if (inst) {
        capture_packet(inst, dgram, "RX");
    }
    
    // Call original ETCP input function
    etcp_conn_input(dgram);
}

// Enhanced packet request function with traffic analysis
static struct ETCP_DGRAM* test_etcp_request_with_capture(struct ETCP_CONN* etcp) {
    struct ETCP_DGRAM* dgram = etcp_request_pkt(etcp);
    
    if (dgram) {
        test_instance_t* inst = NULL;
        
        // Determine which instance this packet belongs to
        if (server_instance && etcp == server_instance->etcp_conn) {
            inst = server_instance;
        } else if (client_instance && etcp == client_instance->etcp_conn) {
            inst = client_instance;
        }
        
        if (inst) {
            capture_packet(inst, dgram, "TX");
        }
    }
    
    return dgram;
}

// Traffic analysis report
static void print_traffic_analysis(test_instance_t* inst, const char* role) {
    if (!inst || !inst->analyzer) return;
    
    pthread_mutex_lock(&inst->analyzer->lock);
    
    printf("\n=== %s Instance Traffic Analysis ===\n", role);
    printf("Total packets captured: %zu\n", inst->analyzer->capture_count);
    
    // Count by packet type
    int payload_count = 0;
    int ack_count = 0;
    int retrans_count = 0;
    int timestamp_count = 0;
    int init_count = 0;
    int other_count = 0;
    
    for (size_t i = 0; i < inst->analyzer->capture_count; i++) {
        packet_capture_t* cap = &inst->analyzer->captures[i];
        
        switch (cap->packet_type) {
            case 0x00: payload_count++; break;      // ETCP_SECTION_PAYLOAD
            case 0x01: ack_count++; break;          // ETCP_SECTION_ACK
            case 0x02: init_count++; break;         // ETCP_INIT_REQUEST
            case 0x03: init_count++; break;         // ETCP_INIT_RESPONSE
            case 0x04: init_count++; break;         // ETCP_CHANNEL_INIT
            case 0x05: init_count++; break;         // ETCP_CHANNEL_RESPONSE
            case 0x06: timestamp_count++; break;    // ETCP_SECTION_TIMESTAMP
            case 0x10: case 0x11: case 0x12: case 0x13: case 0x14: case 0x15: case 0x16: case 0x17: case 0x18: case 0x19: case 0x1A: case 0x1B: case 0x1C: case 0x1D: case 0x1E: case 0x1F:
                retrans_count++; break; // ETCP_SECTION_RETRANS variants
            default: other_count++; break;
        }
    }
    
    printf("Packet breakdown:\n");
    printf("  Payload packets: %d\n", payload_count);
    printf("  ACK packets: %d\n", ack_count);
    printf("  Retransmission requests: %d\n", retrans_count);
    printf("  Timestamp packets: %d\n", timestamp_count);
    printf("  Initialization packets: %d\n", init_count);
    printf("  Other packets: %d\n", other_count);
    
    // Show packet ID analysis for payload packets
    printf("\nPayload packet IDs:\n");
    int last_id = -1;
    int gap_count = 0;
    
    for (size_t i = 0; i < inst->analyzer->capture_count; i++) {
        packet_capture_t* cap = &inst->analyzer->captures[i];
        if (cap->packet_type == 0x00 && cap->packet_id != 0xFFFF) {
            printf("  ID: %u, Len: %u bytes\n", cap->packet_id, cap->data_len);
            
            if (last_id != -1 && cap->packet_id != last_id + 1) {
                gap_count++;
            }
            last_id = cap->packet_id;
        }
    }
    
    if (gap_count > 0) {
        printf("  WARNING: Found %d gaps in packet ID sequence!\n", gap_count);
    }
    
    // Show timing analysis
    if (inst->analyzer->capture_count > 1) {
        uint64_t first_ts = inst->analyzer->captures[0].timestamp;
        uint64_t last_ts = inst->analyzer->captures[inst->analyzer->capture_count-1].timestamp;
        uint64_t duration = (last_ts > first_ts) ? (last_ts - first_ts) : 1;
        
        printf("\nTiming analysis:\n");
        printf("  Test duration: %llu time units (0.1ms each)\n", (unsigned long long)duration);
        printf("  Average packet rate: %.2f packets/time unit\n", 
               (double)inst->analyzer->capture_count / (double)duration);
    }
    
    pthread_mutex_unlock(&inst->analyzer->lock);
}

// Test instance thread function
static void* test_instance_thread(void* arg) {
    test_instance_t* inst = (test_instance_t*)arg;
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Starting %s instance thread (port %u)", 
                inst->is_server ? "server" : "client", inst->listen_port);
    
    // Main event loop
    while (inst->running) {
        uasync_poll(inst->instance->ua, 10); // 10 time units timeout
        
        // Process any pending ETCP packets
        if (inst->etcp_conn) {
            struct ETCP_DGRAM* dgram = test_etcp_request_with_capture(inst->etcp_conn);
            if (dgram) {
                // Send the packet through the socket
                if (inst->socket && inst->socket->fd >= 0) {
                    struct sockaddr_in peer_addr;
                    memset(&peer_addr, 0, sizeof(peer_addr));
                    peer_addr.sin_family = AF_INET;
                    peer_addr.sin_port = htons(inst->peer_port);
                    peer_addr.sin_addr.s_addr = inet_addr("127.0.0.1");
                    
                    // Note: In real implementation, this would send through etcp_encrypt_send
                    DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "Would send packet to peer port %u (len=%u)", 
                                inst->peer_port, dgram->data_len);
                }
            }
        }
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "%s instance thread exiting", inst->is_server ? "server" : "client");
    return NULL;
}

// Create test instance
static test_instance_t* create_test_instance(uint16_t listen_port, uint16_t peer_port, 
                                           uint64_t node_id, int is_server) {
    test_instance_t* inst = calloc(1, sizeof(test_instance_t));
    if (!inst) return NULL;
    
    inst->listen_port = listen_port;
    inst->peer_port = peer_port;
    inst->node_id = node_id;
    inst->is_server = is_server;
    inst->running = 1;
    
    // Create traffic analyzer
    inst->analyzer = calloc(1, sizeof(traffic_analyzer_t));
    if (inst->analyzer) {
        inst->analyzer->max_captures = 1000;
        inst->analyzer->captures = calloc(inst->analyzer->max_captures, sizeof(packet_capture_t));
        pthread_mutex_init(&inst->analyzer->lock, NULL);
    }
    
    // Create UTUN instance directly without config file
    inst->instance = calloc(1, sizeof(struct UTUN_INSTANCE));
    if (!inst->instance) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to allocate UTUN instance");
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    
    // Initialize the instance manually
    inst->instance->node_id = node_id;
    inst->instance->running = 1;
    inst->instance->log_fp = stderr;
    
    // Initialize crypto keys
    const char* priv_key_hex = "67b705a92b41bcaae105af2d6a17743faa7b26ccebba8b3b9b0af05e9cd1d5fb";
    const char* pub_key_hex = "1c55e4ccae7c4470707759086738b10681bf88b81f198cc2ab54a647d1556e17c65e6b1833e0c771e5a39382c03067c388915a4c732191bc130480f20f8e00b9";
    
    // Parse keys
    for (int i = 0; i < 32; i++) {
        sscanf(&priv_key_hex[i*2], "%2hhx", &inst->instance->my_keys.private_key[i]);
    }
    for (int i = 0; i < 64; i++) {
        sscanf(&pub_key_hex[i*2], "%2hhx", &inst->instance->my_keys.public_key[i]);
    }
    
    // Create uasync context
    inst->instance->ua = uasync_create();
    if (!inst->instance->ua) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create uasync context");
        free(inst->instance);
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    
    // Initialize uasync instance
    uasync_init_instance(inst->instance->ua);
    
    // Create packet pool
    inst->instance->pkt_pool = memory_pool_init(1600); // 1600 byte packets
    if (!inst->instance->pkt_pool) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create packet pool");
        uasync_destroy(inst->instance->ua);
        free(inst->instance);
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    if (!inst->instance) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create UTUN instance");
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    
    // Create ETCP connection
    inst->etcp_conn = etcp_connection_create(inst->instance);
    if (!inst->etcp_conn) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create ETCP connection");
        utun_instance_destroy(inst->instance);
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    
    // Create ETCP socket (UDP socket for traffic)
    int sock_fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock_fd < 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create UDP socket");
        etcp_connection_close(inst->etcp_conn);
        utun_instance_destroy(inst->instance);
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    
    struct sockaddr_in local_addr;
    memset(&local_addr, 0, sizeof(local_addr));
    local_addr.sin_family = AF_INET;
    local_addr.sin_port = htons(listen_port);
    local_addr.sin_addr.s_addr = INADDR_ANY;
    
    if (bind(sock_fd, (struct sockaddr*)&local_addr, sizeof(local_addr)) < 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to bind UDP socket to port %u", listen_port);
        close(sock_fd);
        etcp_connection_close(inst->etcp_conn);
        utun_instance_destroy(inst->instance);
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    
    // Create socket structure
    inst->socket = etcp_socket_add(inst->instance, (struct sockaddr_storage*)&local_addr, 0, 0, 0);
    if (!inst->socket) {
        DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to create ETCP socket");
        close(sock_fd);
        etcp_connection_close(inst->etcp_conn);
        utun_instance_destroy(inst->instance);
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Created %s instance (listen=%u, peer=%u, node_id=%llu)",
                is_server ? "server" : "client", listen_port, peer_port, 
                (unsigned long long)node_id);
    
    return inst;
}

// Cleanup test instance
static void destroy_test_instance(test_instance_t* inst) {
    if (!inst) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Destroying %s instance", inst->is_server ? "server" : "client");
    
    inst->running = 0;
    
    if (inst->thread) {
        pthread_join(inst->thread, NULL);
    }
    
    if (inst->etcp_conn) {
        etcp_connection_close(inst->etcp_conn);
    }
    
    if (inst->instance) {
        if (inst->instance->pkt_pool) {
            // Memory pool will be cleaned up with instance
        }
        if (inst->instance->ua) {
            uasync_destroy(inst->instance->ua);
        }
        free(inst->instance);
    }
    
    if (inst->analyzer) {
        // Free captured packet data
        for (size_t i = 0; i < inst->analyzer->capture_count; i++) {
            if (inst->analyzer->captures[i].data_copy) {
                free(inst->analyzer->captures[i].data_copy);
            }
        }
        free(inst->analyzer->captures);
        pthread_mutex_destroy(&inst->analyzer->lock);
        free(inst->analyzer);
    }
    
    free(inst);
}

// Main test function
int main(int argc, char* argv[]) {
    printf("=== ETCP Traffic Flow Debugging Test ===\n");
    
    // Initialize debug system
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_DEBUG);
    debug_enable_category(DEBUG_CATEGORY_ETCP);
    debug_enable_timestamp(1);
    debug_enable_function_name(1);
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Starting ETCP traffic flow debugging test");
    
    // Create server instance - use dynamic high ports to avoid conflicts
    uint16_t server_port = 30000 + (getpid() % 1000);  // Dynamic port based on PID
    uint16_t client_port = server_port + 1;
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Creating server instance on port %u...", server_port);
    server_instance = create_test_instance(server_port, client_port, 0x1111111111111111ULL, 1);
    if (!server_instance) {
        fprintf(stderr, "Failed to create server instance on port %u\n", server_port);
        return 1;
    }
    
    // Create client instance
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Creating client instance on port %u...", client_port);
    client_instance = create_test_instance(client_port, server_port, 0x2222222222222222ULL, 0);
    if (!client_instance) {
        fprintf(stderr, "Failed to create client instance\n");
        destroy_test_instance(server_instance);
        return 1;
    }
    
    // Start instance threads
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Starting instance threads...");
    pthread_create(&server_instance->thread, NULL, test_instance_thread, server_instance);
    pthread_create(&client_instance->thread, NULL, test_instance_thread, client_instance);
    
    // Let instances run for a bit to establish connection
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Running test for 5 seconds...");
    sleep(5);
    
    // Initiate connection from client
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Initiating connection from client...");
    
    // Create initial connection request
    struct ETCP_LINK* client_link = etcp_link_new(client_instance->etcp_conn, client_instance->socket, 
                                                   (struct sockaddr_storage*)&(struct sockaddr_in){
                                                       .sin_family = AF_INET,
                                                       .sin_port = htons(9001),
                                                       .sin_addr.s_addr = inet_addr("127.0.0.1")
                                                   }, 0);
    if (client_link) {
        DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Client link created successfully");
        
        // Send initial packet to trigger connection
        struct ETCP_DGRAM* init_dgram = test_etcp_request_with_capture(client_instance->etcp_conn);
        if (init_dgram) {
            DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Initial connection packet requested");
        }
    }
    
    // Run for another 5 seconds to see traffic flow
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Continuing test for 5 more seconds...");
    sleep(5);
    
    // Generate some test traffic
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Generating test traffic...");
    
    // Create test packets and queue them
    for (int i = 0; i < 10; i++) {
        if (client_instance->etcp_conn) {
            struct ll_entry* entry = queue_entry_new(100); // 100 byte packets
            if (entry) {
                // Fill with test data
                uint8_t* data = (uint8_t*)ll_entry_data(entry);
                for (int j = 0; j < 100; j++) {
                    data[j] = (uint8_t)(i * 10 + j); // Pattern data
                }
                queue_entry_put(client_instance->etcp_conn->input_queue, entry);
                DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "Queued test packet %d", i);
            }
        }
    }
    
    // Run for final 5 seconds
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Final test phase - 5 seconds...");
    sleep(5);
    
    // Stop instances
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Stopping test instances...");
    server_instance->running = 0;
    client_instance->running = 0;
    
    // Wait for threads to finish
    pthread_join(server_instance->thread, NULL);
    pthread_join(client_instance->thread, NULL);
    
    // Print traffic analysis
    print_traffic_analysis(server_instance, "Server");
    print_traffic_analysis(client_instance, "Client");
    
    // Cleanup
    destroy_test_instance(server_instance);
    destroy_test_instance(client_instance);
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Test completed");
    printf("\n=== Test completed ===\n");
    
    return 0;
}