// test_etcp_stress.c - Stress test for ETCP with packet loss, delay, and reordering
#include "etcp.h"
#include "u_async.h"
#include "ll_queue.h"
#include "simple_uasync.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <assert.h>
#include <time.h>

#define NUM_PACKETS 10000
#define MIN_PACKET_SIZE 1
#define MAX_PACKET_SIZE 1300
#define LOSS_PROBABILITY 0.0      // 0% packet loss for testing reordering
#define REORDER_PROBABILITY 0.1   // 10% reordering for testing
#define QUEUE_MAX_SIZE 5000       // Max packets in delay queue
#define MAX_DELAY_MS 100          // Up to 100ms delay for reordering effect
#define TIME_BASE_MS 0.1          // uasync timebase is 0.1ms

// Packet in the network delay queue
typedef struct delayed_packet {
    uint8_t* data;
    uint16_t len;
    uint32_t delivery_time;       // When to deliver (in timebase units)
    struct delayed_packet* next;
} delayed_packet_t;

// Network emulator structure
typedef struct {
    struct ETCP_CONN* sender;              // ETCP instance that sends
    struct ETCP_CONN* receiver;            // ETCP instance that receives
    delayed_packet_t* queue;     // Delay queue (sorted by delivery time)
    int queue_size;
    uint32_t current_time;       // Current time in timebase units
    uint32_t packets_sent;
    uint32_t packets_lost;
    uint32_t packets_reordered;
    uint32_t packets_delivered;
    uint8_t running;
    void* timer_id;
} network_emulator_t;

// Forward declarations
static void sender_tx_callback(struct ETCP_CONN* epkt, uint8_t* data, uint16_t len, void* arg);
static void receiver_tx_callback(struct ETCP_CONN* epkt, uint8_t* data, uint16_t len, void* arg);
static void deliver_packets(network_emulator_t* net);
static void free_delay_queue(delayed_packet_t* queue);

// Random number generator (simple LCG)
static uint32_t random_state = 123456789;
static uint32_t random_next(void) {
    random_state = random_state * 1103515245 + 12345;
    return random_state;
}

static double random_double(void) {
    return (double)random_next() / (double)UINT32_MAX;
}



// Sender's TX callback - called when ETCP wants to send a packet
static void sender_tx_callback(struct ETCP_CONN* epkt, uint8_t* data, uint16_t len, void* arg) {
    (void)epkt;
    network_emulator_t* net = (network_emulator_t*)arg;
    if (!net || !net->running) return;
    
    net->packets_sent++;
    
    // 10% packet loss - just ignore the packet, ETCP will free the data
    if (random_double() < LOSS_PROBABILITY) {
        net->packets_lost++;
        return;
    }
    
    // Calculate delivery time (current time + random delay up to MAX_DELAY_MS)
    uint32_t delay_ms = (uint32_t)(random_double() * MAX_DELAY_MS);
    uint32_t delivery_time = net->current_time + (delay_ms * 10); // Convert ms to timebase
    
    // Create delayed packet - need to copy data since ETCP owns the original
    delayed_packet_t* pkt = malloc(sizeof(delayed_packet_t));
    if (!pkt) {
        return; // Memory allocation failed, packet is lost
    }
    
    pkt->data = malloc(len);
    if (!pkt->data) {
        free(pkt);
        net->packets_lost++; // Count as loss due to memory failure
        return;
    }
    
    memcpy(pkt->data, data, len);
    pkt->len = len;
    pkt->delivery_time = delivery_time;
    pkt->next = NULL;
    
    // Insert into delay queue
    delayed_packet_t** pp = &net->queue;
    
    // 30% chance to insert at random position (reordering)
    if (random_double() < REORDER_PROBABILITY && net->queue_size > 1) {
        net->packets_reordered++;
        int insert_pos = random_next() % (net->queue_size + 1);
        for (int i = 0; i < insert_pos && *pp; i++) {
            pp = &(*pp)->next;
        }
    } else {
        // Normal insertion (sorted by delivery time)
        while (*pp && (*pp)->delivery_time < delivery_time) {
            pp = &(*pp)->next;
        }
    }
    
    pkt->next = *pp;
    *pp = pkt;
    net->queue_size++;
    
    // Limit queue size (drop first packet if needed)
    if (net->queue_size > QUEUE_MAX_SIZE) {
        delayed_packet_t* first = net->queue;
        if (first) {
            net->queue = first->next;
            free(first->data);
            free(first);
            net->queue_size--;
            net->packets_lost++; // Count as loss due to queue overflow
        }
    }
}

// Receiver's TX callback - called when receiver wants to send ACKs or retransmission requests
static void receiver_tx_callback(struct ETCP_CONN* epkt, uint8_t* data, uint16_t len, void* arg) {
    (void)epkt;
    network_emulator_t* net = (network_emulator_t*)arg;
    if (!net || !net->running) return;
    
    // Forward ACKs/retrans requests directly to sender with minimal delay (1 timebase unit)
    // This allows sender to receive ACKs and retransmit lost packets
    simple_uasync_advance_time(1);
    net->current_time = simple_uasync_get_time();
    etcp_rx_input(net->sender, data, len);
    
    // Note: We don't track these in statistics since they're control packets
}

