#include "utun_test_framework.h"
#include "utun_instance.h"
#include "../lib/u_async.h"
#include "../lib/debug_config.h"
#include "../lib/ll_queue.h"
#include "test_udp_socket.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/time.h>

// External mutex from test_udp_socket.c
extern pthread_mutex_t g_registry_mutex;

#ifndef DEBUG_CATEGORY_TEST
#define DEBUG_CATEGORY_TEST (1 << 30)
#endif

// Default test configurations
const struct test_instance_config TEST_SERVER_CONFIG = {
    .config_file = "test_server.conf",
    .node_id = 0x1111222233334444,
    .port_base = 40000,
    .enable_tun = false,
    .enable_socket_hooks = true,
    .enable_packet_capture = true,
    .log_file = NULL
};

const struct test_instance_config TEST_CLIENT_CONFIG = {
    .config_file = "test_client.conf",
    .node_id = 0x8888777766665555,
    .port_base = 40100,
    .enable_tun = false,
    .enable_socket_hooks = true,
    .enable_packet_capture = true,
    .log_file = NULL
};

// Forward declarations for internal functions
static int test_tun_create_hook(void* priv_data);
static ssize_t test_tun_read_hook(int fd, void *buf, size_t count, void* priv_data);
static ssize_t test_tun_write_hook(int fd, const void *buf, size_t count, void* priv_data);
static int test_socket_create_hook(int domain, int type, int protocol, void* context);
static void test_packet_capture_hook(struct UTUN_INSTANCE* instance, const uint8_t* packet, size_t len,
                                    const char* direction, void* context);
static void test_cleanup_captured_packets(struct test_instance* test);

// Create test instance with virtual components
struct test_instance* test_create_instance(const struct test_instance_config* config) {
    if (!config) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "NULL configuration provided");
        return NULL;
    }
    
    struct test_instance* test = calloc(1, sizeof(struct test_instance));
    if (!test) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to allocate test instance");
        return NULL;
    }
    
    // Copy configuration
    test->config = *config;
    
    // Create uasync instance
    test->ua = uasync_create();
    if (!test->ua) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to create uasync instance");
        free(test);
        return NULL;
    }
    
    // Setup test hooks based on configuration
    if (config->enable_tun || config->enable_socket_hooks || config->enable_packet_capture) {
        test->hooks.test_context = test;
        
        if (config->enable_tun) {
            // Create virtual TUN
            char vtun_name[32];
            snprintf(vtun_name, sizeof(vtun_name), "vtun%ld", config->node_id & 0xFFFF);
            test->vtun = virtual_tun_create(vtun_name);
            if (!test->vtun) {
                DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to create virtual TUN");
                uasync_destroy(test->ua);
                free(test);
                return NULL;
            }
            
            test->hooks.tun_create_override = test_tun_create_hook;
            test->hooks.tun_read_override = test_tun_read_hook;
            test->hooks.tun_write_override = test_tun_write_hook;
        }
        
        if (config->enable_socket_hooks) {
            test->hooks.socket_create_override = test_socket_create_hook;
            test->hooks.sendto_override = NULL; // Will be handled by virtual UDP sockets
            test->hooks.recvfrom_override = NULL; // Will be handled by virtual UDP sockets
        }
        
        if (config->enable_packet_capture) {
            test->hooks.packet_captured = test_packet_capture_hook;
        }
        
        // Install hooks
        #ifdef TEST_BUILD
        utun_test_hooks_set(&test->hooks);
        test->hooks_installed = true;
        #endif
    }
    
    // Allocate packet capture buffer
    if (config->enable_packet_capture) {
        test->captured_packets.capacity = 1000;
        test->captured_packets.packets = calloc(test->captured_packets.capacity, sizeof(uint8_t*));
        test->captured_packets.lengths = calloc(test->captured_packets.capacity, sizeof(size_t));
        test->captured_packets.directions = calloc(test->captured_packets.capacity, sizeof(char*));
        
        if (!test->captured_packets.packets || !test->captured_packets.lengths || !test->captured_packets.directions) {
            DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to allocate packet capture buffers");
            if (test->vtun) virtual_tun_destroy(test->vtun);
            uasync_destroy(test->ua);
            free(test);
            return NULL;
        }
    }
    
    // Allocate UDP socket tracking
    test->udp_socket_capacity = 16;
    test->udp_sockets = calloc(test->udp_socket_capacity, sizeof(struct test_udp_socket*));
    if (!test->udp_sockets) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to allocate UDP socket tracking");
        test_cleanup_captured_packets(test);
        if (test->vtun) virtual_tun_destroy(test->vtun);
        uasync_destroy(test->ua);
        free(test);
        return NULL;
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Created test instance: node_id=0x%016lX, config=%s",
               config->node_id, config->config_file);
    return test;
}

