// test_etcp_simple_traffic.c - Simplified ETCP traffic flow debugging without sockets
// This test focuses purely on ETCP protocol traffic analysis using direct packet injection

#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 <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <time.h>
#include <errno.h>

// Simplified traffic analysis - focus on ETCP protocol packets only
typedef struct {
    uint64_t timestamp;
    uint16_t packet_id;
    uint8_t packet_type;
    uint16_t data_len;
    uint8_t data[1500];  // Copy of packet data for analysis
} packet_capture_t;

typedef struct {
    packet_capture_t captures[1000];
    size_t capture_count;
    pthread_mutex_t lock;
} traffic_analyzer_t;

// Test instance - simplified version without actual sockets
typedef struct {
    struct UTUN_INSTANCE* instance;
    struct ETCP_CONN* etcp_conn;
    traffic_analyzer_t* analyzer;
    uint64_t node_id;
    int is_server;
    int running;
} test_instance_t;

// Global test state
static test_instance_t* server_instance = NULL;
static test_instance_t* client_instance = NULL;

// Enhanced packet analysis function
static void analyze_etcp_packet(struct ETCP_CONN* etcp, uint8_t* data, uint16_t len, const char* direction) {
    if (!etcp || !data || len == 0) {
        DEBUG_WARN(DEBUG_CATEGORY_ETCP, "analyze_etcp_packet: invalid parameters (etcp=%p, data=%p, len=%u)", 
                  etcp, data, len);
        return;
    }
    
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "PACKET %s: etcp=%p, len=%u, first_byte=0x%02x", 
                direction, etcp, len, data[0]);
    
    // Parse ETCP sections
    uint16_t pos = 0;
    int section_count = 0;
    
    while (pos < len) {
        if (pos + 3 > len) {
            DEBUG_WARN(DEBUG_CATEGORY_ETCP, "Invalid section header at pos %u (need 3 bytes, have %u)", 
                      pos, len - pos);
            break;
        }
        
        uint8_t section_type = data[pos];
        uint16_t section_len = (data[pos+1] << 8) | data[pos+2];
        
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "Section %d: type=0x%02x, len=%u at pos=%u", 
                    section_count, section_type, section_len, pos);
        
        if (pos + 3 + section_len > len || section_len > 1500) {
            DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Section length %u exceeds packet boundary (pos=%u, total=%u) or max size", 
                       section_len, pos, len);
            break;
        }
        
        // Analyze specific section types
        uint8_t* section_data = data + pos + 3;
        
        switch (section_type) {
            case ETCP_SECTION_PAYLOAD: {
                if (section_len >= 2) {
                    uint16_t packet_id = (section_data[0] << 8) | section_data[1];
                    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "  PAYLOAD: packet_id=%u, payload_len=%u", 
                                packet_id, section_len - 2);
                }
                break;
            }
            case ETCP_SECTION_ACK: {
                if (section_len >= 1) {
                    uint8_t ack_count = section_data[0];
                    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "  ACK: count=%u", ack_count);
                    for (int i = 0; i < ack_count && i*4+5 <= section_len; i++) {
                        uint16_t ack_id = (section_data[1+i*4] << 8) | section_data[2+i*4];
                        uint16_t ack_ts = (section_data[3+i*4] << 8) | section_data[4+i*4];
                        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "    ACK[%d]: id=%u, ts=%u", i, ack_id, ack_ts);
                    }
                }
                break;
            }
            case ETCP_SECTION_TIMESTAMP: {
                if (section_len == 2) {
                    uint16_t ts = (section_data[0] << 8) | section_data[1];
                    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "  TIMESTAMP: ts=%u", ts);
                }
                break;
            }
            case ETCP_INIT_REQUEST: {
                if (section_len >= 76) { // node_id(8) + mtu(2) + keepalive(2) + pubkey(64)
                    uint64_t node_id = 0;
                    for (int i = 0; i < 8; i++) {
                        node_id = (node_id << 8) | section_data[i];
                    }
                    uint16_t mtu = (section_data[8] << 8) | section_data[9];
                    uint16_t keepalive = (section_data[10] << 8) | section_data[11];
                    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "  INIT_REQUEST: node_id=%llu, mtu=%u, keepalive=%u",
                              (unsigned long long)node_id, mtu, keepalive);
                }
                break;
            }
            case ETCP_INIT_RESPONSE: {
                if (section_len >= 10) { // node_id(8) + mtu(2)
                    uint64_t node_id = 0;
                    for (int i = 0; i < 8; i++) {
                        node_id = (node_id << 8) | section_data[i];
                    }
                    uint16_t mtu = (section_data[8] << 8) | section_data[9];
                    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "  INIT_RESPONSE: node_id=%llu, mtu=%u",
                              (unsigned long long)node_id, mtu);
                }
                break;
            }
            default: {
                if (section_type >= ETCP_SECTION_RETRANS && section_type <= ETCP_SECTION_RETRANS + 0x1F) {
                    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "  RETRANS: base_id=%u", section_type - ETCP_SECTION_RETRANS);
                    // Parse retransmission requests
                    for (int i = 0; i < section_len; i += 2) {
                        if (i+1 < section_len) {
                            uint16_t req_id = (section_data[i] << 8) | section_data[i+1];
                            DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "    RETRANS_REQ: id=%u", req_id);
                        }
                    }
                } else {
                    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "  UNKNOWN: type=0x%02x, len=%u", section_type, section_len);
                }
                break;
            }
        }
        
        pos += 3 + section_len;
        section_count++;
    }
    
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "Packet analysis complete: %d sections processed", section_count);
}

