// test_etcp_crypto.c - Test ETCP encryption/decryption with secure channel
#include "../src/secure_channel.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>

// Forward declaration for ETCP_CONN structure
struct ETCP_CONN {
    struct ETCP_CONN* next;
    int mtu;
    uint8_t state;
    struct secure_channel crypto_ctx;
    uint64_t peer_node_id;
    void* input_queue;
    void* output_queue;
    void* rx_list;
    void* sent_list;
    uint16_t rtt_last;
    uint16_t rtt_avg_10;
    uint16_t rtt_avg_100;
    uint16_t jitter;
    uint16_t bandwidth;
    uint32_t bytes_sent_total;
    uint16_t last_sent_timestamp;
    uint32_t bytes_allowed;
    uint32_t retransmissions_count;
    uint32_t ack_packets_count;
    uint32_t control_packets_count;
    uint32_t total_packets_sent;
    uint32_t unique_packets_sent;
    uint32_t bytes_received_total;
    uint16_t next_tx_id;
    uint16_t last_sent_id;
    uint16_t last_rx_id;
    uint16_t last_delivered_id;
    void* next_tx_timer;
    void* retransmit_timer;
    uint16_t rtt_history[100];
    uint8_t rtt_history_idx;
    uint8_t rtt_history_count;
    uint16_t pending_ack_ids[32];
    uint16_t pending_ack_timestamps[32];
    uint8_t pending_ack_count;
    uint16_t pending_retransmit_ids[32];
    uint8_t pending_retransmit_count;
    uint32_t unacked_bytes;
    uint32_t window_size;
    uint16_t last_acked_id;
    uint16_t last_rx_ack_id;
    uint16_t retrans_timer_period;
    uint16_t next_retrans_time;
    uint8_t window_blocked;
    uint16_t oldest_missing_id;
    uint16_t missing_since_time;
};

// Test configuration
#define TEST_DATA "Hello, encrypted world!"
#define TEST_DATA_LEN 23

// Simple test for secure channel encryption/decryption
static int test_secure_channel_crypto(void) {
    printf("=== Testing Secure Channel Crypto ===\n");
    
    // Create test keys
    struct SC_MYKEYS server_keys, client_keys;
    
    // Use fixed test keys
    for (int i = 0; i < SC_PRIVKEY_SIZE; i++) {
        server_keys.private_key[i] = i & 0xFF;
        client_keys.private_key[i] = (i + 128) & 0xFF;
    }
    
    // Generate corresponding public keys (simplified for test)
    for (int i = 0; i < SC_PUBKEY_SIZE; i++) {
        server_keys.public_key[i] = (i * 2) & 0xFF;
        client_keys.public_key[i] = (i * 2 + 1) & 0xFF;
    }
    
    // Initialize server context
    sc_context_t server_ctx;
    if (sc_init_ctx(&server_ctx, &server_keys) != SC_OK) {
        printf("ERROR: Failed to initialize server crypto context\n");
        return -1;
    }
    printf("✓ Server crypto context initialized\n");
    
    // Initialize client context
    sc_context_t client_ctx;
    if (sc_init_ctx(&client_ctx, &client_keys) != SC_OK) {
        printf("ERROR: Failed to initialize client crypto context\n");
        return -1;
    }
    printf("✓ Client crypto context initialized\n");
    
    // For this simple test, we'll manually set peer keys, session ready, and session key
    // to bypass the ECC key exchange which requires proper ECC initialization
    memcpy(client_ctx.peer_public_key, server_keys.public_key, SC_PUBKEY_SIZE);
    memcpy(server_ctx.peer_public_key, client_keys.public_key, SC_PUBKEY_SIZE);
    client_ctx.peer_key_set = 1;
    client_ctx.session_ready = 1;  // This is crucial for encryption to work
    server_ctx.peer_key_set = 1;
    server_ctx.session_ready = 1;  // This is crucial for encryption to work
    
    // Set test session keys manually (16 bytes for AES-128)
    // In a real implementation, this would be derived from ECDH shared secret
    uint8_t test_session_key[SC_SESSION_KEY_SIZE];
    for (int i = 0; i < SC_SESSION_KEY_SIZE; i++) {
        test_session_key[i] = i + 100;
    }
    memcpy(client_ctx.session_key, test_session_key, SC_SESSION_KEY_SIZE);
    memcpy(server_ctx.session_key, test_session_key, SC_SESSION_KEY_SIZE);
    
    printf("✓ Peer public keys and session keys set manually\n");
    
    // Test data
    uint8_t plaintext[] = TEST_DATA;
    uint8_t ciphertext[256];
    uint8_t decrypted[256];
    size_t ciphertext_len, decrypted_len;
    
    // Test encryption from client to server
    printf("\n=== Testing Client to Server Encryption ===\n");
    if (sc_encrypt(&client_ctx, plaintext, TEST_DATA_LEN, ciphertext, &ciphertext_len) != SC_OK) {
        printf("ERROR: Encryption failed\n");
        return -1;
    }
    printf("✓ Encrypted %zu bytes to %zu bytes\n", (size_t)TEST_DATA_LEN, ciphertext_len);
    
    // Test decryption by server
    if (sc_decrypt(&server_ctx, ciphertext, ciphertext_len, decrypted, &decrypted_len) != SC_OK) {
        printf("ERROR: Decryption failed\n");
        return -1;
    }
    printf("✓ Decrypted %zu bytes\n", decrypted_len);
    
    // Verify decrypted data
    if (decrypted_len != TEST_DATA_LEN || memcmp(plaintext, decrypted, TEST_DATA_LEN) != 0) {
        printf("ERROR: Decrypted data doesn't match original\n");
        printf("  Original:  '%.*s'\n", TEST_DATA_LEN, plaintext);
        printf("  Decrypted: '%.*s'\n", (int)decrypted_len, decrypted);
        return -1;
    }
    printf("✓ Decrypted data matches original\n");
    
    // Test encryption from server to client
    printf("\n=== Testing Server to Client Encryption ===\n");
    if (sc_encrypt(&server_ctx, plaintext, TEST_DATA_LEN, ciphertext, &ciphertext_len) != SC_OK) {
        printf("ERROR: Server encryption failed\n");
        return -1;
    }
    printf("✓ Server encrypted %zu bytes to %zu bytes\n", (size_t)TEST_DATA_LEN, ciphertext_len);
    
    if (sc_decrypt(&client_ctx, ciphertext, ciphertext_len, decrypted, &decrypted_len) != SC_OK) {
        printf("ERROR: Client decryption failed\n");
        return -1;
    }
    printf("✓ Client decrypted %zu bytes\n", decrypted_len);
    
    if (decrypted_len != TEST_DATA_LEN || memcmp(plaintext, decrypted, TEST_DATA_LEN) != 0) {
        printf("ERROR: Client decrypted data doesn't match original\n");
        return -1;
    }
    printf("✓ Client decrypted data matches original\n");
    
    // Test with different data
    printf("\n=== Testing with Different Data ===\n");
    uint8_t test_data2[] = "The quick brown fox jumps over the lazy dog";
    size_t test_data2_len = sizeof(test_data2) - 1;
    
    if (sc_encrypt(&client_ctx, test_data2, test_data2_len, ciphertext, &ciphertext_len) != SC_OK) {
        printf("ERROR: Encryption of test data 2 failed\n");
        return -1;
    }
    
    if (sc_decrypt(&server_ctx, ciphertext, ciphertext_len, decrypted, &decrypted_len) != SC_OK) {
        printf("ERROR: Decryption of test data 2 failed\n");
        return -1;
    }
    
    if (decrypted_len != test_data2_len || memcmp(test_data2, decrypted, test_data2_len) != 0) {
        printf("ERROR: Test data 2 decryption mismatch\n");
        return -1;
    }
    printf("✓ Different data test passed\n");
    
    printf("\n=== All Crypto Tests Passed! ===\n");
    return 0;
}

