#include "timeout_heap.h"
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <time.h>
#include <stdint.h>

#define TEST_ASSERT(cond, msg) do { \
    if (!(cond)) { \
        fprintf(stderr, "FAIL: %s:%d: %s\n", __FILE__, __LINE__, (msg)); \
        exit(1); \
    } \
} while (0)

// Element structure for large random test
typedef struct {
    uint64_t expiration;
    void *data;
    int cancelled;
} HeapElement;

// Helper: compare two entries (for qsort)
static int compare_entries(const void *a, const void *b) {
    const TimeoutEntry *ea = (const TimeoutEntry *)a;
    const TimeoutEntry *eb = (const TimeoutEntry *)b;
    if (ea->expiration < eb->expiration) return -1;
    if (ea->expiration > eb->expiration) return 1;
    return 0;
}

// Test 1: basic operations
static void test_basic(void) {
    printf("Test 1: basic operations...\n");
    TimeoutHeap *h = timeout_heap_create(10);
    TEST_ASSERT(h != NULL, "heap creation");
    TEST_ASSERT(h->size == 0, "initial size zero");
    
    // Push one element
    int data1 = 42;
    int ret = timeout_heap_push(h, 100, &data1);
    TEST_ASSERT(ret == 0, "push success");
    TEST_ASSERT(h->size == 1, "size after push");
    
    // Peek
    TimeoutEntry entry;
    ret = timeout_heap_peek(h, &entry);
    TEST_ASSERT(ret == 0, "peek success");
    TEST_ASSERT(entry.expiration == 100, "peek expiration");
    TEST_ASSERT(entry.data == &data1, "peek data");
    TEST_ASSERT(entry.deleted == 0, "peek deleted flag");
    
    // Pop
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == 0, "pop success");
    TEST_ASSERT(entry.expiration == 100, "pop expiration");
    TEST_ASSERT(entry.data == &data1, "pop data");
    TEST_ASSERT(h->size == 0, "size after pop");
    
    // Pop empty
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == -1, "pop empty returns -1");
    
    timeout_heap_destroy(h);
    printf("  Passed\n");
}

// Test 2: ordering with random values
static void test_ordering(size_t num_elements) {
    printf("Test 2: ordering with %zu elements...\n", num_elements);
    TimeoutHeap *h = timeout_heap_create(4); // small initial capacity
    TEST_ASSERT(h != NULL, "heap creation");
    
    // Generate random expirations and store them
    TimeoutEntry *expected = malloc(num_elements * sizeof(TimeoutEntry));
    TEST_ASSERT(expected != NULL, "alloc expected array");
    
    srand((unsigned int)time(NULL));
    for (size_t i = 0; i < num_elements; i++) {
        // Generate unique expiration to avoid ties
        uint64_t exp = i * 1000000ULL + (rand() % 1000000);
        int *data = malloc(sizeof(int));
        *data = (int)i;
        int ret = timeout_heap_push(h, exp, data);
        TEST_ASSERT(ret == 0, "push success");
        expected[i].expiration = exp;
        expected[i].data = data;
        expected[i].deleted = 0;
    }
    
    // Sort expected by expiration
    qsort(expected, num_elements, sizeof(TimeoutEntry), compare_entries);
    
    // Pop and verify order
    for (size_t i = 0; i < num_elements; i++) {
        TimeoutEntry entry;
        int ret = timeout_heap_pop(h, &entry);
        TEST_ASSERT(ret == 0, "pop success");
        TEST_ASSERT(entry.expiration == expected[i].expiration, "expiration order");
        TEST_ASSERT(entry.data == expected[i].data, "data matches");
        TEST_ASSERT(entry.deleted == 0, "not deleted");
        free(entry.data); // cleanup
    }
    
    // Heap should be empty
    TimeoutEntry entry;
    int ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == -1, "heap empty after all pops");
    
    free(expected);
    timeout_heap_destroy(h);
    printf("  Passed\n");
}

// Test 3: cancel operation
static void test_cancel(void) {
    printf("Test 3: cancel operation...\n");
    TimeoutHeap *h = timeout_heap_create(10);
    TEST_ASSERT(h != NULL, "heap creation");
    
    int data1 = 1, data2 = 2, data3 = 3;
    timeout_heap_push(h, 100, &data1);
    timeout_heap_push(h, 200, &data2);
    timeout_heap_push(h, 300, &data3);
    
    // Cancel middle element
    int ret = timeout_heap_cancel(h, 200, &data2);
    TEST_ASSERT(ret == 0, "cancel success");
    
    // Pop and verify we get data1 then data3, not data2
    TimeoutEntry entry;
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == 0, "pop 1 success");
    TEST_ASSERT(entry.expiration == 100, "first expiration");
    TEST_ASSERT(entry.data == &data1, "first data");
    
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == 0, "pop 2 success");
    TEST_ASSERT(entry.expiration == 300, "second expiration");
    TEST_ASSERT(entry.data == &data3, "second data");
    
    // Heap empty
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == -1, "heap empty");
    
    timeout_heap_destroy(h);
    printf("  Passed\n");
}

