// test_etcp.c - Unit tests for ETCP protocol
#include "etcp.h"
#include "u_async.h"
#include "ll_queue.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <assert.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)

// Mock callback storage
typedef struct {
    uint8_t* data;
    uint16_t len;
    struct ETCP_CONN* epkt;
} mock_packet_t;

#define MAX_MOCK_PACKETS 100
static mock_packet_t mock_packets[MAX_MOCK_PACKETS];
static int mock_packet_count = 0;
static uasync_t* test_ua = NULL;

static void reset_mock_packets(void) {
    for (int i = 0; i < mock_packet_count; i++) {
        free(mock_packets[i].data);
        mock_packets[i].data = NULL;
    }
    mock_packet_count = 0;
}

static void mock_tx_callback(struct ETCP_CONN* epkt, uint8_t* data, uint16_t len, void* arg) {
    (void)arg;
    assert(mock_packet_count < MAX_MOCK_PACKETS);
    mock_packets[mock_packet_count].data = malloc(len);
    assert(mock_packets[mock_packet_count].data);
    memcpy(mock_packets[mock_packet_count].data, data, len);
    mock_packets[mock_packet_count].len = len;
    mock_packets[mock_packet_count].epkt = epkt;
    mock_packet_count++;
}



// Test 1: Basic initialization and cleanup
int test_init_free(void) {
    printf("\n=== Test 1: Initialization and cleanup ===\n");
    
    struct ETCP_CONN* epkt = etcp_create(test_ua);
    TEST_ASSERT(epkt != NULL, "etcp_create returns non-NULL");
    TEST_ASSERT(epkt->tx_queue != NULL, "tx_queue created");
    TEST_ASSERT(epkt->output_queue != NULL, "output_queue created");
    
    etcp_free(epkt);
    TEST_ASSERT(1, "etcp_free completes without crash");
    
    return 0;
}

// Test 2: Set callback
int test_set_callback(void) {
    printf("\n=== Test 2: Set callback ===\n");
    
    struct ETCP_CONN* epkt = etcp_create(test_ua);
    TEST_ASSERT(epkt != NULL, "etcp_create");
    
    etcp_set_callback(epkt, mock_tx_callback, NULL);
    // Callback set, no easy way to verify except through tx
    etcp_free(epkt);
    
    return 0;
}

// Test 3: Put data into tx queue
int test_tx_put(void) {
    printf("\n=== Test 3: TX queue put ===\n");
    
    struct ETCP_CONN* epkt = etcp_create(test_ua);
    TEST_ASSERT(epkt != NULL, "etcp_create");
    
    uint8_t test_data[] = {0x01, 0x02, 0x03, 0x04, 0x05};
    int result = etcp_tx_put(epkt, test_data, sizeof(test_data));
    TEST_ASSERT(result == 0, "etcp_tx_put succeeds");
    
    int queue_size = etcp_tx_queue_size(epkt);
    TEST_ASSERT(queue_size == 1, "tx queue size is 1");
    
    etcp_free(epkt);
    
    return 0;
}

// Test 4: Simple packet transmission (without bandwidth limit)
int test_simple_tx(void) {
    printf("\n=== Test 4: Simple transmission ===\n");
    
    reset_mock_packets();
    
    struct ETCP_CONN* epkt = etcp_create(test_ua);
    TEST_ASSERT(epkt != NULL, "etcp_create");
    
    // Set high bandwidth to avoid limiting
    etcp_set_bandwidth(epkt, 65535);
    
    etcp_set_callback(epkt, mock_tx_callback, NULL);
    
    uint8_t test_data[] = "Hello ETCP!";
    int result = etcp_tx_put(epkt, test_data, sizeof(test_data));
    TEST_ASSERT(result == 0, "etcp_tx_put succeeds");
    
    // Transmission may happen via queue callback
    // We can't easily verify transmission in this simple test
    // Just ensure no crash occurred
    TEST_ASSERT(etcp_tx_queue_size(epkt) >= 0, "queue size non-negative");
    
    etcp_free(epkt);
    
    return 0;
}