// Deliver packets whose delivery time has arrived
static void deliver_packets(network_emulator_t* net) {
    while (net->queue && net->queue->delivery_time <= net->current_time) {
        delayed_packet_t* pkt = net->queue;
        net->queue = pkt->next;
        
        // Deliver to receiver
        etcp_rx_input(net->receiver, pkt->data, pkt->len);
        net->packets_delivered++;
        
        free(pkt->data);
        free(pkt);
        net->queue_size--;
    }
}

// Free delay queue
static void free_delay_queue(delayed_packet_t* queue) {
    while (queue) {
        delayed_packet_t* next = queue->next;
        free(queue->data);
        free(queue);
        queue = next;
    }
}

// Generate random packet data
static void generate_packet_data(uint8_t* buffer, uint16_t size, uint32_t seq) {
    // Fill with pattern: sequence number + random data
    for (uint16_t i = 0; i < size; i++) {
        if (i < 4) {
            // First 4 bytes: sequence number
            buffer[i] = (seq >> (8 * i)) & 0xFF;
        } else {
            // Rest: pseudo-random data based on sequence and position
            buffer[i] = (uint8_t)((seq * 7919 + i * 104729) % 256);
        }
    }
}

// Verify received packet
static int verify_packet_data(const uint8_t* data, uint16_t size, uint32_t seq) {
    if (size < 4) return 0;
    
    // Check sequence number
    uint32_t received_seq = 0;
    for (int i = 0; i < 4; i++) {
        received_seq |= ((uint32_t)data[i]) << (8 * i);
    }
    
    if (received_seq != seq) return 0;
    
    // Verify the rest of the data
    for (uint16_t i = 4; i < size; i++) {
        uint8_t expected = (uint8_t)((seq * 7919 + i * 104729) % 256);
        if (data[i] != expected) return 0;
    }
    
    return 1;
}

