#include "u_async.h"
#include "ll_queue.h"
#include "pkt_normalizer.h"
#include "settings.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <assert.h>
#include <time.h>

#define TEST_ASSERT(cond, msg) \
    do { \
        if (!(cond)) { \
            printf("FAIL: %s (line %d)\n", msg, __LINE__); \
            return 1; \
        } else { \
            printf("PASS: %s\n", msg); \
        } \
    } while(0)

/* Структура для хранения тестового пакета */
typedef struct {
    uint8_t* data;
    size_t len;
} test_packet_t;

#define MAX_TEST_PACKETS 20
static test_packet_t sent_packets[MAX_TEST_PACKETS];
static test_packet_t received_packets[MAX_TEST_PACKETS];
static int sent_count = 0;
static int received_count = 0;

static void reset_test_data(void) {
    for (int i = 0; i < sent_count; i++) {
        free(sent_packets[i].data);
        sent_packets[i].data = NULL;
    }
    for (int i = 0; i < received_count; i++) {
        free(received_packets[i].data);
        received_packets[i].data = NULL;
    }
    sent_count = 0;
    received_count = 0;
}

static void add_sent_packet(const uint8_t* data, size_t len) {
    assert(sent_count < MAX_TEST_PACKETS);
    sent_packets[sent_count].data = malloc(len);
    assert(sent_packets[sent_count].data);
    memcpy(sent_packets[sent_count].data, data, len);
    sent_packets[sent_count].len = len;
    sent_count++;
}

static void add_received_packet(const uint8_t* data, size_t len) {
    assert(received_count < MAX_TEST_PACKETS);
    received_packets[received_count].data = malloc(len);
    assert(received_packets[received_count].data);
    memcpy(received_packets[received_count].data, data, len);
    received_packets[received_count].len = len;
    received_count++;
}

static int compare_packets(void) {
    if (sent_count != received_count) {
        printf("Packet count mismatch: sent=%d, received=%d\n", sent_count, received_count);
        return -1;
    }
    for (int i = 0; i < sent_count; i++) {
        if (sent_packets[i].len != received_packets[i].len) {
            printf("Packet %d length mismatch: sent=%zu, received=%zu\n", 
                   i, sent_packets[i].len, received_packets[i].len);
            return -1;
        }
        if (memcmp(sent_packets[i].data, received_packets[i].data, sent_packets[i].len) != 0) {
            printf("Packet %d data mismatch\n", i);
            return -1;
        }
    }
    return 0;
}

/* Создать тестовый пакет со случайным содержимым */
static ll_entry_t* create_test_packet(size_t data_size) {
    ll_entry_t* entry = queue_entry_new(data_size);
    if (!entry) {
        return NULL;
    }
    uint8_t* data = ll_entry_data(entry);
    for (size_t i = 0; i < data_size; i++) {
        data[i] = rand() & 0xFF;
    }
    return entry;
}

/* Коллбэк для пересылки пакетов от упаковщика к распаковщику */
static void forward_cb(ll_queue_t* q, ll_entry_t* unused, void* arg) {
    (void)unused;
    ll_queue_t* target = (ll_queue_t*)arg;
    while (queue_entry_count(q) > 0) {
        ll_entry_t* e = queue_entry_get(q);
        queue_entry_put(target, e); // Transfer ownership
    }
    queue_resume_callback(q);
}

/* Коллбэк для сбора полученных пакетов */
static void receive_cb(ll_queue_t* q, ll_entry_t* unused, void* arg) {
    (void)unused;
    (void)arg;
    while (queue_entry_count(q) > 0) {
        ll_entry_t* e = queue_entry_get(q);
        uint8_t* data = ll_entry_data(e);
        size_t len = ll_entry_size(e);
        add_received_packet(data, len);
        queue_entry_free(e);
    }
    queue_resume_callback(q);
}

/* Принудительно обработать все очереди, вызывая коллбэки для непустых очередей
 * Это нужно потому что в тестовой среде нет uasync_mainloop
 */