// Destroy test instance
void test_destroy_instance(struct test_instance* test) {
    if (!test) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Destroying test instance: node_id=0x%016lX",
               test->config.node_id);
    
    // Stop instance if running
    test_stop_instance(test);
    
    // Clear hooks if installed
    if (test->hooks_installed) {
        #ifdef TEST_BUILD
        utun_test_hooks_clear();
        #endif
        test->hooks_installed = false;
    }
    
    // Destroy virtual components
    if (test->vtun) {
        virtual_tun_destroy(test->vtun);
    }
    
    // Destroy UDP sockets
    for (int i = 0; i < test->udp_socket_count; i++) {
        if (test->udp_sockets[i]) {
            test_udp_socket_destroy(test->udp_sockets[i]);
        }
    }
    free(test->udp_sockets);
    
    // Cleanup captured packets
    test_cleanup_captured_packets(test);
    
    // Destroy uasync
    if (test->ua) {
        uasync_destroy(test->ua);
    }
    
    // Destroy UTUN instance
    if (test->utun) {
        utun_instance_destroy(test->utun);
    }
    
    free(test);
}

// Start test instance
int test_start_instance(struct test_instance* test) {
    if (!test) return -1;
    if (test->utun) return 0; // Already started
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Starting test instance: node_id=0x%016lX", test->config.node_id);
    
    // Determine flags for instance creation
    uint32_t flags = 0;
    if (!test->config.enable_tun) {
        flags |= UTUN_CREATE_NO_TUN;
    }
    if (test->config.enable_socket_hooks || test->config.enable_packet_capture) {
        flags |= UTUN_CREATE_TEST_MODE;
    }
    if (test->config.enable_socket_hooks) {
        flags |= UTUN_CREATE_NO_SOCKET_BIND;
    }
    if (test->config.enable_tun) {
        flags |= UTUN_CREATE_ALLOW_TUN_FAILURE; // Allow virtual TUN to fail gracefully
    }
    
    // Create UTUN instance
    test->utun = utun_instance_create_ex(test->ua, test->config.config_file, test->config.log_file, flags);
    if (!test->utun) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to create UTUN instance");
        return -1;
    }
    
    // Initialize instance
    if (utun_instance_init(test->utun) < 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to initialize UTUN instance");
        utun_instance_destroy(test->utun);
        test->utun = NULL;
        return -1;
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Test instance started successfully");
    return 0;
}

// Stop test instance
void test_stop_instance(struct test_instance* test) {
    if (!test || !test->utun) return;
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Stopping test instance: node_id=0x%016lX", test->config.node_id);
    
    utun_instance_stop(test->utun);
}

// Inject packet into TUN interface
int test_inject_tun_packet(struct test_instance* test, const uint8_t* packet, size_t len) {
    if (!test || !test->vtun || !packet || len == 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Invalid parameters for TUN packet injection");
        return -1;
    }
    
    return virtual_tun_inject_packet(test->vtun, packet, len);
}