// Stress test main function
int main(void) {
    printf("Starting ETCP stress test...\n");
    printf("Parameters:\n");
    printf("  Packets: %d\n", NUM_PACKETS);
    printf("  Size range: %d-%d bytes\n", MIN_PACKET_SIZE, MAX_PACKET_SIZE);
    printf("  Loss probability: %.1f%%\n", LOSS_PROBABILITY * 100);
    printf("  Reorder probability: %.1f%%\n", REORDER_PROBABILITY * 100);
    printf("  Max delay: %d ms\n", MAX_DELAY_MS);
    printf("  Queue size: %d packets\n", QUEUE_MAX_SIZE);
    
    // Seed random number generator
    random_state = (uint32_t)time(NULL);
    
    // Initialize uasync instance
    uasync_t* ua = uasync_create();
    if (!ua) {
        fprintf(stderr, "Failed to create uasync instance\n");
        return 1;
    }
    uasync_init_instance(ua);
    
    // Create network emulator
    network_emulator_t net = {0};
    net.running = 1;
    net.current_time = 0;
    
    // Create ETCP instances
    net.sender = etcp_create(ua);
    net.receiver = etcp_create(ua);
    
    if (!net.sender || !net.receiver) {
        printf("ERROR: Failed to create ETCP instances\n");
        return 1;
    }
    
    // Set up callbacks
    etcp_set_callback(net.sender, sender_tx_callback, &net);
    etcp_set_callback(net.receiver, receiver_tx_callback, &net);
    
    // Set reasonable bandwidth for stress test
    etcp_set_bandwidth(net.sender, 50000); // 50k bytes per timebase
    etcp_set_bandwidth(net.receiver, 50000);
    
    // Don't start network timer - we'll advance time manually
    
    printf("\nGenerating and sending %d packets...\n", NUM_PACKETS);
    
    // Generate and send packets
    uint32_t packets_generated = 0;
    uint32_t bytes_generated = 0;
    uint32_t last_time_print = 0;
    
    while (packets_generated < NUM_PACKETS) {
        // Generate random packet size
        uint16_t size = MIN_PACKET_SIZE + (random_next() % (MAX_PACKET_SIZE - MIN_PACKET_SIZE + 1));
        
        // Allocate and fill packet data
        uint8_t* data = malloc(size);
        if (!data) {
            printf("ERROR: Memory allocation failed\n");
            break;
        }
        
        generate_packet_data(data, size, packets_generated);
        
        // Send via ETCP
        if (etcp_tx_put(net.sender, data, size) != 0) {
            printf("ERROR: Failed to queue packet %u\n", packets_generated);
            free(data);
            break;
        }
        
        free(data); // etcp_tx_put makes its own copy
        packets_generated++;
        bytes_generated += size;
        
        // Periodically print progress
        if (packets_generated % 1000 == 0) {
            printf("  Sent %u packets, %u bytes\n", packets_generated, bytes_generated);
        }
        
        // Advance time to allow ETCP to send packets (bandwidth limiting)
        // and process any expired timers (retransmissions, etc.)
        simple_uasync_advance_time(1); // Advance by 1 timebase unit (0.1ms)
        net.current_time = simple_uasync_get_time(); // Keep in sync
        
        // Deliver any packets whose time has come
        deliver_packets(&net);
        
        // Periodically print time progress
        if (net.current_time - last_time_print >= 1000) { // Every 100ms
            printf("  Time: %.1f ms, Queue: %d packets\n", 
                   net.current_time / 10.0, net.queue_size);
            last_time_print = net.current_time;
        }
    }
    
    printf("Finished sending %u packets (%u bytes)\n", packets_generated, bytes_generated);
    
    // Let network deliver remaining packets and wait for retransmissions
    printf("\nDelivering remaining packets and waiting for retransmissions...\n");
    uint32_t start_time = net.current_time;
    for (int i = 0; i < 200000 && (net.queue_size > 0 || i < 1000); i++) {
        simple_uasync_advance_time(10); // Advance 1ms
        net.current_time = simple_uasync_get_time(); // Keep in sync
        deliver_packets(&net);
        
        // Periodically advance more time to speed up retransmission timeouts
        if (i % 10 == 0) {
            simple_uasync_advance_time(100); // Additional 10ms
            net.current_time = simple_uasync_get_time();
            deliver_packets(&net);
        }
        
        if (i % 500 == 0) {
            uint32_t elapsed = net.current_time - start_time;
            printf("  Queue: %d, Time: %.1f ms (elapsed: %.1f ms)\n", 
                   net.queue_size, net.current_time / 10.0, elapsed / 10.0);
        }
    }
    
    // No network timer to stop since we're not using one
    net.running = 0;
    
    // Free any remaining packets in queue
    free_delay_queue(net.queue);
    net.queue = NULL;
    net.queue_size = 0;
    
    // Collect and verify received packets
    printf("\nVerifying received packets...\n");
    
    ll_queue_t* output_queue = etcp_get_output_queue(net.receiver);
    uint32_t packets_received = 0;
    uint32_t bytes_received = 0;
    uint32_t correct_packets = 0;
    uint32_t max_received_seq = 0;
    
    ll_entry_t* entry;
    while ((entry = queue_entry_get(output_queue)) != NULL) {
        uint8_t* data = ll_entry_data(entry);
        uint16_t size = ll_entry_size(entry);
        
        if (size >= 4) {
            uint32_t seq = 0;
            for (int i = 0; i < 4; i++) {
                seq |= ((uint32_t)data[i]) << (8 * i);
            }
            
            if (verify_packet_data(data, size, seq)) {
                correct_packets++;
                if (seq > max_received_seq) {
                    max_received_seq = seq;
                }
            }
        }
        
        packets_received++;
        bytes_received += size;
        queue_entry_free(entry);
    }
    
    // Print statistics
    printf("\n=== Statistics ===\n");
    printf("Packets generated:  %u\n", packets_generated);
    printf("Packets sent:       %u (via ETCP)\n", net.packets_sent);
    printf("Packets lost:       %u (%.1f%%)\n", net.packets_lost, 
           (net.packets_sent > 0) ? (100.0 * net.packets_lost / net.packets_sent) : 0.0);
    printf("Packets reordered:  %u (%.1f%% of delivered)\n", net.packets_reordered,
           (net.packets_delivered > 0) ? (100.0 * net.packets_reordered / net.packets_delivered) : 0.0);
    printf("Packets delivered:  %u (to receiver)\n", net.packets_delivered);
    printf("Packets received:   %u (in output queue)\n", packets_received);
    printf("Correct packets:    %u (%.1f%%)\n", correct_packets,
           (packets_received > 0) ? (100.0 * correct_packets / packets_received) : 0.0);
    printf("Bytes generated:    %u\n", bytes_generated);
    printf("Bytes received:     %u\n", bytes_received);
    
    // Check for missing packets
    uint32_t expected_received = packets_generated - net.packets_lost;
    if (packets_received < expected_received) {
        printf("\nWARNING: Received %u packets, expected ~%u (some may be in flight)\n",
               packets_received, expected_received);
    } else if (packets_received > expected_received) {
        printf("\nWARNING: Received %u packets, expected ~%u (duplicates?)\n",
               packets_received, expected_received);
    }
    
    // Cleanup
    etcp_free(net.sender);
    etcp_free(net.receiver);
    uasync_destroy(ua);
    
    printf("\nStress test completed.\n");
    
    // Consider test successful if we received at least some packets
    // (with losses, we won't get all of them)
    return (correct_packets > 0) ? 0 : 1;
}