static void process_queues(pkt_normalizer_pair* pair) {
    int iterations = 0;
    int processed;
    
    /* Сначала принудительно сбросим все флаги callback_suspended, потому что
     * uasync таймауты не работают в тестовой среде */
    pair->packer->input->callback_suspended = 0;
    pair->packer->output->callback_suspended = 0;
    pair->unpacker->input->callback_suspended = 0;
    pair->unpacker->output->callback_suspended = 0;
    
    do {
        processed = 0;
        iterations++;
        
        /* Обработать входную очередь упаковщика */
        if (pair->packer->input->callback && 
            queue_entry_count(pair->packer->input) > 0) {
            pair->packer->input->callback(pair->packer->input, 
                                        pair->packer->input->head, 
                                        pair->packer->input->callback_arg);
            processed = 1;
        }
        
        /* Обработать выходную очередь упаковщика (forward_cb) */
        if (pair->packer->output->callback && 
            queue_entry_count(pair->packer->output) > 0) {
            pair->packer->output->callback(pair->packer->output, 
                                         pair->packer->output->head, 
                                         pair->packer->output->callback_arg);
            processed = 1;
        }
        
        /* Обработать входную очередь распаковщика */
        if (pair->unpacker->input->callback && 
            queue_entry_count(pair->unpacker->input) > 0) {
            pair->unpacker->input->callback(pair->unpacker->input, 
                                          pair->unpacker->input->head, 
                                          pair->unpacker->input->callback_arg);
            processed = 1;
        }
        
        /* Обработать выходную очередь распаковщика (receive_cb) */
        if (pair->unpacker->output->callback && 
            queue_entry_count(pair->unpacker->output) > 0) {
            pair->unpacker->output->callback(pair->unpacker->output, 
                                           pair->unpacker->output->head, 
                                           pair->unpacker->output->callback_arg);
            processed = 1;
        }
        
        if (iterations > 50) {
            printf("ERROR: process_queues infinite loop detected\n");
            break;
        }
    } while (processed);
}

/* Тест 1: мелкие пакеты */
static int test_small_packets(pkt_normalizer_pair* pair) {
    printf("\n--- Test 1: small packets ---\n");
    reset_test_data();
    pkt_normalizer_reset_error_count(pair->unpacker);
    
    const char* packets[] = {"hello", "world", "test"};
    for (size_t i = 0; i < sizeof(packets)/sizeof(packets[0]); i++) {
        size_t len = strlen(packets[i]);
        ll_entry_t* entry = queue_entry_new(len);
        TEST_ASSERT(entry != NULL, "create packet");
        memcpy(ll_entry_data(entry), packets[i], len);
        add_sent_packet((const uint8_t*)packets[i], len);
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put into packer input");
    }
    
    process_queues(pair);
    
    TEST_ASSERT(compare_packets() == 0, "small packets comparison");
    return 0;
}

/* Тест 2: средний пакет (не фрагментируется) */
static int test_medium_packet(pkt_normalizer_pair* pair) {
    printf("\n--- Test 2: medium packet ---\n");
    reset_test_data();
    pkt_normalizer_reset_error_count(pair->unpacker);
    
    size_t data_size = 500;
    ll_entry_t* entry = create_test_packet(data_size);
    TEST_ASSERT(entry != NULL, "create medium packet");
    
    uint8_t* data = ll_entry_data(entry);
    add_sent_packet(data, data_size);
    
    TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put medium packet");
    
    process_queues(pair);
    
    TEST_ASSERT(compare_packets() == 0, "medium packet comparison");
    return 0;
}

