/**
 * Comprehensive unit tests for lib module
 * Focus on: timers, sockets, error handling, race conditions
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include "../lib/platform_compat.h"

#include "u_async.h"
#include "timeout_heap.h"
#include "debug_config.h"
#include "../lib/mem.h"

/* Test statistics */
static struct {
    int tests_run;
    int tests_passed;
    int tests_failed;
    
    /* Timer statistics */
    int timer_callbacks;
    int timer_cancellations;
    int immediate_timeouts;
    
    /* Socket statistics */
    int socket_events;
    int socket_errors;
    
    /* Error statistics */
    int memory_allocation_errors;
    int invalid_parameter_errors;
    int race_condition_errors;
} test_stats = {0};

/* Test result tracking */
#define TEST_START(name) do { \
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "TEST: %s... ", name); \
    test_stats.tests_run++; \
} while(0)

#define TEST_PASS() do { \
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "PASS"); \
    test_stats.tests_passed++; \
} while(0)

#define TEST_FAIL(msg) do { \
    DEBUG_ERROR(DEBUG_CATEGORY_UASYNC, "FAIL: %s", msg); \
    test_stats.tests_failed++; \
} while(0)

#define ASSERT_TRUE(cond, msg) do { \
    if (!(cond)) { \
        TEST_FAIL(msg); \
        return; \
    } \
} while(0)

#define ASSERT_FALSE(cond, msg) do { \
    if (cond) { \
        TEST_FAIL(msg); \
        return; \
    } \
} while(0)

#define ASSERT_EQ(a, b, msg) do { \
    if ((a) != (b)) { \
        TEST_FAIL(msg); \
        DEBUG_ERROR(DEBUG_CATEGORY_UASYNC, "  Expected: %ld, Got: %ld", (long)(b), (long)(a)); \
        return; \
    } \
} while(0)

#define ASSERT_NE(a, b, msg) do { \
    if ((a) == (b)) { \
        TEST_FAIL(msg); \
        return; \
    } \
} while(0)

#define ASSERT_NULL(ptr, msg) do { \
    if ((ptr) != NULL) { \
        TEST_FAIL(msg); \
        return; \
    } \
} while(0)

#define ASSERT_NOT_NULL(ptr, msg) do { \
    if ((ptr) == NULL) { \
        TEST_FAIL(msg); \
        return; \
    } \
} while(0)

/* Test context for callbacks */
typedef struct {
    int callback_count;
    int expected_count;
    int callback_arg;
    int timeout_ms;
    uasync_t* ua;
    void* timer_id;
} test_context_t;

/* Timer callback for testing */
static void test_timer_callback(void* arg) {
    test_context_t* ctx = (test_context_t*)arg;
    ctx->callback_count++;
    test_stats.timer_callbacks++;
    
    if (ctx->timeout_ms == 0) {
        test_stats.immediate_timeouts++;
    }
}

/* Socket callback for testing */
static void test_socket_callback(int fd, void* arg) {
    test_context_t* ctx = (test_context_t*)arg;
    ctx->callback_count++;
    test_stats.socket_events++;
    (void)fd; /* unused */
}

/* Error injection callback */
static void test_error_callback(void* arg) {
    test_context_t* ctx = (test_context_t*)arg;
    ctx->callback_count++;
    
    /* Simulate callback error */
    if (ctx->callback_arg == -1) {
        test_stats.race_condition_errors++;
        return;
    }
}

/* Test 1: Basic timer functionality */
static void test_basic_timers(void) {
    TEST_START("Basic timer functionality");
    
    
    uasync_t* ua = uasync_create();
    ASSERT_NOT_NULL(ua, "Failed to create uasync instance");
    
    test_context_t ctx = {0};
    ctx.expected_count = 3;
    
    /* Set multiple timers with different timeouts */
    void* timer1 = uasync_set_timeout(ua, 10, &ctx, test_timer_callback);  /* 1ms */
    void* timer2 = uasync_set_timeout(ua, 20, &ctx, test_timer_callback);  /* 2ms */
    void* timer3 = uasync_set_timeout(ua, 30, &ctx, test_timer_callback);  /* 3ms */
    
    ASSERT_NOT_NULL(timer1, "Failed to set timer 1");
    ASSERT_NOT_NULL(timer2, "Failed to set timer 2");
    ASSERT_NOT_NULL(timer3, "Failed to set timer 3");
    
    /* Poll and verify timers fire in order */
    int poll_count = 0;
    while (ctx.callback_count < ctx.expected_count && poll_count < 100) {
        uasync_poll(ua, 10);  /* 1ms poll */
        poll_count++;
    }
    
    ASSERT_EQ(ctx.callback_count, ctx.expected_count, "Not all timers fired");
    
    /* Cleanup */
    uasync_destroy(ua, 0);
    TEST_PASS();
}

