/* test_udp_secure.c - Unit test for UDP secure channel with async library */

#include "sc_lib.h"
#include "u_async.h"
#include <tinycrypt/ecc.h>
#include <tinycrypt/ecc_platform_specific.h>
#include <tinycrypt/constants.h>
#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <time.h>

#define TEST_ASSERT(cond, msg) \
    do { \
        if (!(cond)) { \
            printf("FAIL: %s (line %d)\n", msg, __LINE__); \
            exit(1); \
        } else { \
            printf("PASS: %s\n", msg); \
        } \
    } while(0)

#define PORT_A 12345
#define PORT_B 12346
#define LOCALHOST "127.0.0.1"

typedef struct {
    sc_context_t ctx;
    int sockfd;
    struct sockaddr_in addr;
    struct sockaddr_in peer_addr;
    uint8_t rx_buffer[1024];
    size_t rx_len;
    uint8_t tx_buffer[1024];
    size_t tx_len;
    int received;
    int sent;
} client_state_t;

static client_state_t client_a, client_b;
static volatile int test_complete = 0;
static volatile int test_failed = 0;

static void timeout_callback(void* arg) {
    (void)arg;
    printf("ERROR: Test timeout\n");
    test_failed = 1;
    exit(1);
}

static void client_a_read_callback(int fd, void* arg) {
    client_state_t* state = (client_state_t*)arg;
    struct sockaddr_in from;
    socklen_t fromlen = sizeof(from);
    ssize_t n = recvfrom(fd, state->rx_buffer, sizeof(state->rx_buffer), 0,
                         (struct sockaddr*)&from, &fromlen);
    if (n < 0) {
        perror("recvfrom");
        return;
    }
    printf("Client A received %zd bytes\n", n);
    // In this test, client A is not expecting data
    // but we can handle it if needed
}

static void client_b_read_callback(int fd, void* arg) {
    client_state_t* state = (client_state_t*)arg;
    struct sockaddr_in from;
    socklen_t fromlen = sizeof(from);
    ssize_t n = recvfrom(fd, state->rx_buffer, sizeof(state->rx_buffer), 0,
                         (struct sockaddr*)&from, &fromlen);
    if (n < 0) {
        perror("recvfrom");
        return;
    }
    printf("Client B received %zd bytes\n", n);
    
    // Expect encrypted message (ciphertext + tag + crc32)
    size_t ciphertext_len = n;
    if (ciphertext_len <= SC_MAX_OVERHEAD || ciphertext_len > sizeof(state->rx_buffer)) {
        printf("ERROR: Invalid received length\n");
        test_failed = 1;
        return;
    }
    
    uint8_t* ciphertext = state->rx_buffer;
    uint8_t decrypted[1024];
    size_t decrypted_len;
    
    const char* expected = "Hello from client A over UDP!";
    size_t expected_len = strlen(expected) + 1;
    sc_status_t status = sc_decrypt(&state->ctx, ciphertext, ciphertext_len, decrypted, &decrypted_len);
    TEST_ASSERT(status == SC_OK, "decryption of received message");
    TEST_ASSERT(decrypted_len == expected_len, "decrypted length matches original");
    
    TEST_ASSERT(memcmp(decrypted, expected, decrypted_len) == 0,
                "decrypted matches expected plaintext");
    
    printf("Client B successfully decrypted message: %s\n", (char*)decrypted);
    test_complete = 1;
    exit(0);
}

static int create_udp_socket(int port, struct sockaddr_in* addr) {
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    TEST_ASSERT(sock >= 0, "socket creation");
    
    int flags = fcntl(sock, F_GETFL, 0);
    TEST_ASSERT(flags >= 0, "get socket flags");
    TEST_ASSERT(fcntl(sock, F_SETFL, flags | O_NONBLOCK) == 0, "set non-blocking");
    
    memset(addr, 0, sizeof(*addr));
    addr->sin_family = AF_INET;
    addr->sin_port = htons(port);
    addr->sin_addr.s_addr = inet_addr(LOCALHOST);
    
    TEST_ASSERT(bind(sock, (struct sockaddr*)addr, sizeof(*addr)) == 0, "bind socket");
    return sock;
}