// Test 4: duplicate expiration times
static void test_duplicate_expirations(void) {
    printf("Test 4: duplicate expirations...\n");
    TimeoutHeap *h = timeout_heap_create(10);
    TEST_ASSERT(h != NULL, "heap creation");
    
    int data1 = 1, data2 = 2, data3 = 3;
    // Push same expiration for data1 and data2
    timeout_heap_push(h, 100, &data1);
    timeout_heap_push(h, 100, &data2);
    timeout_heap_push(h, 200, &data3);
    
    // Pop should give both items with expiration 100 (order not guaranteed)
    TimeoutEntry entry;
    int ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == 0, "first pop success");
    TEST_ASSERT(entry.expiration == 100, "first pop expiration");
    // Record which data pointer we got
    void *first_data = entry.data;
    TEST_ASSERT(first_data == &data1 || first_data == &data2, "first pop data matches one of duplicates");
    
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == 0, "second pop success");
    TEST_ASSERT(entry.expiration == 100, "second pop expiration");
    void *second_data = entry.data;
    TEST_ASSERT(second_data == &data1 || second_data == &data2, "second pop data matches one of duplicates");
    TEST_ASSERT(first_data != second_data, "two different data pointers");
    
    // Third pop should give expiration 200
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == 0, "third pop success");
    TEST_ASSERT(entry.expiration == 200, "third pop expiration");
    TEST_ASSERT(entry.data == &data3, "third pop data");
    
    // Heap empty
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == -1, "heap empty");
    
    timeout_heap_destroy(h);
    printf("  Passed\n");
}

// Test 5: deleted flag handling (multiple deleted at root)
static void test_deleted_root(void) {
    printf("Test 5: deleted root handling...\n");
    TimeoutHeap *h = timeout_heap_create(10);
    TEST_ASSERT(h != NULL, "heap creation");
    
    int data1 = 1, data2 = 2, data3 = 3, data4 = 4;
    timeout_heap_push(h, 100, &data1);
    timeout_heap_push(h, 50, &data2);  // earliest
    timeout_heap_push(h, 150, &data3);
    timeout_heap_push(h, 75, &data4);
    
    // Cancel the earliest (root)
    timeout_heap_cancel(h, 50, &data2);
    
    // Peek should skip deleted root and give next earliest (75)
    TimeoutEntry entry;
    int ret = timeout_heap_peek(h, &entry);
    TEST_ASSERT(ret == 0, "peek success after root deleted");
    TEST_ASSERT(entry.expiration == 75, "peek expiration after skip");
    TEST_ASSERT(entry.data == &data4, "peek data after skip");
    
    // Pop should also skip deleted root
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == 0, "pop success after root deleted");
    TEST_ASSERT(entry.expiration == 75, "pop expiration after skip");
    
    // Next pop should be 100
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == 0, "pop second");
    TEST_ASSERT(entry.expiration == 100, "second pop expiration");
    
    // Cancel 150, then push another earlier
    timeout_heap_cancel(h, 150, &data3);
    int data5 = 5;
    timeout_heap_push(h, 60, &data5); // earlier than 150 but after 100
    
    // Pop should give 60
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == 0, "pop after new push");
    TEST_ASSERT(entry.expiration == 60, "expiration of newly pushed");
    
    // Heap empty (150 deleted)
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == -1, "heap empty after deleted remaining");
    
    timeout_heap_destroy(h);
    printf("  Passed\n");
}

// Test 6: capacity growth
static void test_growth(void) {
    printf("Test 6: capacity growth...\n");
    TimeoutHeap *h = timeout_heap_create(2); // tiny capacity
    TEST_ASSERT(h != NULL, "heap creation");
    TEST_ASSERT(h->capacity == 2, "initial capacity");
    
    // Push 10 elements
    for (int i = 0; i < 10; i++) {
        int ret = timeout_heap_push(h, i * 10, NULL);
        TEST_ASSERT(ret == 0, "push success");
    }
    TEST_ASSERT(h->size == 10, "size after pushes");
    TEST_ASSERT(h->capacity >= 10, "capacity grown");
    
    // Pop all, verify order
    uint64_t prev = 0;
    for (int i = 0; i < 10; i++) {
        TimeoutEntry entry;
        int ret = timeout_heap_pop(h, &entry);
        TEST_ASSERT(ret == 0, "pop success");
        TEST_ASSERT(entry.expiration == prev, "order after growth");
        prev = entry.expiration + 10;
    }
    
    timeout_heap_destroy(h);
    printf("  Passed\n");
}

// Comparison function for HeapElement
static int compare_elements(const void *a, const void *b) {
    const HeapElement *ea = (const HeapElement *)a;
    const HeapElement *eb = (const HeapElement *)b;
    if (ea->expiration < eb->expiration) return -1;
    if (ea->expiration > eb->expiration) return 1;
    return 0;
}

