/*
 * test_chacha20_extended.c - Test ChaCha20 extended counter implementation
 * 
 * This test verifies that the extended counter properly handles:
 * 1. Counter overflow at 2^32 blocks (256GB boundary)
 * 2. Correct keystream generation across the overflow boundary
 * 3. No duplicate keystream blocks
 */

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

#define TEST_BLOCK_SIZE 64
#define BLOCKS_NEAR_OVERFLOW 10  // Test blocks around overflow point

// Test helper: Compare two blocks for equality
int blocks_equal(const uint8_t* block1, const uint8_t* block2, size_t len) {
    return memcmp(block1, block2, len) == 0;
}

// Test 1: Verify extended counter handles overflow correctly
int test_counter_overflow() {
    printf("Test 1: Counter overflow handling\n");
    printf("  Testing counter transition from 0xFFFFFFFF to 0x00000000...\n");
    
    uint8_t key[32];
    uint8_t nonce[8];
    uint8_t input[TEST_BLOCK_SIZE];
    uint8_t output1[TEST_BLOCK_SIZE];
    uint8_t output2[TEST_BLOCK_SIZE];
    uint8_t output3[TEST_BLOCK_SIZE];
    
    // Initialize test data
    memset(key, 0xAA, 32);
    memset(nonce, 0xBB, 8);
    memset(input, 0, TEST_BLOCK_SIZE);
    
    // Test at counter_low = 0xFFFFFFFE, counter_high = 0
    uint32_t counter_low = 0xFFFFFFFE;
    uint32_t counter_high = 0;
    
    printf("  Block at counter_low=0xFFFFFFFE, counter_high=0...\n");
    if (chacha20_encrypt_extended(key, counter_low, counter_high, nonce, 
                                   input, output1, TEST_BLOCK_SIZE) != 0) {
        printf("  ❌ FAILED: Error at counter_low=0xFFFFFFFE\n");
        return 1;
    }
    
    // Test at counter_low = 0xFFFFFFFF, counter_high = 0
    counter_low = 0xFFFFFFFF;
    printf("  Block at counter_low=0xFFFFFFFF, counter_high=0...\n");
    if (chacha20_encrypt_extended(key, counter_low, counter_high, nonce, 
                                   input, output2, TEST_BLOCK_SIZE) != 0) {
        printf("  ❌ FAILED: Error at counter_low=0xFFFFFFFF\n");
        return 1;
    }
    
    // Test at counter_low = 0x00000000, counter_high = 1 (after overflow)
    counter_low = 0x00000000;
    counter_high = 1;
    printf("  Block at counter_low=0x00000000, counter_high=1...\n");
    if (chacha20_encrypt_extended(key, counter_low, counter_high, nonce, 
                                   input, output3, TEST_BLOCK_SIZE) != 0) {
        printf("  ❌ FAILED: Error at counter_low=0x00000000, counter_high=1\n");
        return 1;
    }
    
    // Verify all three blocks are different (no keystream reuse)
    if (blocks_equal(output1, output2, TEST_BLOCK_SIZE)) {
        printf("  ❌ FAILED: Blocks at 0xFFFFFFFE and 0xFFFFFFFF are identical!\n");
        return 1;
    }
    
    if (blocks_equal(output2, output3, TEST_BLOCK_SIZE)) {
        printf("  ❌ FAILED: Blocks at 0xFFFFFFFF,0 and 0x00000000,1 are identical!\n");
        return 1;
    }
    
    if (blocks_equal(output1, output3, TEST_BLOCK_SIZE)) {
        printf("  ❌ FAILED: Blocks at 0xFFFFFFFE,0 and 0x00000000,1 are identical!\n");
        return 1;
    }
    
    printf("  ✓ All blocks are unique across overflow boundary\n");
    printf("  ✓ PASSED\n\n");
    return 0;
}

// Test 2: Simulate processing data that crosses 256GB boundary
int test_large_file_simulation() {
    printf("Test 2: Large file simulation (256GB+ boundary)\n");
    printf("  Simulating processing across 256GB boundary...\n");
    
    uint8_t key[32];
    uint8_t nonce[8];
    uint8_t input[1024];
    uint8_t output[1024];
    
    // Initialize test data
    memset(key, 0x55, 32);
    memset(nonce, 0x77, 8);
    for (int i = 0; i < 1024; i++) {
        input[i] = i & 0xFF;
    }
    
    // Simulate being at 256GB - 512 bytes (just before overflow)
    // 256GB = 2^32 blocks * 64 bytes = 274,877,906,944 bytes
    // Block number at 256GB - 512 bytes = 2^32 - 8 blocks
    uint32_t counter_low = 0xFFFFFFF8;  // 2^32 - 8
    uint32_t counter_high = 0;
    
    printf("  Processing 1KB starting at block 0xFFFFFFF8 (256GB - 512 bytes)...\n");
    
    // This should cross the overflow boundary
    int result = chacha20_encrypt_extended(key, counter_low, counter_high, nonce, 
                                           input, output, 1024);
    
    if (result != 0) {
        printf("  ❌ FAILED: Error processing data across 256GB boundary\n");
        return 1;
    }
    
    printf("  ✓ Successfully processed data across 256GB boundary\n");
    printf("  ✓ PASSED\n\n");
    return 0;
}

