// test_etcp_crypto.c - Test ETCP encryption/decryption with secure channel
#include "etcp.h"
#include "config_parser.h"
#include "secure_channel.h"
#include "etcp_connections.h"
#include "packet_buffer.h"
#include "packet_pool.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <fcntl.h>
#include <errno.h>

// Test configuration
#define TEST_PORT 42345
#define CLIENT_PORT 42346
#define TEST_DATA "Hello, encrypted world!"
#define TEST_DATA_LEN 23

// Helper: create UDP socket bound to port
static int create_udp_socket(int port) {
    int fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (fd < 0) return -1;
    
    int flags = fcntl(fd, F_GETFL, 0);
    fcntl(fd, F_SETFL, flags | O_NONBLOCK);
    
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(port);
    
    if (bind(fd, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        close(fd);
        return -1;
    }
    
    return fd;
}

// Helper: generate test keys
static void generate_test_keys(uint8_t* server_priv, uint8_t* server_pub,
                               uint8_t* client_priv, uint8_t* client_pub) {
    // Use fixed test keys for reproducibility
    const char* server_priv_hex = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
    const char* server_pub_hex  = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
    const char* client_priv_hex = "fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210";
    const char* client_pub_hex  = "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";
    
    for (int i = 0; i < 32; i++) {
        sscanf(server_priv_hex + i*2, "%2hhx", &server_priv[i]);
        sscanf(server_pub_hex  + i*2, "%2hhx", &server_pub[i]);
        sscanf(client_priv_hex + i*2, "%2hhx", &client_priv[i]);
        sscanf(client_pub_hex  + i*2, "%2hhx", &client_pub[i]);
    }
}

// Helper: hex to binary
static int hex_to_bin(const char* hex, uint8_t* bin, size_t len) {
    if (strlen(hex) != len * 2) return -1;
    for (size_t i = 0; i < len; i++) {
        if (sscanf(hex + i*2, "%2hhx", &bin[i]) != 1) return -1;
    }
    return 0;
}

int main(void) {
    printf("=== ETCP Crypto Test ===\n");
    
    // Generate test keys
    uint8_t server_priv[32], server_pub[32];
    uint8_t client_priv[32], client_pub[32];
    generate_test_keys(server_priv, server_pub, client_priv, client_pub);
    
    // Create UDP sockets
    int server_fd = create_udp_socket(TEST_PORT);
    int client_fd = create_udp_socket(CLIENT_PORT);
    if (server_fd < 0 || client_fd < 0) {
        printf("ERROR: Failed to create sockets\n");
        return 1;
    }
    
    // Create ETCP instances
    struct ETCP_CONN* server_etcp = calloc(1, sizeof(struct ETCP_CONN));
    struct ETCP_CONN* client_etcp = calloc(1, sizeof(struct ETCP_CONN));
    if (!server_etcp || !client_etcp) {
        printf("ERROR: Failed to allocate ETCP instances\n");
        return 1;
    }
    
    // Initialize secure channels
    server_etcp->crypto_ctx = calloc(1, sizeof(sc_context_t));
    client_etcp->crypto_ctx = calloc(1, sizeof(sc_context_t));
    if (!server_etcp->crypto_ctx || !client_etcp->crypto_ctx) {
        printf("ERROR: Failed to allocate crypto contexts\n");
        return 1;
    }
    
    // Set keys
    memcpy(server_etcp->crypto_ctx->private_key, server_priv, 32);
    memcpy(server_etcp->crypto_ctx->public_key, server_pub, 32);
    server_etcp->crypto_ctx->initialized = 1;
    
    memcpy(client_etcp->crypto_ctx->private_key, client_priv, 32);
    memcpy(client_etcp->crypto_ctx->public_key, client_pub, 32);
    client_etcp->crypto_ctx->initialized = 1;
    
    // Client sets server's public key as peer
    memcpy(client_etcp->peer_public_key, server_pub, 32);
    client_etcp->has_peer_key = 1;
    sc_set_peer_public_key(client_etcp->crypto_ctx, server_pub);
    
    // Server will get client's public key from INIT packet
    
    // Create connections managers
    struct ETCP_CONNECTIONS* server_conns = etcp_connections_init(server_etcp, "127.0.0.1", 0);
    struct ETCP_CONNECTIONS* client_conns = etcp_connections_init(client_etcp, "127.0.0.1", 0);
    if (!server_conns || !client_conns) {
        printf("ERROR: Failed to create connections\n");
        return 1;
    }
    
    // Create client link (client initiates connection)
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr("127.0.0.1");
    server_addr.sin_port = htons(TEST_PORT);
    
    struct ETCP_LINK* client_link = etcp_link_new(client_etcp, &client_conns->socket, client_conns,
                                                  (struct sockaddr*)&server_addr, sizeof(server_addr));
    if (!client_link) {
        printf("ERROR: Failed to create client link\n");
        return 1;
    }
    
    // Send INIT from client to server
    printf("Sending INIT from client...\n");
    if (etcp_link_send_init(client_link, 1500, 30) < 0) {
        printf("ERROR: Failed to send INIT\n");
        return 1;
    }
    
    // Server should receive INIT and create link
    usleep(100000); // 100ms
    
    struct sockaddr_in from_addr;
    socklen_t from_len = sizeof(from_addr);
    uint8_t buffer[2048];
    ssize_t received = recvfrom(server_fd, buffer, sizeof(buffer), MSG_DONTWAIT,
                               (struct sockaddr*)&from_addr, &from_len);
    
    if (received > 0) {
        printf("Server received %zd bytes\n", received);
        
        // Process packet on server - use packet pool
        struct packet_buffer* pkt = packet_pool_get(&server_etcp->packet_pool);
        if (pkt) {
            memcpy(packet_data(pkt), buffer, received);
            pkt->metadata.data_len = received;
            pkt->metadata.remote_addr = *(struct sockaddr_storage*)&from_addr;
            pkt->metadata.s = &server_conns->socket;
            
            if (etcp_input(pkt, &server_conns->socket, server_conns) == 0) {
                printf("Server processed INIT successfully\n");
                
                // Check that server got client's public key
                if (server_etcp->has_peer_key) {
                    printf("Server received client's public key\n");
                }
            }
            
            packet_pool_put(&server_etcp->packet_pool, pkt);
        }
    }
    
    // Test encrypted data transfer
    printf("\nTesting encrypted data transfer...\n");
    
    const char* test_data = TEST_DATA;
    if (etcp_link_send(client_etcp, client_link, (const uint8_t*)test_data, TEST_DATA_LEN) == 0) {
        printf("Client sent encrypted data\n");
    }
    
    usleep(100000);
    
    received = recvfrom(server_fd, buffer, sizeof(buffer), MSG_DONTWAIT,
                       (struct sockaddr*)&from_addr, &from_len);
    if (received > 0) {
        printf("Server received %zd encrypted bytes\n", received);
        // In real scenario, etcp_input would decrypt and process
    }
    
    // Print statistics
    printf("\n=== Statistics ===\n");
    size_t enc_err, dec_err, send_err, total_enc, total_dec;
    etcp_connections_get_crypto_stats(server_conns, &enc_err, &dec_err, &send_err, NULL, &total_enc, &total_dec);
    printf("Server: encrypted=%zu, decrypted=%zu, errors=%zu/%zu/%zu\n",
           total_enc, total_dec, enc_err, dec_err, send_err);
    
    etcp_connections_get_crypto_stats(client_conns, &enc_err, &dec_err, &send_err, NULL, &total_enc, &total_dec);
    printf("Client: encrypted=%zu, decrypted=%zu, errors=%zu/%zu/%zu\n",
           total_enc, total_dec, enc_err, dec_err, send_err);
    
    // Cleanup
    close(server_fd);
    close(client_fd);
    free(server_etcp->crypto_ctx);
    free(client_etcp->crypto_ctx);
    free(server_etcp);
    free(client_etcp);
    
    printf("\n=== Test completed ===\n");
    return 0;
}