// Simulate packet reception for analysis
static void simulate_packet_rx(test_instance_t* inst, uint8_t* data, uint16_t len) {
    if (!inst || !inst->etcp_conn || !data) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "SIMULATE RX: inst=%s, len=%u", 
                inst->is_server ? "server" : "client", len);
    
    // Analyze the packet first
    analyze_etcp_packet(inst->etcp_conn, data, len, "SIM_RX");
    
    // Create a mock ETCP_DGRAM for processing
    struct ETCP_DGRAM mock_dgram;
    memset(&mock_dgram, 0, sizeof(mock_dgram));
    mock_dgram.link = NULL; // Will be set by connection processing
    mock_dgram.data_len = len > 1500 ? 1500 : len;
    mock_dgram.timestamp = get_current_timestamp();
    if (mock_dgram.data_len > 0) {
        memcpy(mock_dgram.data, data, mock_dgram.data_len);
    }
    
    // Process through ETCP input
    etcp_conn_input(&mock_dgram);
}

// Simulate packet transmission for analysis
static void simulate_packet_tx(test_instance_t* inst) {
    if (!inst || !inst->etcp_conn) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "SIMULATE TX: inst=%s", 
                inst->is_server ? "server" : "client");
    
    struct ETCP_DGRAM* dgram = etcp_request_pkt(inst->etcp_conn);
    if (dgram) {
        DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "Generated packet: len=%u", dgram->data_len);
        analyze_etcp_packet(inst->etcp_conn, dgram->data, dgram->data_len, "SIM_TX");
        
        // In a real scenario, this would be sent over the network
        // For now, we just analyze it and free it
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "Would send %u bytes to peer", dgram->data_len);
    } else {
        DEBUG_TRACE(DEBUG_CATEGORY_ETCP, "No packets available to send");
    }
}