// Inject UDP packet
int test_inject_udp_packet(struct test_instance* test, int socket_fd,
                          const uint8_t* packet, size_t len,
                          const struct sockaddr* src_addr, socklen_t addr_len) {
    if (!test || !packet || len == 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Invalid parameters for UDP packet injection");
        return -1;
    }
    
    struct test_udp_socket* sock = test_udp_socket_find_by_fd(socket_fd);
    if (!sock) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Socket fd=%d not found", socket_fd);
        return -1;
    }
    
    return test_udp_socket_inject(sock, packet, len, src_addr, addr_len);
}

// Get captured packets
size_t test_get_captured_packets(struct test_instance* test, 
                                const char* direction,
                                uint8_t*** packets, size_t** lengths) {
    if (!test || !direction || !packets || !lengths) return 0;
    
    // Count packets matching direction
    size_t count = 0;
    for (size_t i = 0; i < test->captured_packets.count; i++) {
        if (strcmp(test->captured_packets.directions[i], direction) == 0) {
            count++;
        }
    }
    
    if (count == 0) return 0;
    
    // Allocate result arrays
    *packets = calloc(count, sizeof(uint8_t*));
    *lengths = calloc(count, sizeof(size_t));
    if (!*packets || !*lengths) {
        free(*packets);
        free(*lengths);
        return 0;
    }
    
    // Fill result arrays
    size_t result_idx = 0;
    for (size_t i = 0; i < test->captured_packets.count; i++) {
        if (strcmp(test->captured_packets.directions[i], direction) == 0) {
            (*packets)[result_idx] = test->captured_packets.packets[i];
            (*lengths)[result_idx] = test->captured_packets.lengths[i];
            result_idx++;
        }
    }
    
    return count;
}

// Get captured packet count
size_t test_get_captured_packet_count(struct test_instance* test, const char* direction) {
    if (!test || !direction) return 0;
    
    size_t count = 0;
    for (size_t i = 0; i < test->captured_packets.count; i++) {
        if (strcmp(test->captured_packets.directions[i], direction) == 0) {
            count++;
        }
    }
    return count;
}

// Clear captured packets
void test_clear_captured_packets(struct test_instance* test) {
    if (!test) return;
    
    for (size_t i = 0; i < test->captured_packets.count; i++) {
        if (test->captured_packets.packets[i]) {
            free(test->captured_packets.packets[i]);
        }
        if (test->captured_packets.directions[i]) {
            free(test->captured_packets.directions[i]);
        }
    }
    test->captured_packets.count = 0;
}

// Wait for packet with timeout
int test_wait_for_packet(struct test_instance* test, const char* direction, int timeout_ms) {
    if (!test || !direction || timeout_ms <= 0) return -1;
    
    size_t initial_count = test_get_captured_packet_count(test, direction);
    size_t current_count = initial_count;
    
    struct timeval start, current;
    gettimeofday(&start, NULL);
    
    while (current_count == initial_count) {
        gettimeofday(&current, NULL);
        int elapsed_ms = (current.tv_sec - start.tv_sec) * 1000 + (current.tv_usec - start.tv_usec) / 1000;
        
        if (elapsed_ms >= timeout_ms) {
            DEBUG_WARN(DEBUG_CATEGORY_TEST, "Timeout waiting for %s packet after %dms", direction, timeout_ms);
            return -1;
        }
        
        usleep(10000); // Sleep 10ms
        current_count = test_get_captured_packet_count(test, direction);
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Detected %s packet after %dms", direction, 
               (current.tv_sec - start.tv_sec) * 1000 + (current.tv_usec - start.tv_usec) / 1000);
    return 0;
}

