/**
 * @file test_wireguard_simple.c
 * @brief WireGuard Protocol Tests (Simplified - minimal dependencies)
 */

#define X25519_SW_TEST 1

#include "se050_wireguard.h"
#include "se050_x25519_sw.h"
#include "se050_chacha20_poly1305.h"
#include "se050_blake2s.h"
#include "se050_hmac_blake2s.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

static int passed = 0;
static int failed = 0;

#define TEST_ASSERT(cond, msg) do { \
    if (cond) { \
        printf("[PASS] %s\n", msg); \
        passed++; \
    } else { \
        printf("[FAIL] %s\n", msg); \
        failed++; \
    } \
} while(0)

/* Test vectors */
static const uint8_t TEST_PRIVATE_KEY[32] = {
    0x77,0x07,0x6d,0x0a,0x73,0x18,0xa5,0x7d,
    0x3c,0x16,0xc1,0x72,0x51,0xb2,0x66,0x45,
    0xdf,0xbc,0x9f,0x03,0xc8,0xf2,0xbc,0x2b,
    0x4f,0x81,0xb0,0x9c,0x0c,0xcb,0x3e,0x9a
};

static const uint8_t TEST_PEER_PUBLIC_KEY[32] = {
    0x5d,0xab,0x08,0x7e,0x62,0x4a,0x8a,0x4b,
    0x79,0xe1,0x7f,0x8b,0x83,0x80,0x0e,0xe6,
    0x6f,0x31,0x97,0xb6,0x93,0x44,0x41,0x76,
    0x61,0x6c,0x39,0x5c,0xbe,0x7a,0x4b,0x82
};

/* Test: X25519 public key derivation */
static void test_x25519_derive(void)
{
    printf("\n--- Test X25519 Public Key Derivation ---\n");
    
    uint8_t public_key[32];
    int ret = se050_x25519_sw_derive_public_key(public_key, TEST_PRIVATE_KEY);
    
    TEST_ASSERT(ret == 0, "Public key derivation returns 0");
    
    /* Verify key is non-zero */
    uint8_t all_zero = 1;
    for (int i = 0; i < 32; i++) {
        if (public_key[i] != 0) all_zero = 0;
    }
    TEST_ASSERT(all_zero == 0, "Public key is non-zero");
}

/* Test: ChaCha20-Poly1305 AEAD */
static void test_chacha20_poly1305(void)
{
    printf("\n--- Test ChaCha20-Poly1305 AEAD ---\n");
    
    const uint8_t key[32] = {0};
    const uint8_t nonce[12] = {0};
    const uint8_t plaintext[] = "test";
    const uint8_t aad[] = "aad";
    
    uint8_t ciphertext[100];
    uint8_t tag[16];
    
    se050_chacha20_poly1305_ctx_t ctx;
    int ret = se050_chacha20_poly1305_init(&ctx, key);
    TEST_ASSERT(ret == 0, "Context initialization returns 0");
    
    ret = se050_chacha20_poly1305_encrypt(&ctx, nonce, plaintext, sizeof(plaintext)-1,
                                            aad, sizeof(aad)-1, ciphertext, tag);
    TEST_ASSERT(ret == 0, "Encryption returns 0");
    
    uint8_t decrypted[100];
    ret = se050_chacha20_poly1305_decrypt(&ctx, nonce, ciphertext, sizeof(plaintext)-1,
                                            aad, sizeof(aad)-1, tag, decrypted);
    TEST_ASSERT(ret == 0, "Decryption returns 0");
    TEST_ASSERT(decrypted[0] == 't' && decrypted[1] == 'e' && 
                decrypted[2] == 's' && decrypted[3] == 't', 
                "Decrypted content matches");
}

/* Test: HMAC-BLAKE2s */
static void test_hmac_blake2s(void)
{
    printf("\n--- Test HMAC-BLAKE2s ---\n");
    
    const uint8_t key[32] = {0};
    const uint8_t data[] = "test message";
    uint8_t mac[32];
    
    int ret = se050_hmac_blake2s(mac, key, 32, data, sizeof(data)-1);
    TEST_ASSERT(ret == 0, "HMAC computation returns 0");
    
    /* Verify MAC is non-zero */
    uint8_t all_zero = 1;
    for (int i = 0; i < 32; i++) {
        if (mac[i] != 0) all_zero = 0;
    }
    TEST_ASSERT(all_zero == 0, "MAC is non-zero");
}

