// test_etcp_link_crypto_working.c - Working test with ETCP_LINK encryption via UDP
#include "etcp.h"
#include "etcp_connections.h"
#include "secure_channel.h"
#include "packet_pool.h"
#include "crc32.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <fcntl.h>
#include <errno.h>

#define TEST_PORT_SERVER 22345
#define TEST_PORT_CLIENT 22346
#define TEST_DATA "Hello, encrypted ETCP_LINK!"

// Fixed test keys (hex strings)
static const char* SERVER_PRIV = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
static const char* SERVER_PUB  = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
static const char* CLIENT_PRIV = "fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210";
static const char* CLIENT_PUB  = "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef";

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;
}

// Create UDP socket
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) {
        perror("bind");
        close(fd);
        return -1;
    }
    
    printf("[SOCK] Created socket on port %d (fd=%d)\n", port, fd);
    return fd;
}

// Helper structure for test instance
typedef struct {
    struct ETCP_CONN* etcp;
    struct ETCP_CONNECTIONS* conns;
    int socket_fd;
    uint64_t node_id;
    const char* name;
} test_inst_t;

// Initialize crypto context
static int init_crypto(struct ETCP_CONN* etcp, const char* priv_hex, const char* pub_hex) {
    etcp->crypto_ctx = calloc(1, sizeof(sc_context_t));
    if (!etcp->crypto_ctx) return -1;
    
    uint8_t priv[32], pub[32];
    if (hex_to_bin(priv_hex, priv, 32) < 0 || hex_to_bin(pub_hex, pub, 32) < 0) {
        return -1;
    }
    
    memcpy(etcp->crypto_ctx->private_key, priv, 32);
    memcpy(etcp->crypto_ctx->public_key, pub, 32);
    etcp->crypto_ctx->initialized = 1;
    return 0;
}

// Create test instance
test_inst_t* create_test_inst(const char* name, uint64_t node_id,
                               const char* priv_hex, const char* pub_hex,
                               const char* peer_pub_hex,
                               int bind_port) {
    test_inst_t* inst = calloc(1, sizeof(test_inst_t));
    if (!inst) return NULL;
    
    inst->name = name;
    inst->node_id = node_id;
    
    // Create ETCP_CONN
    inst->etcp = calloc(1, sizeof(struct ETCP_CONN));
    if (!inst->etcp) goto error;
    
    inst->etcp->node_id = node_id;
    inst->etcp->state = 1;
    
    // Init crypto
    if (init_crypto(inst->etcp, priv_hex, pub_hex) < 0) {
        fprintf(stderr, "%s: Failed to init crypto\n", name);
        goto error;
    }
    
    // Set peer key if provided
    if (peer_pub_hex) {
        uint8_t peer_pub[32];
        if (hex_to_bin(peer_pub_hex, peer_pub, 32) == 0) {
            memcpy(inst->etcp->peer_public_key, peer_pub, 32);
            inst->etcp->has_peer_key = 1;
            sc_set_peer_public_key(inst->etcp->crypto_ctx, peer_pub);
        }
    }
    
    // Init packet pool
    packet_pool_init(&inst->etcp->packet_pool);
    
    // Create UDP socket
    inst->socket_fd = create_udp_socket(bind_port);
    if (inst->socket_fd < 0) goto error;
    
    // Create connections
    inst->conns = etcp_connections_init(inst->etcp, "127.0.0.1", bind_port);
    if (!inst->conns) goto error;
    
    printf("[INST] %s: Created on port %d (node_id: %llx)\n", 
           name, bind_port, (unsigned long long)node_id);
    return inst;
    
error:
    if (inst) {
        if (inst->etcp) {
            free(inst->etcp->crypto_ctx);
            packet_pool_destroy(&inst->etcp->packet_pool);
        }
        free(inst->etcp);
        if (inst->socket_fd >= 0) close(inst->socket_fd);
    }
    free(inst);
    return NULL;
}

// Destroy test instance
void destroy_test_inst(test_inst_t* inst) {
    if (!inst) return;
    
    etcp_destroy(inst->etcp);
    if (inst->socket_fd >= 0) close(inst->socket_fd);
    free(inst);
}

// Create link between instances
struct ETCP_LINK* create_test_link(test_inst_t* local, test_inst_t* remote,
                                    int remote_port) {
    struct sockaddr_in remote_addr;
    memset(&remote_addr, 0, sizeof(remote_addr));
    remote_addr.sin_family = AF_INET;
    remote_addr.sin_addr.s_addr = inet_addr("127.0.0.1");
    remote_addr.sin_port = htons(remote_port);
    
    struct ETCP_LINK* link = etcp_link_new(local->etcp, &local->conns->socket, local->conns,
                                           (struct sockaddr*)&remote_addr, sizeof(remote_addr));
    if (!link) {
        fprintf(stderr, "Failed to create link\n");
        return NULL;
    }
    
    if (etcp_link_add_to_connections(local->conns, link) < 0) {
        fprintf(stderr, "Failed to add link to connections\n");
        etcp_link_close(link);
        return NULL;
    }
    
    return link;
}

