/*
 * ChaCha20 Test Suite - RFC 8439 Reference Test Vectors
 * 
 * This test suite validates our ChaCha20 implementation against the official
 * test vectors from RFC 8439 "ChaCha20 and Poly1305 for IETF Protocols".
 */

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include "../nostr_core/nostr_chacha20.h"

// Helper function to convert hex string to bytes
static int hex_to_bytes(const char* hex, uint8_t* bytes, size_t len) {
    for (size_t i = 0; i < len; i++) {
        if (sscanf(hex + 2*i, "%2hhx", &bytes[i]) != 1) {
            return -1;
        }
    }
    return 0;
}

// Helper function to convert bytes to hex string for display
static void bytes_to_hex(const uint8_t* bytes, size_t len, char* hex) {
    for (size_t i = 0; i < len; i++) {
        sprintf(hex + 2*i, "%02x", bytes[i]);
    }
    hex[2*len] = '\0';
}

// Helper function to compare byte arrays
static int bytes_equal(const uint8_t* a, const uint8_t* b, size_t len) {
    return memcmp(a, b, len) == 0;
}

// Test 1: ChaCha Quarter Round (RFC 8439 Section 2.1.1)
static int test_quarter_round() {
    printf("=== Test 1: ChaCha Quarter Round ===\n");
    
    uint32_t state[16] = {0};
    state[0] = 0x11111111;
    state[1] = 0x01020304;
    state[2] = 0x9b8d6f43;
    state[3] = 0x01234567;
    
    printf("Input:  a=0x%08x, b=0x%08x, c=0x%08x, d=0x%08x\n",
           state[0], state[1], state[2], state[3]);
    
    chacha20_quarter_round(state, 0, 1, 2, 3);
    
    printf("Output: a=0x%08x, b=0x%08x, c=0x%08x, d=0x%08x\n",
           state[0], state[1], state[2], state[3]);
    
    // Expected values from RFC 8439
    uint32_t expected[4] = {0xea2a92f4, 0xcb1cf8ce, 0x4581472e, 0x5881c4bb};
    
    if (state[0] == expected[0] && state[1] == expected[1] && 
        state[2] == expected[2] && state[3] == expected[3]) {
        printf("✅ Quarter round test PASSED\n\n");
        return 1;
    } else {
        printf("❌ Quarter round test FAILED\n");
        printf("Expected: a=0x%08x, b=0x%08x, c=0x%08x, d=0x%08x\n\n",
               expected[0], expected[1], expected[2], expected[3]);
        return 0;
    }
}

// Test 2: ChaCha20 Block Function - RFC 8439 Appendix A.1 Test Vectors
static int test_chacha20_block_1() {
    printf("=== Test 2: ChaCha20 Block Function - RFC 8439 Test Vectors ===\n");
    
    uint8_t key[32] = {0};
    uint8_t nonce[12] = {0};
    uint8_t output0[64], output1[64];
    
    printf("Key: all zeros\n");
    printf("Nonce: all zeros\n");
    
    // Test counter = 0 (RFC 8439 Appendix A.1 Test Vector #1)
    printf("\nTesting counter = 0:\n");
    chacha20_block(key, 0, nonce, output0);
    char hex_output0[129];
    bytes_to_hex(output0, 64, hex_output0);
    printf("Counter=0 output: %s\n", hex_output0);
    
    // Test counter = 1 (RFC 8439 Appendix A.1 Test Vector #2)
    printf("\nTesting counter = 1:\n");
    chacha20_block(key, 1, nonce, output1);
    char hex_output1[129];
    bytes_to_hex(output1, 64, hex_output1);
    printf("Counter=1 output: %s\n", hex_output1);
    
    // Expected for counter=0 from RFC 8439 Appendix A.1 Test Vector #1
    const char* expected_counter0_hex = 
        "76b8e0ada0f13d90405d6ae55386bd28"
        "bdd219b8a08ded1aa836efcc8b770dc7"
        "da41597c5157488d7724e03fb8d84a37"
        "6a43b8f41518a11cc387b669b2ee6586";
    
    // Expected for counter=1 from RFC 8439 Appendix A.1 Test Vector #2
    const char* expected_counter1_hex = 
        "9f07e7be5551387a98ba977c732d080d"
        "cb0f29a048e3656912c6533e32ee7aed"
        "29b721769ce64e43d57133b074d839d5"
        "31ed1f28510afb45ace10a1f4b794d6f";
    
    uint8_t expected0[64], expected1[64];
    hex_to_bytes(expected_counter0_hex, expected0, 64);
    hex_to_bytes(expected_counter1_hex, expected1, 64);
    
    printf("\nExpected counter=0: %s\n", expected_counter0_hex);
    printf("Expected counter=1: %s\n", expected_counter1_hex);
    
    int test0_pass = bytes_equal(output0, expected0, 64);
    int test1_pass = bytes_equal(output1, expected1, 64);
    
    if (test0_pass) printf("✅ Counter=0 test PASSED\n");
    else printf("❌ Counter=0 test FAILED\n");
    
    if (test1_pass) printf("✅ Counter=1 test PASSED\n");
    else printf("❌ Counter=1 test FAILED\n");
    
    printf("\n");
    return test0_pass && test1_pass; // Both tests must pass
}