// Create simplified test instance
static test_instance_t* create_simple_instance(uint64_t node_id, int is_server) {
    test_instance_t* inst = calloc(1, sizeof(test_instance_t));
    if (!inst) return NULL;
    
    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) {
        pthread_mutex_init(&inst->analyzer->lock, NULL);
    }
    
    // Create UTUN instance directly without TUN device
    inst->instance = calloc(1, sizeof(struct UTUN_INSTANCE));
    if (!inst->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 (hardcoded for test)
    const char* priv_key_hex = "67b705a92b41bcaae105af2d6a17743faa7b26ccebba8b3b9b0af05e9cd1d5fb";
    const char* pub_key_hex = "1c55e4ccae7c4470707759086738b10681bf88b81f198cc2ab54a647d1556e17c65e6b1833e0c771e5a39382c03067c388915a4c732191bc130480f20f8e00b9";
    
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "Parsing crypto keys: priv_len=%zu, pub_len=%zu", 
                strlen(priv_key_hex), strlen(pub_key_hex));
    
    // Parse keys
    for (int i = 0; i < 32; i++) {
        if (sscanf(&priv_key_hex[i*2], "%2hhx", &inst->instance->my_keys.private_key[i]) != 1) {
            DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to parse private key byte %d", i);
        }
    }
    for (int i = 0; i < 64; i++) {
        if (sscanf(&pub_key_hex[i*2], "%2hhx", &inst->instance->my_keys.public_key[i]) != 1) {
            DEBUG_ERROR(DEBUG_CATEGORY_ETCP, "Failed to parse public key byte %d", i);
        }
    }
    
    DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "Keys parsed successfully");
    
    // Create uasync context
    inst->instance->ua = uasync_create();
    if (!inst->instance->ua) {
        free(inst->instance);
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    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) {
        uasync_destroy(inst->instance->ua);
        free(inst->instance);
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    
    // Create ETCP connection
    inst->etcp_conn = etcp_connection_create(inst->instance);
    if (!inst->etcp_conn) {
        memory_pool_destroy(inst->instance->pkt_pool);
        uasync_destroy(inst->instance->ua);
        free(inst->instance);
        free(inst->analyzer);
        free(inst);
        return NULL;
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Created %s instance (node_id=%llu)",
                is_server ? "server" : "client", (unsigned long long)node_id);
    
    return inst;
}

// Cleanup simplified test instance
static void destroy_simple_instance(test_instance_t* inst) {
    if (!inst) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Destroying %s instance", inst->is_server ? "server" : "client");
    
    if (inst->etcp_conn) {
        etcp_connection_close(inst->etcp_conn);
    }
    
    if (inst->instance) {
        if (inst->instance->pkt_pool) {
            memory_pool_destroy(inst->instance->pkt_pool);
        }
        if (inst->instance->ua) {
            uasync_destroy(inst->instance->ua);
        }
        free(inst->instance);
    }
    
    if (inst->analyzer) {
        pthread_mutex_destroy(&inst->analyzer->lock);
        free(inst->analyzer);
    }
    
    free(inst);
}

// Generate test packets manually
static void generate_test_packets(test_instance_t* inst, int count) {
    if (!inst || !inst->etcp_conn) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Generating %d test packets for %s instance", 
                count, inst->is_server ? "server" : "client");
    
    for (int i = 0; i < count; i++) {
        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+1) * 10 + j); // Pattern data
            }
            queue_entry_put(inst->etcp_conn->input_queue, entry);
            DEBUG_DEBUG(DEBUG_CATEGORY_ETCP, "Queued test packet %d for %s", 
                        i+1, inst->is_server ? "server" : "client");
        }
    }
}