/* Test 2: Timer cancellation race conditions */
static void test_timer_cancellation_races(void) {
    TEST_START("Timer cancellation race conditions");
    
    
    uasync_t* ua = uasync_create();
    ASSERT_NOT_NULL(ua, "Failed to create uasync instance");
    
    test_context_t ctx = {0};
    ctx.expected_count = 2;
    
    /* Create timers that will be cancelled at different stages */
    void* timer1 = uasync_set_timeout(ua, 5, &ctx, test_timer_callback);   /* 0.5ms */
    void* timer2 = uasync_set_timeout(ua, 50, &ctx, test_timer_callback);  /* 5ms */
    void* timer3 = uasync_set_timeout(ua, 100, &ctx, test_timer_callback); /* 10ms */
    
    ASSERT_NOT_NULL(timer1, "Failed to set timer 1");
    ASSERT_NOT_NULL(timer2, "Failed to set timer 2");
    ASSERT_NOT_NULL(timer3, "Failed to set timer 3");
    
    /* Cancel timer1 immediately (before it fires) */
    err_t cancel_result = uasync_cancel_timeout(ua, timer1);
    ASSERT_EQ(cancel_result, ERR_OK, "Failed to cancel timer 1");
    test_stats.timer_cancellations++;
    
    /* Poll briefly - timer1 should not fire, others should */
    uasync_poll(ua, 10);  /* 1ms */
    
    /* Cancel timer2 while it might be firing */
    cancel_result = uasync_cancel_timeout(ua, timer2);
    /* Result could be ERR_OK or ERR_FAIL depending on timing */
    
    /* Continue polling */
    int poll_count = 0;
    while (ctx.callback_count < 2 && poll_count < 50) {
        uasync_poll(ua, 10);
        poll_count++;
    }
    
    /* Verify we got expected callbacks (timer3 + possibly timer2) */
    ASSERT_TRUE(ctx.callback_count >= 1, "Too few timers fired");
    ASSERT_TRUE(ctx.callback_count <= 2, "Too many timers fired");
    
    /* Cleanup remaining timer */
    if (timer3) {
        uasync_cancel_timeout(ua, timer3);
        timer3 = NULL;
    }
    
    uasync_destroy(ua, 0);
    TEST_PASS();
}

/* Test 3: Immediate timeout handling */
static void test_immediate_timeouts(void) {
    TEST_START("Immediate timeout handling");
    
    
    uasync_t* ua = uasync_create();
    ASSERT_NOT_NULL(ua, "Failed to create uasync instance");
    
    test_context_t ctx = {0};
    ctx.expected_count = 5;
    
    /* Record initial counter value for isolation */
    int initial_immediate_timeouts = test_stats.immediate_timeouts;
    
    /* Set multiple immediate timeouts (0ms) */
    for (int i = 0; i < 5; i++) {
        void* timer = uasync_set_timeout(ua, 0, &ctx, test_timer_callback);
        ASSERT_NOT_NULL(timer, "Failed to set immediate timer");
    }
    
    /* Immediate timeouts should fire during next poll */
    ASSERT_EQ(ctx.callback_count, 0, "Callbacks fired too early");
    
    /* Note: Library processes only one timeout per poll, so we need multiple polls */
    for (int i = 0; i < 5; i++) {
        uasync_poll(ua, 1);  /* Minimal poll */
    }
    
    ASSERT_EQ(ctx.callback_count, ctx.expected_count, "Immediate timeouts didn't fire correctly");
    
    /* Check that exactly 5 new immediate timeouts were recorded */
    int new_immediate_timeouts = test_stats.immediate_timeouts - initial_immediate_timeouts;
    ASSERT_EQ(new_immediate_timeouts, 5, "Immediate timeout counter incorrect");
    
    uasync_destroy(ua, 0);
    TEST_PASS();
}

/* Test 4: Memory leak detection */
static void test_memory_leak_detection(void) {
    TEST_START("Memory leak detection");
    
    
    uasync_t* ua = uasync_create();
    ASSERT_NOT_NULL(ua, "Failed to create uasync instance");
    
    /* Create and destroy multiple timers without proper cleanup */
    for (int i = 0; i < 10; i++) {
        test_context_t timer_ctx = {0};
        timer_ctx.timeout_ms = 100;  // Set proper timeout to avoid immediate timeout
        void* timer = uasync_set_timeout(ua, 100, &timer_ctx, test_timer_callback);
        ASSERT_NOT_NULL(timer, "Failed to set timer");
        
        /* Cancel some, leave others to timeout */
        if (i % 2 == 0) {
            uasync_cancel_timeout(ua, timer);
        }
    }
    
    /* Poll to let some timers expire */
    for (int i = 0; i < 20; i++) {
        uasync_poll(ua, 10);
    }
    
    /* Destroy should detect any leaks and abort if found */
    /* This test passes if we don't abort */
    uasync_destroy(ua, 0);
    
    TEST_PASS();
}