// Test 3: ChaCha20 Block Function - Test Vector #2 (RFC 8439 Appendix A.1)
static int test_chacha20_block_2() {
    printf("=== Test 3: ChaCha20 Block Function - Different Key ===\n");
    
    // Key with last byte = 1, all-zero nonce, counter = 1
    uint8_t key[32] = {0};
    key[31] = 0x01; // Last byte = 1
    uint8_t nonce[12] = {0};
    uint32_t counter = 1;
    uint8_t output[64];
    
    printf("Key: all zeros except last byte = 0x01\n");
    printf("Nonce: all zeros\n");
    printf("Counter: %u\n", counter);
    
    int result = chacha20_block(key, counter, nonce, output);
    if (result != 0) {
        printf("❌ ChaCha20 block function failed\n\n");
        return 0;
    }
    
    // Expected output from RFC 8439 Appendix A.1 Test Vector #3
    const char* expected_hex = 
        "3aeb5224ecf849929b9d828db1ced4dd"
        "832025e8018b8160b82284f3c949aa5a"
        "8eca00bbb4a73bdad192b5c42f73f2fd"
        "4e273644c8b36125a64addeb006c13a0";
    
    uint8_t expected[64];
    hex_to_bytes(expected_hex, expected, 64);
    
    char hex_output[129];
    bytes_to_hex(output, 64, hex_output);
    printf("Output: %s\n", hex_output);
    
    if (bytes_equal(output, expected, 64)) {
        printf("✅ ChaCha20 block test #2 PASSED\n\n");
        return 1;
    } else {
        printf("❌ ChaCha20 block test #2 FAILED\n");
        printf("Expected: %s\n\n", expected_hex);
        return 0;
    }
}

// Test 4: ChaCha20 Encryption - "Sunscreen" Test (RFC 8439 Section 2.4.2)
static int test_chacha20_encryption() {
    printf("=== Test 4: ChaCha20 Encryption - Sunscreen Text ===\n");
    
    // Key and nonce from RFC 8439 Section 2.4.2
    const char* key_hex = 
        "000102030405060708090a0b0c0d0e0f"
        "101112131415161718191a1b1c1d1e1f";
    const char* nonce_hex = "000000000000004a00000000";
    
    uint8_t key[32];
    uint8_t nonce[12];
    hex_to_bytes(key_hex, key, 32);
    hex_to_bytes(nonce_hex, nonce, 12);
    
    // Test plaintext (first part of "Sunscreen" text)
    const char* plaintext = 
        "Ladies and Gentlemen of the class of '99: If I could offer you "
        "only one tip for the future, sunscreen would be it.";
    size_t plaintext_len = strlen(plaintext);
    
    printf("Plaintext: \"%.50s...\"\n", plaintext);
    printf("Length: %zu bytes\n", plaintext_len);
    
    uint8_t ciphertext[256];
    uint32_t counter = 1;
    
    int result = chacha20_encrypt(key, counter, nonce, 
                                 (const uint8_t*)plaintext, 
                                 ciphertext, plaintext_len);
    
    if (result != 0) {
        printf("❌ ChaCha20 encryption failed\n\n");
        return 0;
    }
    
    // Expected ciphertext (first 64 bytes) from RFC 8439
    const char* expected_hex = 
        "6e2e359a2568f98041ba0728dd0d6981"
        "e97e7aec1d4360c20a27afccfd9fae0b"
        "f91b65c5524733ab8f593dabcd62b357"
        "1639d624e65152ab8f530c359f0861d8";
    
    uint8_t expected[64];
    hex_to_bytes(expected_hex, expected, 64);
    
    char hex_output[129];
    bytes_to_hex(ciphertext, 64, hex_output);
    printf("Ciphertext (first 64 bytes): %s\n", hex_output);
    
    if (bytes_equal(ciphertext, expected, 64)) {
        printf("✅ ChaCha20 encryption test PASSED\n");
        
        // Test decryption (should get back original plaintext)
        uint8_t decrypted[256];
        result = chacha20_encrypt(key, counter, nonce, ciphertext, decrypted, plaintext_len);
        
        if (result == 0 && memcmp(plaintext, decrypted, plaintext_len) == 0) {
            printf("✅ ChaCha20 decryption test PASSED\n\n");
            return 1;
        } else {
            printf("❌ ChaCha20 decryption test FAILED\n\n");
            return 0;
        }
    } else {
        printf("❌ ChaCha20 encryption test FAILED\n");
        printf("Expected: %s\n\n", expected_hex);
        return 0;
    }
}