/* Тест 3: большой пакет с фрагментацией */
static int test_fragmentation(pkt_normalizer_pair* pair) {
    printf("\n--- Test 3: fragmentation ---\n");
    reset_test_data();
    pkt_normalizer_reset_error_count(pair->unpacker);
    
    /* Сохранить оригинальный размер фрагмента */
    int original_fragment_size = settings.max_fragment_size;
    
    /* Установить маленький размер фрагмента для тестирования */
    settings.max_fragment_size = 300;
    
    size_t data_size = 1000; /* Должен быть фрагментирован */
    ll_entry_t* entry = create_test_packet(data_size);
    TEST_ASSERT(entry != NULL, "create big packet");
    
    uint8_t* data = ll_entry_data(entry);
    add_sent_packet(data, data_size);
    
    TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put big packet");
    
    process_queues(pair);
    
    /* Проверить, что не было ошибок сборки */
    TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no fragmentation errors");
    
    /* Восстановить оригинальный размер */
    settings.max_fragment_size = original_fragment_size;
    
    TEST_ASSERT(compare_packets() == 0, "fragmented packet comparison");
    return 0;
}

/* Тест 4: несколько пакетов разного размера */
static int test_mixed_packets(pkt_normalizer_pair* pair) {
    printf("\n--- Test 4: mixed size packets ---\n");
    reset_test_data();
    pkt_normalizer_reset_error_count(pair->unpacker);
    
    /* Отправить пакеты разных размеров */
    size_t sizes[] = {10, 100, 50, 300, 5};
    for (size_t i = 0; i < sizeof(sizes)/sizeof(sizes[0]); i++) {
        ll_entry_t* entry = create_test_packet(sizes[i]);
        TEST_ASSERT(entry != NULL, "create mixed packet");
        
        uint8_t* data = ll_entry_data(entry);
        add_sent_packet(data, sizes[i]);
        
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put mixed packet");
    }
    
    process_queues(pair);
    
    TEST_ASSERT(compare_packets() == 0, "mixed packets comparison");
    return 0;
}

/* Тест 6: стресс-тест со случайными размерами пакетов */
static int test_stress_random(pkt_normalizer_pair* pair) {
    printf("\n--- Test 6: stress test with random packets ---\n");
    const int NUM_PACKETS = 100;
    const int MAX_PACKET_SIZE = 4000;
    int original_fragment_size = settings.max_fragment_size;
    
    /* Часть 1: случайные пакеты без изменения размера фрагмента */
    for (int i = 0; i < NUM_PACKETS; i++) {
        reset_test_data();
        pkt_normalizer_reset_error_count(pair->unpacker);
        
        /* Случайный размер пакета от 1 до MAX_PACKET_SIZE */
        size_t size = (rand() % MAX_PACKET_SIZE) + 1;
        
        ll_entry_t* entry = create_test_packet(size);
        TEST_ASSERT(entry != NULL, "create random packet");
        
        uint8_t* data = ll_entry_data(entry);
        add_sent_packet(data, size);
        
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put random packet");
        
        process_queues(pair);
        
        /* Проверить, что пакет корректно обработан */
        TEST_ASSERT(compare_packets() == 0, "random packet comparison");
        TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no errors during random packet");
    }
    
    /* Часть 2: фрагментация со случайным размером фрагмента */
    const int FRAGMENT_TESTS = 50;
    for (int i = 0; i < FRAGMENT_TESTS; i++) {
        reset_test_data();
        pkt_normalizer_reset_error_count(pair->unpacker);
        
        /* Случайный размер фрагмента от 50 до 500 */
        int frag_size = (rand() % 451) + 50;
        settings.max_fragment_size = frag_size;
        
        /* Случайный размер пакета от frag_size+1 до 3000 (чтобы гарантировать фрагментацию) */
        size_t packet_size = (rand() % (3000 - frag_size)) + frag_size + 1;
        
        ll_entry_t* entry = create_test_packet(packet_size);
        TEST_ASSERT(entry != NULL, "create fragmentable packet");
        
        uint8_t* data = ll_entry_data(entry);
        add_sent_packet(data, packet_size);
        
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put fragmentable packet");
        
        process_queues(pair);
        
        /* Проверить, что пакет корректно обработан */
        TEST_ASSERT(compare_packets() == 0, "fragmentable packet comparison");
        TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no errors during fragmentation");
    }
    
    settings.max_fragment_size = original_fragment_size;
    return 0;
}