// Test 8: large number of elements, random order, with some cancellations
static void test_large_random(size_t num_elements, size_t cancel_percent) {
    printf("Test 8: large random test (%zu elements, cancel %zu%%)...\n", 
           num_elements, cancel_percent);
    TimeoutHeap *h = timeout_heap_create(100);
    TEST_ASSERT(h != NULL, "heap creation");
    
    // Generate elements with unique expiration values
    HeapElement *elems = malloc(num_elements * sizeof(HeapElement));
    TEST_ASSERT(elems != NULL, "alloc elems");
    
    // Create array of unique expiration values
    uint64_t *expirations = malloc(num_elements * sizeof(uint64_t));
    TEST_ASSERT(expirations != NULL, "alloc expirations");
    for (size_t i = 0; i < num_elements; i++) {
        expirations[i] = i; // base unique value
    }
    // Shuffle using Fisher-Yates
    srand((unsigned int)time(NULL) ^ 0x1234);
    for (size_t i = num_elements - 1; i > 0; i--) {
        size_t j = rand() % (i + 1);
        uint64_t tmp = expirations[i];
        expirations[i] = expirations[j];
        expirations[j] = tmp;
    }
    
    for (size_t i = 0; i < num_elements; i++) {
        elems[i].expiration = expirations[i];
        elems[i].data = malloc(1); // unique pointer
        elems[i].cancelled = 0;
        int ret = timeout_heap_push(h, elems[i].expiration, elems[i].data);
        TEST_ASSERT(ret == 0, "push success");
    }
    free(expirations);
    
    // Cancel some random elements
    size_t cancel_count = num_elements * cancel_percent / 100;
    for (size_t i = 0; i < cancel_count; i++) {
        size_t idx = rand() % num_elements;
        if (!elems[idx].cancelled) {
            int ret = timeout_heap_cancel(h, elems[idx].expiration, elems[idx].data);
            TEST_ASSERT(ret == 0, "cancel success");
            elems[idx].cancelled = 1;
        }
    }
    
    // Build array of non-cancelled elements, sort by expiration
    HeapElement *remaining = malloc(num_elements * sizeof(HeapElement));
    TEST_ASSERT(remaining != NULL, "alloc remaining");
    size_t remain_count = 0;
    for (size_t i = 0; i < num_elements; i++) {
        if (!elems[i].cancelled) {
            remaining[remain_count++] = elems[i];
        }
    }
    qsort(remaining, remain_count, sizeof(HeapElement), compare_elements);
    
    // Pop and verify order
    uint64_t last_expiration = 0;
    for (size_t i = 0; i < remain_count; i++) {
        TimeoutEntry entry;
        int ret = timeout_heap_pop(h, &entry);
        TEST_ASSERT(ret == 0, "pop success");
        TEST_ASSERT(entry.expiration == remaining[i].expiration, "expiration order");
        TEST_ASSERT(entry.data == remaining[i].data, "data matches");
        TEST_ASSERT(entry.expiration >= last_expiration, "non-decreasing");
        last_expiration = entry.expiration;
        free(entry.data);
    }
    
    // Heap empty
    TimeoutEntry entry;
    int ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == -1, "heap empty after all pops");
    
    // Free cancelled elements data
    for (size_t i = 0; i < num_elements; i++) {
        if (elems[i].cancelled) {
            free(elems[i].data);
        }
    }
    
    free(elems);
    free(remaining);
    timeout_heap_destroy(h);
    printf("  Passed\n");
}

// Test 7: peek should not remove
static void test_peek_no_remove(void) {
    printf("Test 7: peek does not remove...\n");
    TimeoutHeap *h = timeout_heap_create(10);
    TEST_ASSERT(h != NULL, "heap creation");
    
    int data1 = 1, data2 = 2;
    timeout_heap_push(h, 100, &data1);
    timeout_heap_push(h, 200, &data2);
    
    TimeoutEntry entry;
    int ret = timeout_heap_peek(h, &entry);
    TEST_ASSERT(ret == 0, "peek success");
    TEST_ASSERT(entry.expiration == 100, "peek expiration");
    TEST_ASSERT(h->size == 2, "size unchanged after peek");
    
    ret = timeout_heap_peek(h, &entry);
    TEST_ASSERT(ret == 0, "second peek success");
    TEST_ASSERT(entry.expiration == 100, "peek same element");
    
    // Pop should still give same element
    ret = timeout_heap_pop(h, &entry);
    TEST_ASSERT(ret == 0, "pop after peek");
    TEST_ASSERT(entry.expiration == 100, "pop expiration");
    
    timeout_heap_destroy(h);
    printf("  Passed\n");
}

int main(void) {
    printf("=== Timeout Heap Tests ===\n");
    
    test_basic();
    test_ordering(100);       // moderate number
    test_ordering(1000);      // more elements
    test_cancel();
    test_duplicate_expirations();
    test_deleted_root();
    test_growth();
    test_peek_no_remove();
    test_large_random(5000, 10);   // 5000 elements, 10% cancelled
    test_large_random(10000, 5);   // 10000 elements, 5% cancelled
    
    printf("\n=== All tests passed! ===\n");
    return 0;
}