// Test 3: Verify extended vs standard ChaCha20 compatibility
int test_compatibility() {
    printf("Test 3: Compatibility with standard ChaCha20\n");
    printf("  Verifying extended mode matches standard mode when counter_high=0...\n");
    
    uint8_t key[32];
    uint8_t nonce_standard[12];
    uint8_t nonce_reduced[8];
    uint8_t input[TEST_BLOCK_SIZE];
    uint8_t output_standard[TEST_BLOCK_SIZE];
    uint8_t output_extended[TEST_BLOCK_SIZE];
    
    // Initialize test data
    memset(key, 0x33, 32);
    memset(nonce_standard, 0x44, 12);
    memcpy(nonce_reduced, nonce_standard + 4, 8);  // Extract last 8 bytes
    memset(input, 0, TEST_BLOCK_SIZE);
    
    uint32_t counter = 42;
    
    // Standard ChaCha20
    if (chacha20_encrypt(key, counter, nonce_standard, input, 
                        output_standard, TEST_BLOCK_SIZE) != 0) {
        printf("  ❌ FAILED: Standard ChaCha20 error\n");
        return 1;
    }
    
    // Extended ChaCha20 with counter_high=0 and matching nonce
    // The extended version builds nonce as [counter_high][nonce_reduced]
    // So we need to ensure the first 4 bytes of nonce_standard are 0
    uint8_t nonce_standard_zero[12] = {0};
    memcpy(nonce_standard_zero + 4, nonce_reduced, 8);
    
    if (chacha20_encrypt(key, counter, nonce_standard_zero, input, 
                        output_standard, TEST_BLOCK_SIZE) != 0) {
        printf("  ❌ FAILED: Standard ChaCha20 error\n");
        return 1;
    }
    
    if (chacha20_encrypt_extended(key, counter, 0, nonce_reduced, input, 
                                  output_extended, TEST_BLOCK_SIZE) != 0) {
        printf("  ❌ FAILED: Extended ChaCha20 error\n");
        return 1;
    }
    
    // Compare outputs
    if (!blocks_equal(output_standard, output_extended, TEST_BLOCK_SIZE)) {
        printf("  ❌ FAILED: Extended mode output differs from standard mode\n");
        printf("  First 16 bytes of standard: ");
        for (int i = 0; i < 16; i++) printf("%02x ", output_standard[i]);
        printf("\n  First 16 bytes of extended: ");
        for (int i = 0; i < 16; i++) printf("%02x ", output_extended[i]);
        printf("\n");
        return 1;
    }
    
    printf("  ✓ Extended mode matches standard mode when counter_high=0\n");
    printf("  ✓ PASSED\n\n");
    return 0;
}

// Test 4: Stress test - verify no errors at extreme counter values
int test_extreme_values() {
    printf("Test 4: Extreme counter values\n");
    printf("  Testing at various extreme counter positions...\n");
    
    uint8_t key[32];
    uint8_t nonce[8];
    uint8_t input[TEST_BLOCK_SIZE];
    uint8_t output[TEST_BLOCK_SIZE];
    
    memset(key, 0x99, 32);
    memset(nonce, 0x66, 8);
    memset(input, 0, TEST_BLOCK_SIZE);
    
    // Test various extreme positions
    struct {
        uint32_t counter_low;
        uint32_t counter_high;
        const char* description;
    } test_cases[] = {
        {0x00000000, 0, "Start of first 256GB segment"},
        {0xFFFFFFFF, 0, "End of first 256GB segment"},
        {0x00000000, 1, "Start of second 256GB segment"},
        {0xFFFFFFFF, 1, "End of second 256GB segment"},
        {0x00000000, 0xFFFF, "Start of segment 65535"},
        {0xFFFFFFFF, 0xFFFF, "End of segment 65535"},
    };
    
    for (size_t i = 0; i < sizeof(test_cases) / sizeof(test_cases[0]); i++) {
        printf("  Testing: %s (0x%08X, 0x%08X)...\n", 
               test_cases[i].description,
               test_cases[i].counter_low,
               test_cases[i].counter_high);
        
        if (chacha20_encrypt_extended(key, test_cases[i].counter_low, 
                                      test_cases[i].counter_high, nonce,
                                      input, output, TEST_BLOCK_SIZE) != 0) {
            printf("  ❌ FAILED at %s\n", test_cases[i].description);
            return 1;
        }
    }
    
    printf("  ✓ All extreme values handled correctly\n");
    printf("  ✓ PASSED\n\n");
    return 0;
}

int main() {
    printf("=================================================================\n");
    printf("ChaCha20 Extended Counter Test Suite\n");
    printf("=================================================================\n\n");
    
    int failures = 0;
    
    failures += test_counter_overflow();
    failures += test_large_file_simulation();
    failures += test_compatibility();
    failures += test_extreme_values();
    
    printf("=================================================================\n");
    if (failures == 0) {
        printf("✓ ALL TESTS PASSED\n");
        printf("=================================================================\n");
        printf("\nThe extended counter implementation is working correctly.\n");
        printf("It can now handle pads larger than 256GB without overflow errors.\n");
        return 0;
    } else {
        printf("❌ %d TEST(S) FAILED\n", failures);
        printf("=================================================================\n");
        return 1;
    }
}