/* Тест 7: граничные случаи и максимальные размеры */
static int test_edge_cases(pkt_normalizer_pair* pair) {
    printf("\n--- Test 7: edge cases and maximum sizes ---\n");
    int original_fragment_size = settings.max_fragment_size;
    
    /* Сбросить данные теста */
    reset_test_data();
    pkt_normalizer_reset_error_count(pair->unpacker);
    
    /* Тест 1: минимальный размер пакета (1 байт) */
    {
        size_t size = 1;
        ll_entry_t* entry = create_test_packet(size);
        TEST_ASSERT(entry != NULL, "create 1-byte packet");
        add_sent_packet(ll_entry_data(entry), size);
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put 1-byte packet");
        process_queues(pair);
        TEST_ASSERT(compare_packets() == 0, "1-byte packet comparison");
        TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no errors for 1-byte packet");
    }
    
    /* Тест 2: граница одно- и двухбайтного заголовка (239/240) */
    {
        reset_test_data();
        pkt_normalizer_reset_error_count(pair->unpacker);
        size_t size = 239;
        ll_entry_t* entry = create_test_packet(size);
        TEST_ASSERT(entry != NULL, "create 239-byte packet");
        add_sent_packet(ll_entry_data(entry), size);
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put 239-byte packet");
        process_queues(pair);
        TEST_ASSERT(compare_packets() == 0, "239-byte packet comparison");
        TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no errors for 239-byte packet");
    }
    {
        reset_test_data();
        pkt_normalizer_reset_error_count(pair->unpacker);
        size_t size = 240;
        ll_entry_t* entry = create_test_packet(size);
        TEST_ASSERT(entry != NULL, "create 240-byte packet");
        add_sent_packet(ll_entry_data(entry), size);
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put 240-byte packet");
        process_queues(pair);
        TEST_ASSERT(compare_packets() == 0, "240-byte packet comparison");
        TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no errors for 240-byte packet");
    }
    
    /* Тест 3: максимальный размер без фрагментации (3839) */
    {
        reset_test_data();
        pkt_normalizer_reset_error_count(pair->unpacker);
        size_t size = 3839;
        ll_entry_t* entry = create_test_packet(size);
        TEST_ASSERT(entry != NULL, "create 3839-byte packet");
        add_sent_packet(ll_entry_data(entry), size);
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put 3839-byte packet");
        process_queues(pair);
        TEST_ASSERT(compare_packets() == 0, "3839-byte packet comparison");
        TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no errors for 3839-byte packet");
    }
    
    /* Тест 4: размер, требующий фрагментации (3840) с обычным max_fragment_size */
    {
        reset_test_data();
        pkt_normalizer_reset_error_count(pair->unpacker);
        size_t size = 3840;
        ll_entry_t* entry = create_test_packet(size);
        TEST_ASSERT(entry != NULL, "create 3840-byte packet");
        add_sent_packet(ll_entry_data(entry), size);
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put 3840-byte packet");
        process_queues(pair);
        TEST_ASSERT(compare_packets() == 0, "3840-byte packet comparison");
        TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no errors for 3840-byte packet");
    }
    
    /* Тест 5: размер чуть меньше max_fragment_size, чтобы не фрагментироваться */
    {
        reset_test_data();
        pkt_normalizer_reset_error_count(pair->unpacker);
        settings.max_fragment_size = 500;
        size_t size = 480; /* достаточно мало, чтобы поместиться с заголовком */
        ll_entry_t* entry = create_test_packet(size);
        TEST_ASSERT(entry != NULL, "create 480-byte packet with fragment size 500");
        add_sent_packet(ll_entry_data(entry), size);
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put packet");
        process_queues(pair);
        TEST_ASSERT(compare_packets() == 0, "packet comparison with fragment size 500");
        TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no errors");
    }
    
    /* Тест 6: размер чуть больше max_fragment_size, чтобы вызвать фрагментацию */
    {
        reset_test_data();
        pkt_normalizer_reset_error_count(pair->unpacker);
        settings.max_fragment_size = 500;
        size_t size = 520; /* потребует фрагментации */
        ll_entry_t* entry = create_test_packet(size);
        TEST_ASSERT(entry != NULL, "create 520-byte packet with fragment size 500");
        add_sent_packet(ll_entry_data(entry), size);
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put packet");
        process_queues(pair);
        TEST_ASSERT(compare_packets() == 0, "packet comparison with fragmentation");
        TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no errors");
    }
    
    /* Тест 7: максимальный размер пакета с фрагментацией (65535) - ограничим 10000 для скорости */
    {
        reset_test_data();
        pkt_normalizer_reset_error_count(pair->unpacker);
        settings.max_fragment_size = 1000;
        size_t size = 10000;
        ll_entry_t* entry = create_test_packet(size);
        TEST_ASSERT(entry != NULL, "create 10000-byte fragmented packet");
        add_sent_packet(ll_entry_data(entry), size);
        TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put large fragmented packet");
        process_queues(pair);
        TEST_ASSERT(compare_packets() == 0, "large fragmented packet comparison");
        TEST_ASSERT(pkt_normalizer_get_error_count(pair->unpacker) == 0, "no errors for large fragmented packet");
    }
    
    settings.max_fragment_size = original_fragment_size;
    return 0;
}