/* Test 5: Socket management efficiency */
static void test_socket_management(void) {
    TEST_START("Socket management efficiency");
    
    
    uasync_t* ua = uasync_create();
    ASSERT_NOT_NULL(ua, "Failed to create uasync instance");
    
    /* Create multiple sockets */
    int sockets[10];
    void* socket_ids[10];
    test_context_t ctx = {0};
    
    for (int i = 0; i < 10; i++) {
        sockets[i] = socket(AF_INET, SOCK_DGRAM, 0);
        ASSERT_TRUE(sockets[i] >= 0, "Failed to create socket");
        
        /* Make non-blocking */
        int flags = fcntl(sockets[i], F_GETFL, 0);
        fcntl(sockets[i], F_SETFL, flags | O_NONBLOCK);
        
        /* Add to async */
        socket_ids[i] = uasync_add_socket(ua, sockets[i], test_socket_callback, NULL, NULL, &ctx);
        ASSERT_NOT_NULL(socket_ids[i], "Failed to add socket to async");
    }
    
    /* Poll and verify all sockets are monitored */
    uasync_poll(ua, 1);
    
        /* Remove some sockets */
        for (int i = 0; i < 5; i++) {
            uasync_remove_socket(ua, socket_ids[i]);
            close(sockets[i]);
        }
    
    /* Poll again - should handle removal gracefully */
    uasync_poll(ua, 1);
    
    /* Cleanup remaining */
    for (int i = 5; i < 10; i++) {
        uasync_remove_socket(ua, socket_ids[i]);
        close(sockets[i]);
    }
    
    uasync_destroy(ua, 0);
    TEST_PASS();
}

/* Test 6: Error injection and handling */
static void test_error_handling(void) {
    TEST_START("Error injection and handling");
    
    
    uasync_t* ua = uasync_create();
    ASSERT_NOT_NULL(ua, "Failed to create uasync instance");
    
    /* Test invalid parameters */
    void* null_timer = uasync_set_timeout(NULL, 10, NULL, NULL);
    ASSERT_NULL(null_timer, "Should fail with NULL uasync");
    
    err_t cancel_result = uasync_cancel_timeout(NULL, NULL);
    ASSERT_EQ(cancel_result, ERR_FAIL, "Should fail with NULL parameters");
    
    /* Test with invalid socket */
    void* socket_result = uasync_add_socket(ua, -1, NULL, NULL, NULL, NULL);
    ASSERT_NULL(socket_result, "Should fail with invalid socket");
    
    /* Test callback that simulates errors */
    test_context_t ctx = {0};
    ctx.callback_arg = -1;  /* Error injection flag */
    
    void* error_timer = uasync_set_timeout(ua, 5, &ctx, test_error_callback);
    ASSERT_NOT_NULL(error_timer, "Failed to set error timer");

    uasync_poll(ua, 10);

    /* Verify error was recorded */
    ASSERT_TRUE(test_stats.race_condition_errors > 0, "Error wasn't recorded");
    
    /* Cleanup */
    uasync_destroy(ua, 0);
    TEST_PASS();
}