// Send encrypted data and print stats
static int test_send_encrypted(struct ETCP_LINK* link, const uint8_t* data, size_t len) {
    printf("[SEND] Sending %zu bytes...\n", len);
    
    size_t before_enc = link->conns->total_encrypted;
    int ret = etcp_link_send(link->etcp, link, data, len);
    size_t after_enc = link->conns->total_encrypted;
    
    if (ret < 0) {
        fprintf(stderr, "[SEND] Failed (enc_errors=%zu)\n", link->conns->encrypt_errors);
        return -1;
    }
    
    if (after_enc > before_enc) {
        printf("[SEND] ✅ Sent and encrypted (%zu bytes)\n", len);
    } else {
        printf("[SEND] ⚠️  Sent without encryption (session not ready)\n");
    }
    
    return 0;
}

// Receive and decrypt
static int test_receive_decrypt(test_inst_t* inst, uint8_t* out_data, size_t* out_len) {
    struct sockaddr_in from_addr;
    socklen_t from_len = sizeof(from_addr);
    uint8_t recv_buffer[2048];
    
    ssize_t received = recvfrom(inst->socket_fd, recv_buffer, sizeof(recv_buffer), 0,
                               (struct sockaddr*)&from_addr, &from_len);
    if (received < 0) {
        if (errno == EAGAIN || errno == EWOULDBLOCK) {
            return 0; // No data
        }
        perror("[RECV] recvfrom");
        return -1;
    }
    
    printf("[RECV] Received %zd bytes\n", received);
    
    // Check if encrypted (simple heuristic: encrypted starts with cmd byte then aes output)
    if (received > 20 && recv_buffer[0] == 0x02) { // cmd byte
        printf("[RECV] Looks like encrypted packet\n");
    }
    
    // For now, just copy raw bytes (later will decrypt)
    memcpy(out_data, recv_buffer, received < 2048 ? received : 2047);
    *out_len = received;
    
    return 1; // Got data
}

int main(void) {
    printf("=== ETCP Link Encryption Test ===\n\n");
    
    // Create instances
    test_inst_t* server = create_test_inst("SERVER", 0x1111111111111111ULL,
                                            SERVER_PRIV, SERVER_PUB, NULL,
                                            TEST_PORT_SERVER);
    test_inst_t* client = create_test_inst("CLIENT", 0x2222222222222222ULL,
                                            CLIENT_PRIV, CLIENT_PUB, SERVER_PUB,
                                            TEST_PORT_CLIENT);
    
    if (!server || !client) {
        fprintf(stderr, "Failed to create instances\n");
        return 1;
    }
    
    // Create link from client to server
    struct ETCP_LINK* client_link = create_test_link(client, server, TEST_PORT_SERVER);
    if (!client_link) return 1;
    
    // Send test data
    printf("\n[TEST] Sending encrypted data...\n");
    if (test_send_encrypted(client_link, (const uint8_t*)TEST_DATA, strlen(TEST_DATA)) < 0) {
        fprintf(stderr, "Send failed\n");
        return 1;
    }
    
    // Server receives
    usleep(50000);
    uint8_t recv_buffer[2048];
    size_t recv_len = 0;
    
    int ret = test_receive_decrypt(server, recv_buffer, &recv_len);
    if (ret > 0) {
        printf("[RECV] Server got %zu bytes\n", recv_len);
        printf("       First 20 bytes: ");
        for (int i = 0; i < 20 && i < recv_len; i++) printf("%02x ", recv_buffer[i]);
        printf("\n");
    }
    
    // Check stats
    printf("\n=== Statistics ===\n");
    size_t enc_err, dec_err, send_err, recv_err, total_enc, total_dec;
    etcp_connections_get_crypto_stats(server->conns, &enc_err, &dec_err, &send_err, &recv_err, &total_enc, &total_dec);
    printf("SERVER: enc=%zu, dec=%zu, err=(enc:%zu,dec:%zu,send:%zu,recv:%zu)\n",
           total_enc, total_dec, enc_err, dec_err, send_err, recv_err);
    
    etcp_connections_get_crypto_stats(client->conns, &enc_err, &dec_err, &send_err, &recv_err, &total_enc, &total_dec);
    printf("CLIENT: enc=%zu, dec=%zu, err=(enc:%zu,dec:%zu,send:%zu,recv:%zu)\n",
           total_enc, total_dec, enc_err, dec_err, send_err, recv_err);
    
    if (total_enc > 0) {
        printf("\n✅ SUCCESS: Data was encrypted and sent!\n");
    } else {
        printf("\n❌ FAIL: No encryption happened\n");
    }
    
    // Cleanup
    destroy_test_inst(server);
    destroy_test_inst(client);
    
    printf("\n=== Test completed ===\n");
    return 0;
}