/**
 * @file test_csprng.c
 * @brief CSPRNG Tests
 */

#include "se050_wireguard.h"
#include "se050_chacha20_poly1305.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)

/* Mock SE050 RNG for testing */
static int mock_se050_rng(uint8_t *out, size_t len, void *ctx)
{
    (void)ctx;
    /* Generate deterministic "seed" for testing */
    for (size_t i = 0; i < len; i++) {
        out[i] = (uint8_t)(i + 0x42);
    }
    return 0;
}

/* Test: CSPRNG initialization */
static void test_csprng_init(void)
{
    printf("\n--- Test CSPRNG Initialization ---\n");
    
    int ret = se050_csprng_init(mock_se050_rng, NULL);
    TEST_ASSERT(ret == 0, "CSPRNG init returns 0");
    
    se050_csprng_cleanup();
}

/* Test: CSPRNG generates non-zero data */
static void test_csprng_output(void)
{
    printf("\n--- Test CSPRNG Output ---\n");
    
    se050_csprng_init(mock_se050_rng, NULL);
    
    uint8_t rand1[32];
    int ret = se050_csprng_random(rand1, 32);
    TEST_ASSERT(ret == 0, "Random generation returns 0");
    
    uint8_t all_zero = 1;
    for (int i = 0; i < 32; i++) {
        if (rand1[i] != 0) all_zero = 0;
    }
    TEST_ASSERT(all_zero == 0, "Random data is non-zero");
    
    se050_csprng_cleanup();
}

/* Test: CSPRNG generates different values */
static void test_csprng_uniqueness(void)
{
    printf("\n--- Test CSPRNG Uniqueness ---\n");
    
    se050_csprng_init(mock_se050_rng, NULL);
    
    uint8_t rand1[32], rand2[32];
    se050_csprng_random(rand1, 32);
    se050_csprng_random(rand2, 32);
    
    TEST_ASSERT(memcmp(rand1, rand2, 32) != 0, "Successive calls produce different values");
    
    se050_csprng_cleanup();
}

/* Test: CSPRNG keypair generation */
static void test_csprng_keypair(void)
{
    printf("\n--- Test CSPRNG Keypair ---\n");
    
    se050_csprng_init(mock_se050_rng, NULL);
    
    uint8_t priv[32], pub[32];
    int ret = se050_wireguard_generate_keypair_csprng(priv, pub);
    TEST_ASSERT(ret == 0, "Keypair generation returns 0");
    
    uint8_t priv_zero = 1, pub_zero = 1;
    for (int i = 0; i < 32; i++) {
        if (priv[i] != 0) priv_zero = 0;
        if (pub[i] != 0) pub_zero = 0;
    }
    TEST_ASSERT(priv_zero == 0, "Private key is non-zero");
    TEST_ASSERT(pub_zero == 0, "Public key is non-zero");
    
    se050_csprng_cleanup();
}

/* Test: CSPRNG cleanup zeros memory */
static void test_csprng_cleanup(void)
{
    printf("\n--- Test CSPRNG Cleanup ---\n");
    
    se050_csprng_init(mock_se050_rng, NULL);
    se050_csprng_cleanup();
    
    /* Try to generate after cleanup - should fail */
    uint8_t rand[32];
    int ret = se050_csprng_random(rand, 32);
    TEST_ASSERT(ret != 0, "Generation fails after cleanup");
}

/* Test: Multiple keypairs from same seed */
static void test_multiple_keypairs(void)
{
    printf("\n--- Test Multiple Keypairs ---\n");
    
    se050_csprng_init(mock_se050_rng, NULL);
    
    uint8_t priv1[32], pub1[32];
    uint8_t priv2[32], pub2[32];
    
    se050_wireguard_generate_keypair_csprng(priv1, pub1);
    se050_wireguard_generate_keypair_csprng(priv2, pub2);
    
    TEST_ASSERT(memcmp(priv1, priv2, 32) != 0, "Multiple private keys are different");
    TEST_ASSERT(memcmp(pub1, pub2, 32) != 0, "Multiple public keys are different");
    
    se050_csprng_cleanup();
}

int main(void)
{
    printf("========================================\n");
    printf(" CSPRNG Test Suite\n");
    printf("========================================\n");
    
    test_csprng_init();
    test_csprng_output();
    test_csprng_uniqueness();
    test_csprng_keypair();
    test_csprng_cleanup();
    test_multiple_keypairs();
    
    printf("\n========================================\n");
    printf(" Results: %d passed, %d failed\n", passed, failed);
    printf("========================================\n");
    
    return failed == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}