/* Test: Session initialization */
static void test_session_init(void)
{
    printf("\n--- Test Session Initialization ---\n");
    
    se050_wireguard_session_t session;
    int ret = se050_wireguard_session_init(&session, TEST_PRIVATE_KEY, TEST_PEER_PUBLIC_KEY);
    
    TEST_ASSERT(ret == 0, "Session init returns 0");
    
    /* Verify public key was derived */
    uint8_t expected_public[32];
    se050_x25519_sw_derive_public_key(expected_public, TEST_PRIVATE_KEY);
    TEST_ASSERT(memcmp(session.public_key, expected_public, 32) == 0, 
                "Public key derived correctly");
    
    se050_wireguard_session_cleanup(&session);
}

/* Test: Key derivation */
static void test_key_derivation(void)
{
    printf("\n--- Test Key Derivation ---\n");
    
    se050_wireguard_session_t session;
    se050_wireguard_session_init(&session, TEST_PRIVATE_KEY, TEST_PEER_PUBLIC_KEY);
    
    uint8_t shared_secret[32] = {0};
    for (int i = 0; i < 32; i++) shared_secret[i] = i;
    
    int ret = se050_wireguard_derive_keys(&session, shared_secret);
    TEST_ASSERT(ret == 0, "Key derivation returns 0");
    TEST_ASSERT(session.handshake_complete == 1, "Handshake marked complete");
    
    uint8_t all_zero = 1;
    for (int i = 0; i < 32; i++) {
        if (session.sending_key[i] != 0 || session.receiving_key[i] != 0) {
            all_zero = 0;
            break;
        }
    }
    TEST_ASSERT(all_zero == 0, "Session keys are non-zero");
    
    se050_wireguard_session_cleanup(&session);
}

/* Test: Packet encryption/decryption */
static void test_encrypt_decrypt(void)
{
    printf("\n--- Test Encryption/Decryption ---\n");
    
    se050_wireguard_session_t session;
    se050_wireguard_session_init(&session, TEST_PRIVATE_KEY, TEST_PEER_PUBLIC_KEY);
    
    uint8_t shared_secret[32] = {0};
    for (int i = 0; i < 32; i++) shared_secret[i] = i;
    se050_wireguard_derive_keys(&session, shared_secret);
    
    const char *plaintext = "Hello, WireGuard!";
    size_t plaintext_len = strlen(plaintext);
    
    uint8_t encrypted[1024];
    size_t encrypted_len;
    
    int ret = se050_wireguard_encrypt_packet(&session, encrypted, &encrypted_len,
                                              (uint8_t*)plaintext, plaintext_len);
    TEST_ASSERT(ret == 0, "Encryption returns 0");
    TEST_ASSERT(encrypted_len == 16 + plaintext_len + 16, 
                "Encrypted length is correct");
    
    uint8_t decrypted[1024];
    size_t decrypted_len;
    
    ret = se050_wireguard_decrypt_packet(&session, decrypted, &decrypted_len,
                                          encrypted, encrypted_len);
    TEST_ASSERT(ret == 0, "Decryption returns 0");
    TEST_ASSERT(decrypted_len == plaintext_len, "Decrypted length matches");
    TEST_ASSERT(memcmp(decrypted, plaintext, plaintext_len) == 0, 
                "Decrypted content matches original");
    
    se050_wireguard_session_cleanup(&session);
}

/* Test: Replay detection */
static void test_replay_detection(void)
{
    printf("\n--- Test Replay Detection ---\n");
    
    se050_wireguard_session_t session;
    se050_wireguard_session_init(&session, TEST_PRIVATE_KEY, TEST_PEER_PUBLIC_KEY);
    
    uint8_t shared_secret[32] = {0};
    for (int i = 0; i < 32; i++) shared_secret[i] = i;
    se050_wireguard_derive_keys(&session, shared_secret);
    
    const char *plaintext = "Test message";
    uint8_t encrypted[1024];
    size_t encrypted_len;
    se050_wireguard_encrypt_packet(&session, encrypted, &encrypted_len,
                                    (uint8_t*)plaintext, strlen(plaintext));
    
    uint8_t decrypted[1024];
    size_t decrypted_len;
    int ret = se050_wireguard_decrypt_packet(&session, decrypted, &decrypted_len,
                                              encrypted, encrypted_len);
    TEST_ASSERT(ret == 0, "First decryption succeeds");
    
    ret = se050_wireguard_decrypt_packet(&session, decrypted, &decrypted_len,
                                          encrypted, encrypted_len);
    TEST_ASSERT(ret != 0, "Replay packet rejected");
    
    se050_wireguard_session_cleanup(&session);
}