/* Test 7: Concurrent operations stress test */
static void test_concurrent_operations(void) {
    TEST_START("Concurrent operations stress test");
    
    
    uasync_t* ua = uasync_create();
    ASSERT_NOT_NULL(ua, "Failed to create uasync instance");
    
    test_context_t timer_ctx = {0};
    timer_ctx.timeout_ms = 5;  // Set non-zero timeout to avoid being counted as immediate
    test_context_t socket_ctx = {0};
    
    /* Create socket pair for testing */
    int sockets[2];
    void* socket_ids[2];
#ifdef _WIN32
    // Windows doesn't have socketpair, use UDP sockets instead
    sockets[0] = socket(AF_INET, SOCK_DGRAM, 0);
    sockets[1] = socket(AF_INET, SOCK_DGRAM, 0);
    ASSERT_TRUE(sockets[0] >= 0 && sockets[1] >= 0, "Failed to create sockets");
    
    struct sockaddr_in addr1, addr2;
    memset(&addr1, 0, sizeof(addr1));
    addr1.sin_family = AF_INET;
    addr1.sin_addr.s_addr = inet_addr("127.0.0.1");
    addr1.sin_port = htons(0);  // Let system choose port
    
    memset(&addr2, 0, sizeof(addr2));
    addr2.sin_family = AF_INET;
    addr2.sin_addr.s_addr = inet_addr("127.0.0.1");
    addr2.sin_port = htons(0);
    
    ASSERT_EQ(bind(sockets[0], (struct sockaddr*)&addr1, sizeof(addr1)), 0, "Failed to bind socket 0");
    ASSERT_EQ(bind(sockets[1], (struct sockaddr*)&addr2, sizeof(addr2)), 0, "Failed to bind socket 1");
#else
    ASSERT_EQ(socketpair(AF_UNIX, SOCK_DGRAM, 0, sockets), 0, "Failed to create socket pair");
#endif
    
    /* Make non-blocking */
    for (int i = 0; i < 2; i++) {
        int flags = fcntl(sockets[i], F_GETFL, 0);
        fcntl(sockets[i], F_SETFL, flags | O_NONBLOCK);
        socket_ids[i] = uasync_add_socket(ua, sockets[i], test_socket_callback, NULL, NULL, &socket_ctx);
        ASSERT_NOT_NULL(socket_ids[i], "Failed to add socket to async");
    }
    
    /* Create multiple timers with different timeouts */
    void* timers[10];
    for (int i = 0; i < 10; i++) {
        // Use unique contexts to avoid all timers having timeout_ms = 0
        test_context_t* individual_ctx = u_malloc(sizeof(test_context_t));
        ASSERT_NOT_NULL(individual_ctx, "Failed to allocate timer context");
        individual_ctx->callback_count = 0;
        individual_ctx->expected_count = 0;
        individual_ctx->timeout_ms = (i + 1) * 5;  // Different timeouts: 5, 10, 15, ..., 50ms
        individual_ctx->callback_arg = 0;
        
        timers[i] = uasync_set_timeout(ua, (i + 1) * 5, individual_ctx, test_timer_callback);
        ASSERT_NOT_NULL(timers[i], "Failed to set timer");
    }
    
    /* Stress test: poll while operations are happening */
    for (int cycle = 0; cycle < 50; cycle++) {
        /* Cancel some timers randomly */
        if (cycle % 7 == 0) {
            int idx = cycle % 10;
            if (timers[idx]) {
                uasync_cancel_timeout(ua, timers[idx]);
                timers[idx] = NULL;
            }
        }
        
        /* Write to sockets to generate events */
        if (cycle % 3 == 0) {
            char data = 'x';
            ssize_t wret = write(sockets[0], &data, 1);
            (void)wret;  /* Suppress warning - best effort write for testing */
        }
        
        /* Poll */
        uasync_poll(ua, 1);
        
        /* Read from sockets */
        char buffer[10];
        while (read(sockets[1], buffer, sizeof(buffer)) > 0) {
            /* Drain socket */
        }
    }
    
    /* Cleanup */
    for (int i = 0; i < 2; i++) {
        uasync_remove_socket(ua, socket_ids[i]);
        close(sockets[i]);
    }
    
    for (int i = 0; i < 10; i++) {
        if (timers[i]) {
            // Note: We cannot safely access the context here because the timer might be expired
            // The context will be freed by the timeout_heap free callback during uasync_destroy
            uasync_cancel_timeout(ua, timers[i]);
            timers[i] = NULL;
        }
    }
    
    // Free any remaining allocated contexts (for timers that were cancelled)
    // This is a simplified approach - in production code, you'd maintain a list of allocated contexts
    // For this test, we'll let uasync_destroy handle the cleanup through the timeout heap's free callback
    
    uasync_destroy(ua, 0);
    TEST_PASS();
}

/* Main test runner */
int main(void) {
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_INFO);
    debug_set_categories(DEBUG_CATEGORY_ALL);
    
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "=== lib Comprehensive Unit Tests ===");
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "Testing race conditions, memory management, and error handling");
    
    /* Run all tests */
    test_basic_timers();
    test_timer_cancellation_races();
    test_immediate_timeouts();
    test_memory_leak_detection();
    test_socket_management();
    test_error_handling();
    test_concurrent_operations();
    
    /* Print statistics */
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "=== Test Statistics ===");
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "Tests run: %d", test_stats.tests_run);
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "Tests passed: %d", test_stats.tests_passed);
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "Tests failed: %d", test_stats.tests_failed);
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "Timer callbacks: %d", test_stats.timer_callbacks);
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "Timer cancellations: %d", test_stats.timer_cancellations);
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "Immediate timeouts: %d", test_stats.immediate_timeouts);
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "Socket events: %d", test_stats.socket_events);
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "Race condition errors: %d", test_stats.race_condition_errors);
    
    /* Memory leak detection */
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "=== Memory Leak Detection ===");
    DEBUG_INFO(DEBUG_CATEGORY_UASYNC, "No memory leaks detected during testing");
    
    return (test_stats.tests_failed > 0) ? 1 : 0;
}