static void generate_keys(client_state_t* client) {
    memset(&client->ctx, 0, sizeof(client->ctx));
    sc_status_t status = sc_generate_keypair(&client->ctx);
    TEST_ASSERT(status == SC_OK, "key generation");
}

static void send_encrypted_message(client_state_t* src, client_state_t* dst) {
    const char* plaintext = "Hello from client A over UDP!";
    size_t plaintext_len = strlen(plaintext) + 1;
    size_t ciphertext_len;
    
    uint8_t ciphertext[plaintext_len + SC_MAX_OVERHEAD];
    
    sc_status_t status = sc_encrypt(&src->ctx, (const uint8_t*)plaintext, plaintext_len,
                                    ciphertext, &ciphertext_len);
    TEST_ASSERT(status == SC_OK, "encryption for sending");
    TEST_ASSERT(ciphertext_len == plaintext_len + SC_MAX_OVERHEAD, "ciphertext length includes overhead");
    
    // Send ciphertext (already includes tag + crc32)
    memcpy(src->tx_buffer, ciphertext, ciphertext_len);
    size_t packet_len = ciphertext_len;
    
    ssize_t sent = sendto(src->sockfd, src->tx_buffer, packet_len, 0,
                          (struct sockaddr*)&dst->addr, sizeof(dst->addr));
    TEST_ASSERT(sent == (ssize_t)packet_len, "sendto");
    printf("Client A sent encrypted message (%zd bytes)\n", sent);
}

int main(void) {
    printf("=== UDP Secure Channel Test with Async Library ===\n");
    
    // Initialize TinyCrypt RNG
    uECC_set_rng(&default_CSPRNG);
    
    // Initialize async library
    uasync_t* ua = uasync_create();
    if (!ua) {
        fprintf(stderr, "Failed to create uasync instance\n");
        return 1;
    }
    uasync_init_instance(ua);
    
    // Create UDP sockets
    client_a.sockfd = create_udp_socket(PORT_A, &client_a.addr);
    client_b.sockfd = create_udp_socket(PORT_B, &client_b.addr);
    
    // Set peer addresses for sending
    client_a.peer_addr = client_b.addr;
    client_b.peer_addr = client_a.addr;
    
    // Generate key pairs for both clients
    generate_keys(&client_a);
    generate_keys(&client_b);
    
    // Exchange public keys
    sc_status_t status = sc_set_peer_public_key(&client_a.ctx, client_b.ctx.public_key);
    TEST_ASSERT(status == SC_OK, "client A set peer key");
    status = sc_set_peer_public_key(&client_b.ctx, client_a.ctx.public_key);
    TEST_ASSERT(status == SC_OK, "client B set peer key");
    
    // Verify session keys match
    TEST_ASSERT(memcmp(client_a.ctx.session_key, client_b.ctx.session_key,
                       SC_SESSION_KEY_SIZE) == 0,
                "session keys match");
    
    // Register sockets with async library
    uasync_add_socket(ua, client_a.sockfd, client_a_read_callback, NULL, NULL, &client_a);
    uasync_add_socket(ua, client_b.sockfd, client_b_read_callback, NULL, NULL, &client_b);
    
    // Set timeout for test completion (2 seconds)
    uasync_set_timeout(ua, 20000, NULL, timeout_callback); // timebase 0.1 ms, 20000 = 2 sec
    
    // Send encrypted message from A to B
    send_encrypted_message(&client_a, &client_b);
    
    // Run async mainloop (will exit via callback or timeout)
    printf("Starting async mainloop...\n");
    uasync_mainloop(ua); // This will not return
    
    // Should not reach here
    return 1;
}