// Main test function
int main(int argc, char* argv[]) {
    printf("=== ETCP Traffic Flow Debugging Test (Simplified) ===\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 simplified ETCP traffic flow debugging test");
    
    // Create server instance
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Creating server instance...");
    server_instance = create_simple_instance(0x1111111111111111ULL, 1);
    if (!server_instance) {
        fprintf(stderr, "Failed to create server instance\n");
        return 1;
    }
    
    // Create client instance
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Creating client instance...");
    client_instance = create_simple_instance(0x2222222222222222ULL, 0);
    if (!client_instance) {
        fprintf(stderr, "Failed to create client instance\n");
        destroy_simple_instance(server_instance);
        return 1;
    }
    
    // Phase 1: Test basic packet generation and analysis
    printf("\n=== Phase 1: Basic Packet Generation ===\n");
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Generating test packets...");
    generate_test_packets(server_instance, 5);
    generate_test_packets(client_instance, 5);
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Simulating packet transmission...");
    
    // Let server generate some packets
    for (int i = 0; i < 3; i++) {
        simulate_packet_tx(server_instance);
        usleep(10000); // 1ms delay
    }
    
    // Let client generate some packets
    for (int i = 0; i < 3; i++) {
        simulate_packet_tx(client_instance);
        usleep(10000); // 1ms delay
    }
    
    // Phase 2: Test connection establishment simulation
    printf("\n=== Phase 2: Connection Establishment Simulation ===\n");
    
    // Simulate INIT_REQUEST from client - smaller version for testing
    uint8_t init_request[] = {
        ETCP_INIT_REQUEST,  // Section type
        0x00, 0x14,         // Section length (20 bytes) - smaller for test
        // Node ID (8 bytes)
        0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22,
        // MTU (2 bytes)
        0x05, 0xDC,         // 1500
        // Keepalive (2 bytes)
        0x00, 0x64,         // 100
        // Fake public key (8 bytes) - just for testing
        0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x11, 0x22
    };
    // Total: 3 + 20 = 23 bytes
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Simulating INIT_REQUEST from client...");
    simulate_packet_rx(server_instance, init_request, sizeof(init_request));
    
    // Simulate INIT_RESPONSE from server
    uint8_t init_response[] = {
        ETCP_INIT_RESPONSE, // Section type
        0x00, 0x0A,         // Section length (10 bytes)
        // Node ID (8 bytes)
        0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11,
        // MTU (2 bytes)
        0x05, 0xDC          // 1500
    };
    // Total: 3 + 10 = 13 bytes
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Simulating INIT_RESPONSE from server...");
    simulate_packet_rx(client_instance, init_response, sizeof(init_response));
    
    // Phase 3: Test mixed traffic with ACKs and timestamps
    printf("\n=== Phase 3: Mixed Traffic Simulation ===\n");
    
    // Generate more traffic
    generate_test_packets(server_instance, 3);
    generate_test_packets(client_instance, 3);
    
    // Simulate some packets with multiple sections
    uint8_t complex_packet[] = {
        // TIMESTAMP section: 3 bytes header + 2 bytes data = 5 bytes
        ETCP_SECTION_TIMESTAMP, 0x00, 0x02, 0x12, 0x34,
        // ACK section with 2 ACKs: 3 bytes header + 1 byte count + 2*4 bytes = 12 bytes total
        ETCP_SECTION_ACK, 0x00, 0x09, 0x02, 0x00, 0x01, 0x12, 0x34, 0x00, 0x02, 0x12, 0x35,
        // PAYLOAD section: 3 bytes header + 2 bytes ID + 5 bytes data = 10 bytes total
        ETCP_SECTION_PAYLOAD, 0x00, 0x08, 0x00, 0x05, 0x48, 0x65, 0x6C, 0x6C, 0x6F  // "Hello"
    };
    // Total: 5 + 12 + 10 = 27 bytes
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Simulating complex packet with multiple sections...");
    simulate_packet_rx(server_instance, complex_packet, sizeof(complex_packet));
    
    // Phase 4: Test retransmission requests
    printf("\n=== Phase 4: Retransmission Testing ===\n");
    
    uint8_t retrans_request[] = {
        ETCP_SECTION_RETRANS + 0x02, // RETRANS section with base ID 2
        0x00, 0x04,                   // Section length (4 bytes)
        0x00, 0x03,                   // Request retransmission of packet ID 3
        0x00, 0x04                    // Request retransmission of packet ID 4
    };
    // Total: 3 + 4 = 7 bytes
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Simulating retransmission request...");
    simulate_packet_rx(client_instance, retrans_request, sizeof(retrans_request));
    
    // Final packet generation
    printf("\n=== Final Phase: Additional Traffic ===\n");
    
    // Generate final packets
    for (int i = 0; i < 5; i++) {
        simulate_packet_tx(server_instance);
        simulate_packet_tx(client_instance);
        usleep(5000); // 0.5ms delay
    }
    
    // Print summary
    printf("\n=== Traffic Analysis Summary ===\n");
    printf("Server instance: processed packets with detailed section analysis\n");
    printf("Client instance: processed packets with detailed section analysis\n");
    printf("Connection establishment: INIT handshake simulated\n");
    printf("Mixed traffic: ACK, TIMESTAMP, PAYLOAD sections processed\n");
    printf("Retransmission: RETRANS requests processed\n");
    
    // Cleanup
    destroy_simple_instance(server_instance);
    destroy_simple_instance(client_instance);
    
    DEBUG_INFO(DEBUG_CATEGORY_ETCP, "Test completed");
    printf("\n=== Test completed successfully ===\n");
    
    return 0;
}