// Two-instance test runner
int test_run_two_instances(const struct test_instance_config* server_config,
                          const struct test_instance_config* client_config,
                          void (*test_scenario)(struct test_instance* server, 
                                              struct test_instance* client),
                          int timeout_ms) {
    if (!server_config || !client_config || !test_scenario) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Invalid parameters for two-instance test");
        return -1;
    }
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Starting two-instance test");
    
    // Create instances
    struct test_instance* server = test_create_instance(server_config);
    struct test_instance* client = test_create_instance(client_config);
    
    if (!server || !client) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to create test instances");
        test_destroy_instance(server);
        test_destroy_instance(client);
        return -1;
    }
    
    // Start instances
    if (test_start_instance(server) < 0 || test_start_instance(client) < 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to start test instances");
        test_destroy_instance(server);
        test_destroy_instance(client);
        return -1;
    }
    
    // Give instances time to initialize
    sleep(1);
    
    // Run test scenario
    int result = 0;
    if (timeout_ms > 0) {
        // TODO: Implement timeout mechanism for test scenario
        test_scenario(server, client);
    } else {
        test_scenario(server, client);
    }
    
    // Cleanup
    test_destroy_instance(server);
    test_destroy_instance(client);
    
    DEBUG_INFO(DEBUG_CATEGORY_TEST, "Two-instance test completed: %s", result == 0 ? "SUCCESS" : "FAILED");
    return result;
}

// Get test instance statistics
void test_get_instance_stats(struct test_instance* test, size_t* tun_packets_sent,
                            size_t* tun_packets_received, size_t* udp_packets_sent,
                            size_t* udp_packets_received) {
    if (!test) return;
    
    if (test->vtun) {
        virtual_tun_get_stats(test->vtun, tun_packets_sent, tun_packets_received, NULL, NULL);
    }
    
    // Aggregate UDP socket statistics
    size_t udp_sent = 0, udp_received = 0;
    for (int i = 0; i < test->udp_socket_count; i++) {
        if (test->udp_sockets[i]) {
            size_t sent, recv;
            test_udp_socket_get_stats(test->udp_sockets[i], &sent, &recv, NULL, NULL, NULL, NULL);
            udp_sent += sent;
            udp_received += recv;
        }
    }
    
    if (udp_packets_sent) *udp_packets_sent = udp_sent;
    if (udp_packets_received) *udp_packets_received = udp_received;
}

// Utility functions
const char* test_get_direction_name(const char* direction) {
    if (!direction) return "unknown";
    
    if (strcmp(direction, "tun_in") == 0) return "TUN incoming";
    if (strcmp(direction, "tun_out") == 0) return "TUN outgoing";
    if (strcmp(direction, "udp_in") == 0) return "UDP incoming";
    if (strcmp(direction, "udp_out") == 0) return "UDP outgoing";
    return direction;
}

int test_compare_packets(const uint8_t* pkt1, size_t len1, const uint8_t* pkt2, size_t len2) {
    if (!pkt1 || !pkt2) return -1;
    if (len1 != len2) return -1;
    return memcmp(pkt1, pkt2, len1);
}

void test_dump_packet(const char* prefix, const uint8_t* packet, size_t len) {
    if (!prefix || !packet || len == 0) return;
    
    printf("%s: %zu bytes\n", prefix, len);
    for (size_t i = 0; i < len; i++) {
        printf("%02X ", packet[i]);
        if ((i + 1) % 16 == 0) printf("\n");
    }
    if (len % 16 != 0) printf("\n");
}

// Internal hook implementations
static int test_tun_create_hook(void* priv_data) {
    struct test_instance* test = (struct test_instance*)priv_data;
    if (!test || !test->vtun) return -1;
    
    // Return the read fd as the TUN device fd
    return virtual_tun_get_read_fd(test->vtun);
}

static ssize_t test_tun_read_hook(int fd, void *buf, size_t count, void* priv_data) {
    struct test_instance* test = (struct test_instance*)priv_data;
    if (!test || !test->vtun) return -1;
    
    return virtual_tun_read_packet(test->vtun, buf, count);
}