/* Тест 5: проверка жизненного цикла */
static int test_lifecycle(void) {
    printf("\n--- Test 5: lifecycle ---\n");
    
    pkt_normalizer_pair* pair = pkt_normalizer_pair_init();
    TEST_ASSERT(pair != NULL, "pair initialization");
    TEST_ASSERT(pair->packer != NULL, "packer initialization");
    TEST_ASSERT(pair->unpacker != NULL, "unpacker initialization");
    
    /* Создать тестовые очереди */
    ll_queue_t* test_queue = queue_new();
    TEST_ASSERT(test_queue != NULL, "test queue creation");
    
    /* Отправить один пакет */
    const char* test_data = "test";
    ll_entry_t* entry = queue_entry_new(strlen(test_data));
    TEST_ASSERT(entry != NULL, "create test packet");
    memcpy(ll_entry_data(entry), test_data, strlen(test_data));
    
    TEST_ASSERT(queue_entry_put(pair->packer->input, entry) == 0, "put test packet");
    
    /* Освободить всё */
    pkt_normalizer_pair_deinit(pair);
    queue_free(test_queue);
    
    return 0;
}

int main(void) {
    int result = 0;
    
    printf("=== Packet Normalizer Unit Test ===\n");
    
    srand((unsigned int)time(NULL));
    uasync_init();
    
    pkt_normalizer_pair* pair = pkt_normalizer_pair_init();
    if (!pair) {
        fprintf(stderr, "Failed to initialize packet normalizer pair\n");
        return 1;
    }
    
    /* Настроить цепочку обработки: packer->output -> unpacker->input -> unpacker->output */
    queue_set_callback(pair->packer->output, forward_cb, pair->unpacker->input);
    queue_set_callback(pair->unpacker->output, receive_cb, NULL);
    
    result |= test_lifecycle();
    result |= test_small_packets(pair);
    result |= test_medium_packet(pair);
    result |= test_fragmentation(pair);
    result |= test_mixed_packets(pair);
    result |= test_stress_random(pair);
    result |= test_edge_cases(pair);
    
    pkt_normalizer_pair_deinit(pair);
    reset_test_data();
    
    if (result == 0) {
        printf("\n=== All tests PASSED ===\n");
    } else {
        printf("\n=== Some tests FAILED ===\n");
    }
    
    return result;
}