// Test ETCP connection crypto
static int test_etcp_connection_crypto(void) {
    printf("\n=== Testing ETCP Connection Crypto ===\n");
    
    // Create a simple ETCP connection
    struct ETCP_CONN* conn = calloc(1, sizeof(struct ETCP_CONN));
    if (!conn) {
        printf("ERROR: Failed to allocate ETCP connection\n");
        return -1;
    }
    
    // Initialize connection basic parameters
    conn->mtu = 1500;
    conn->state = 1; // initialized
    conn->peer_node_id = 0; // no peer yet
    
    // Create test keys
    struct SC_MYKEYS keys;
    for (int i = 0; i < SC_PRIVKEY_SIZE; i++) {
        keys.private_key[i] = i & 0xFF;
        keys.public_key[i] = (i + 64) & 0xFF;
    }
    
    // Initialize crypto context
    if (sc_init_ctx(&conn->crypto_ctx, &keys) != SC_OK) {
        printf("ERROR: Failed to initialize connection crypto\n");
        free(conn);
        return -1;
    }
    printf("✓ ETCP connection crypto initialized\n");
    
    // Test that we can use the crypto context for basic operations
    uint8_t test_data[] = "ETCP connection test";
    uint8_t encrypted[256];
    uint8_t decrypted[256];
    size_t encrypted_len, decrypted_len;
    
    // Self-encryption test (set our own public key as peer and session key)
    memcpy(conn->crypto_ctx.peer_public_key, keys.public_key, SC_PUBKEY_SIZE);
    conn->crypto_ctx.peer_key_set = 1;
    conn->crypto_ctx.session_ready = 1;
    
    // Set test session key for self-encryption
    memcpy(conn->crypto_ctx.session_key, keys.public_key, SC_SESSION_KEY_SIZE); // Use public key as session key for test
    
    if (sc_encrypt(&conn->crypto_ctx, test_data, sizeof(test_data)-1, encrypted, &encrypted_len) != SC_OK) {
        printf("ERROR: Connection encryption failed\n");
        free(conn);
        return -1;
    }
    printf("✓ Connection encryption works\n");
    
    if (sc_decrypt(&conn->crypto_ctx, encrypted, encrypted_len, decrypted, &decrypted_len) != SC_OK) {
        printf("ERROR: Connection decryption failed\n");
        free(conn);
        return -1;
    }
    
    if (decrypted_len != sizeof(test_data)-1 || memcmp(test_data, decrypted, sizeof(test_data)-1) != 0) {
        printf("ERROR: Connection decryption mismatch\n");
        free(conn);
        return -1;
    }
    printf("✓ Connection decryption works\n");
    
    free(conn);
    printf("✓ ETCP connection crypto test passed\n");
    return 0;
}

int main(void) {
    printf("=== ETCP Crypto Test Suite ===\n");
    
    int result1 = test_secure_channel_crypto();
    if (result1 != 0) {
        printf("ERROR: Secure channel crypto test failed\n");
        return 1;
    }
    
    int result2 = test_etcp_connection_crypto();
    if (result2 != 0) {
        printf("ERROR: ETCP connection crypto test failed\n");
        return 1;
    }
    
    printf("\n🎉 All crypto tests passed successfully!\n");
    return 0;
}