static ssize_t test_tun_write_hook(int fd, const void *buf, size_t count, void* priv_data) {
    struct test_instance* test = (struct test_instance*)priv_data;
    if (!test || !test->vtun) return -1;
    
    ssize_t result = virtual_tun_write_packet(test->vtun, buf, count);
    
    // Capture packet if enabled
    if (test->config.enable_packet_capture && result > 0) {
        test_packet_capture_hook(test->utun, buf, result, "tun_out", test);
    }
    
    return result;
}

static int test_socket_create_hook(int domain, int type, int protocol, void* context) {
    struct test_instance* test = (struct test_instance*)context;
    if (!test || !test->config.enable_socket_hooks) return -1;
    
    // Create virtual UDP socket
    struct test_udp_socket* sock = test_udp_socket_create(domain);
    if (!sock) return -1;
    
    // Add to tracking
    pthread_mutex_lock(&g_registry_mutex);
    if (test->udp_socket_count >= test->udp_socket_capacity) {
        // Expand array
        int new_capacity = test->udp_socket_capacity * 2;
        struct test_udp_socket** new_sockets = realloc(test->udp_sockets, 
                                                       new_capacity * sizeof(struct test_udp_socket*));
        if (!new_sockets) {
            pthread_mutex_unlock(&g_registry_mutex);
            test_udp_socket_destroy(sock);
            return -1;
        }
        test->udp_sockets = new_sockets;
        test->udp_socket_capacity = new_capacity;
    }
    test->udp_sockets[test->udp_socket_count++] = sock;
    pthread_mutex_unlock(&g_registry_mutex);
    
    return test_udp_socket_get_fd(sock);
}

static void test_packet_capture_hook(struct UTUN_INSTANCE* instance, const uint8_t* packet, size_t len,
                                    const char* direction, void* context) {
    struct test_instance* test = (struct test_instance*)context;
    if (!test || !test->config.enable_packet_capture || !packet || len == 0) return;
    
    // Expand capture buffer if needed
    if (test->captured_packets.count >= test->captured_packets.capacity) {
        size_t new_capacity = test->captured_packets.capacity * 2;
        uint8_t** new_packets = realloc(test->captured_packets.packets, 
                                       new_capacity * sizeof(uint8_t*));
        size_t* new_lengths = realloc(test->captured_packets.lengths, 
                                     new_capacity * sizeof(size_t));
        char** new_directions = realloc(test->captured_packets.directions, 
                                       new_capacity * sizeof(char*));
        
        if (!new_packets || !new_lengths || !new_directions) {
            DEBUG_ERROR(DEBUG_CATEGORY_TEST, "Failed to expand packet capture buffer");
            return;
        }
        
        test->captured_packets.packets = new_packets;
        test->captured_packets.lengths = new_lengths;
        test->captured_packets.directions = new_directions;
        test->captured_packets.capacity = new_capacity;
    }
    
    // Store packet
    uint8_t* packet_copy = malloc(len);
    if (!packet_copy) return;
    
    memcpy(packet_copy, packet, len);
    test->captured_packets.packets[test->captured_packets.count] = packet_copy;
    test->captured_packets.lengths[test->captured_packets.count] = len;
    test->captured_packets.directions[test->captured_packets.count] = strdup(direction);
    test->captured_packets.count++;
    
    DEBUG_DEBUG(DEBUG_CATEGORY_TEST, "Captured packet: %s, %zu bytes", direction, len);
}

static void test_cleanup_captured_packets(struct test_instance* test) {
    if (!test) return;
    
    for (size_t i = 0; i < test->captured_packets.count; i++) {
        if (test->captured_packets.packets[i]) {
            free(test->captured_packets.packets[i]);
        }
        if (test->captured_packets.directions[i]) {
            free(test->captured_packets.directions[i]);
        }
    }
    
    free(test->captured_packets.packets);
    free(test->captured_packets.lengths);
    free(test->captured_packets.directions);
    
    memset(&test->captured_packets, 0, sizeof(test->captured_packets));
}