// Test 5: ChaCha20 Edge Cases and Additional Validation
static int test_chacha20_edge_cases() {
    printf("=== Test 5: ChaCha20 Edge Cases and Additional Validation ===\n");
    
    uint8_t key[32];
    uint8_t nonce[12];
    uint8_t input[64];
    uint8_t output[64];
    
    // Initialize test data
    memset(key, 0, 32);
    memset(nonce, 0, 12);
    memset(input, 0xAA, 64);  // Fill with test pattern
    
    printf("Testing various scenarios...\n");
    
    // Test 1: Counter = 0
    int result1 = chacha20_encrypt(key, 0, nonce, input, output, 64);
    printf("Counter=0 (64 bytes): %s\n", result1 == 0 ? "PASS" : "FAIL");
    
    // Test 2: Counter = 1
    int result2 = chacha20_encrypt(key, 1, nonce, input, output, 64);
    printf("Counter=1 (64 bytes): %s\n", result2 == 0 ? "PASS" : "FAIL");
    
    // Test 3: Empty data (0 bytes)
    int result3 = chacha20_encrypt(key, 0, nonce, input, output, 0);
    printf("Zero-length data: %s\n", result3 == 0 ? "PASS" : "FAIL");
    
    // Test 4: Single byte
    int result4 = chacha20_encrypt(key, 0, nonce, input, output, 1);
    printf("Single byte: %s\n", result4 == 0 ? "PASS" : "FAIL");
    
    // Test 5: Partial block (35 bytes)
    int result5 = chacha20_encrypt(key, 0, nonce, input, output, 35);
    printf("Partial block (35 bytes): %s\n", result5 == 0 ? "PASS" : "FAIL");
    
    // Test 6: Multi-block (128 bytes)
    uint8_t large_input[128];
    uint8_t large_output[128];
    memset(large_input, 0x55, 128);
    int result6 = chacha20_encrypt(key, 0, nonce, large_input, large_output, 128);
    printf("Multi-block (128 bytes): %s\n", result6 == 0 ? "PASS" : "FAIL");
    
    if (result1 == 0 && result2 == 0 && result3 == 0 && 
        result4 == 0 && result5 == 0 && result6 == 0) {
        printf("✅ All edge case tests PASSED\n\n");
        return 1;
    } else {
        printf("❌ Some edge case tests FAILED\n\n");
        return 0;
    }
}

// Main test function
int main() {
    printf("🧪 ChaCha20 Test Suite - RFC 8439 Reference Vectors\n");
    printf("=====================================================\n\n");
    
    int passed = 0;
    int total = 5;
    
    if (test_quarter_round()) passed++;
    if (test_chacha20_block_1()) passed++;
    if (test_chacha20_block_2()) passed++;
    if (test_chacha20_encryption()) passed++;
    if (test_chacha20_edge_cases()) passed++;
    
    printf("=== Test Summary ===\n");
    printf("Passed: %d/%d tests\n", passed, total);
    
    if (passed == total) {
        printf("🎉 ALL CHACHA20 TESTS PASSED! 🎉\n");
        printf("\nOur ChaCha20 implementation is RFC 8439 compliant and ready for production use.\n");
        printf("✅ Quarter round operations work correctly\n");
        printf("✅ Block function matches reference vectors\n");
        printf("✅ Encryption/decryption is bidirectional\n");
        printf("✅ Edge cases handled properly\n");
        return 0;
    } else {
        printf("❌ Some tests failed. ChaCha20 implementation needs fixes.\n");
        return 1;
    }
}