/* Test: MAC computation */
static void test_mac_computation(void)
{
    printf("\n--- Test MAC Computation ---\n");
    
    se050_wireguard_session_t session;
    se050_wireguard_session_init(&session, TEST_PRIVATE_KEY, TEST_PEER_PUBLIC_KEY);
    
    uint8_t test_data[64] = {0};
    uint8_t mac1[16], mac2[16];
    
    int ret = se050_wireguard_compute_mac1(&session, test_data, sizeof(test_data), mac1);
    TEST_ASSERT(ret == 0, "MAC1 computation returns 0");
    
    ret = se050_wireguard_compute_mac2(&session, mac1, test_data, sizeof(test_data), mac2);
    TEST_ASSERT(ret == 0, "MAC2 computation returns 0");
    
    uint8_t mac1_zero = 1, mac2_zero = 1;
    for (int i = 0; i < 16; i++) {
        if (mac1[i] != 0) mac1_zero = 0;
        if (mac2[i] != 0) mac2_zero = 0;
    }
    TEST_ASSERT(mac1_zero == 0, "MAC1 is non-zero");
    TEST_ASSERT(mac2_zero == 0, "MAC2 is non-zero");
    
    se050_wireguard_session_cleanup(&session);
}

/* Test: Key generation */
static void test_key_generation(void)
{
    printf("\n--- Test Key Generation ---\n");
    
    uint8_t private_key[32], public_key[32];
    
    int ret = se050_wireguard_generate_keypair(private_key, public_key);
    TEST_ASSERT(ret == 0, "Key generation returns 0");
    
    uint8_t private_zero = 1, public_zero = 1;
    for (int i = 0; i < 32; i++) {
        if (private_key[i] != 0) private_zero = 0;
        if (public_key[i] != 0) public_zero = 0;
    }
    TEST_ASSERT(private_zero == 0, "Private key is non-zero");
    TEST_ASSERT(public_zero == 0, "Public key is non-zero");
    
    uint8_t expected_public[32];
    se050_x25519_sw_derive_public_key(expected_public, private_key);
    TEST_ASSERT(memcmp(public_key, expected_public, 32) == 0, 
                "Public key derived from private key");
}

/* Test: Session cleanup zeros memory */
static void test_session_cleanup(void)
{
    printf("\n--- Test Session Cleanup ---\n");
    
    se050_wireguard_session_t session;
    se050_wireguard_session_init(&session, TEST_PRIVATE_KEY, TEST_PEER_PUBLIC_KEY);
    
    uint8_t private_copy[32];
    memcpy(private_copy, session.private_key, 32);
    
    se050_wireguard_session_cleanup(&session);
    
    uint8_t all_zero = 1;
    for (int i = 0; i < 32; i++) {
        if (session.private_key[i] != 0) all_zero = 0;
    }
    TEST_ASSERT(all_zero == 1, "Session memory zeroed after cleanup");
}

/* Test: Invalid inputs */
static void test_invalid_inputs(void)
{
    printf("\n--- Test Invalid Inputs ---\n");
    
    se050_wireguard_session_t session;
    
    int ret = se050_wireguard_session_init(NULL, TEST_PRIVATE_KEY, TEST_PEER_PUBLIC_KEY);
    TEST_ASSERT(ret == -1, "NULL session rejected");
    
    ret = se050_wireguard_session_init(&session, NULL, TEST_PEER_PUBLIC_KEY);
    TEST_ASSERT(ret == -1, "NULL private key rejected");
    
    ret = se050_wireguard_session_init(&session, TEST_PRIVATE_KEY, NULL);
    TEST_ASSERT(ret == -1, "NULL peer public key rejected");
}

int main(void)
{
    printf("========================================\n");
    printf(" WireGuard Protocol Test Suite\n");
    printf("========================================\n");
    
    test_x25519_derive();
    test_chacha20_poly1305();
    test_hmac_blake2s();
    test_session_init();
    test_key_derivation();
    test_encrypt_decrypt();
    test_replay_detection();
    test_mac_computation();
    test_key_generation();
    test_session_cleanup();
    test_invalid_inputs();
    
    printf("\n========================================\n");
    printf(" Results: %d passed, %d failed\n", passed, failed);
    printf("========================================\n");
    
    return failed == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}