// Test 5: Packet parsing (rx_input)
int test_rx_input(void) {
    printf("\n=== Test 5: RX input parsing ===\n");
    
    struct ETCP_CONN* epkt = etcp_create(test_ua);
    TEST_ASSERT(epkt != NULL, "etcp_create");
    
    // Create a simple packet: id=1, timestamp=100, hdr=0, payload "test"
    uint8_t packet[] = {
        0x00, 0x01, // id = 1
        0x00, 0x64, // timestamp = 100
        0x00,       // hdr = 0 (payload)
        't', 'e', 's', 't'
    };
    
    int result = etcp_rx_input(epkt, packet, sizeof(packet));
    TEST_ASSERT(result == 0, "etcp_rx_input succeeds");
    
    // Check that output queue has the data
    ll_queue_t* output = etcp_get_output_queue(epkt);
    TEST_ASSERT(output != NULL, "output queue exists");
    
    int output_count = queue_entry_count(output);
    TEST_ASSERT(output_count == 1, "output queue has 1 packet");
    
    // Verify payload
    ll_entry_t* entry = queue_entry_get(output);
    TEST_ASSERT(entry != NULL, "got entry from output queue");
    
    uint8_t* data = ll_entry_data(entry);
    size_t data_len = ll_entry_size(entry);
    TEST_ASSERT(data_len == 4, "payload length is 4");
    TEST_ASSERT(memcmp(data, "test", 4) == 0, "payload matches");
    
    queue_entry_free(entry);
    
    etcp_free(epkt);
    
    return 0;
}

// Test 6: Packet reordering
int test_reordering(void) {
    printf("\n=== Test 6: Packet reordering ===\n");
    
    struct ETCP_CONN* epkt = etcp_create(test_ua);
    TEST_ASSERT(epkt != NULL, "etcp_create");
    
    // Create packets with IDs 1, 2, 3
    uint8_t packet1[] = {
        0x00, 0x01, // id = 1
        0x00, 0x10, // timestamp
        0x00,       // hdr = 0
        'a'
    };
    
    uint8_t packet2[] = {
        0x00, 0x02, // id = 2
        0x00, 0x20, // timestamp
        0x00,       // hdr = 0
        'b'
    };
    
    uint8_t packet3[] = {
        0x00, 0x03, // id = 3
        0x00, 0x30, // timestamp
        0x00,       // hdr = 0
        'c'
    };
    
    // Receive in wrong order: 2, 1, 3
    etcp_rx_input(epkt, packet2, sizeof(packet2));
    etcp_rx_input(epkt, packet1, sizeof(packet1));
    etcp_rx_input(epkt, packet3, sizeof(packet3));
    
    // Check output queue - should have all 3 packets in correct order
    ll_queue_t* output = etcp_get_output_queue(epkt);
    TEST_ASSERT(queue_entry_count(output) == 3, "all 3 packets in output");
    
    // Verify order: 1, 2, 3
    ll_entry_t* entry;
    char expected[] = {'a', 'b', 'c'};
    int idx = 0;
    
    while ((entry = queue_entry_get(output)) != NULL) {
        uint8_t* data = ll_entry_data(entry);
        size_t len = ll_entry_size(entry);
        TEST_ASSERT(len == 1, "payload length 1");
        TEST_ASSERT(data[0] == expected[idx], "correct packet order");
        idx++;
        queue_entry_free(entry);
    }
    
    TEST_ASSERT(idx == 3, "all packets processed");
    
    etcp_free(epkt);
    
    return 0;
}

// Test 7: Metrics update
int test_metrics(void) {
    printf("\n=== Test 7: Metrics ===\n");
    
    struct ETCP_CONN* epkt = etcp_create(test_ua);
    TEST_ASSERT(epkt != NULL, "etcp_create");
    
    // Initial metrics should be zero
    TEST_ASSERT(etcp_get_rtt(epkt) == 0, "initial RTT is 0");
    TEST_ASSERT(etcp_get_jitter(epkt) == 0, "initial jitter is 0");
    
    // Send a packet with ACK to update metrics
    // This requires more complex setup with round-trip
    
    etcp_free(epkt);
    
    return 0;
}

int main(void) {
    printf("Starting ETCP tests...\n");
    
    // Initialize uasync for timers
    uasync_t* test_ua = uasync_create();
    TEST_ASSERT(test_ua != NULL, "create uasync instance");
    uasync_init_instance(test_ua);
    
    int failures = 0;
    
    failures += test_init_free();
    failures += test_set_callback();
    failures += test_tx_put();
    failures += test_simple_tx();
    failures += test_rx_input();
    failures += test_reordering();
    failures += test_metrics();
    
    printf("\n=== Summary ===\n");
    if (failures == 0) {
        printf("All tests passed!\n");
    } else {
        printf("%d test(s) failed.\n", failures);
    }
    
    reset_mock_packets();
    
    uasync_destroy(test_ua);
    
    return failures == 0 ? 0 : 1;
}