/*
 * NIP-44 Encryption/Decryption Test
 * 
 * Test suite for NIP-44 versioned encryption functionality
 * Uses known test vectors and cross-implementation compatibility tests
 */

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

// Test vectors for NIP-44 with proper key pairs
typedef struct {
    const char* name;
    const char* sender_private_key_hex;
    const char* recipient_private_key_hex;  // FIX: Need proper private key, not public key
    const char* plaintext;
    const char* expected_encrypted; // Optional - for known test vectors
} nip44_test_vector_t;

// Known decryption-only test vectors from nostr-tools (for cross-compatibility testing)
// Note: NIP-44 encryption is non-deterministic - ciphertext varies each time
// These vectors test our ability to decrypt known good ciphertext from reference implementations
static nip44_test_vector_t decryption_test_vectors[] = {
    {
        "Decryption test: single char 'a'",
        "0000000000000000000000000000000000000000000000000000000000000001",  // sec1
        "0000000000000000000000000000000000000000000000000000000000000002",  // sec2
        "a",
        "AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABee0G5VSK0/9YypIObAtDKfYEAjD35uVkHyB0F4DwrcNaCXlCWZKaArsGrY6M9wnuTMxWfp1RTN9Xga8no+kF5Vsb"
    },
    {
        "Decryption test: emoji",
        "0000000000000000000000000000000000000000000000000000000000000002",  // sec1
        "0000000000000000000000000000000000000000000000000000000000000001",  // sec2
        "🍕🫃",
        "AvAAAAAAAAAAAAAAAAAAAPAAAAAAAAAAAAAAAAAAAAAPSKSK6is9ngkX2+cSq85Th16oRTISAOfhStnixqZziKMDvB0QQzgFZdjLTPicCJaV8nDITO+QfaQ61+KbWQIOO2Yj"
    },
    {
        "Decryption test: wide unicode",
        "5c0c523f52a5b6fad39ed2403092df8cebc36318b39383bca6c00808626fab3a",  // sec1
        "4b22aa260e4acb7021e32f38a6cdf4b673c6a277755bfce287e370c924dc936d",  // sec2
        "表ポあA鷗ŒéＢ逍Üߪąñ丂㐀𠀀",
        "ArY1I2xC2yDwIbuNHN/1ynXdGgzHLqdCrXUPMwELJPc7s7JqlCMJBAIIjfkpHReBPXeoMCyuClwgbT419jUWU1PwaNl4FEQYKCDKVJz+97Mp3K+Q2YGa77B6gpxB/lr1QgoqpDf7wDVrDmOqGoiPjWDqy8KzLueKDcm9BVP8xeTJIxs="
    }
};

// Round-trip test vectors with proper key pairs
static nip44_test_vector_t test_vectors[] = {
    {
        "Basic short message",
        "91ba716fa9e7ea2fcbad360cf4f8e0d312f73984da63d90f524ad61a6a1e7dbe",  // Working keys from simple_nip44_test
        "96f6fa197aa07477ab88f6981118466ae3a982faab8ad5db9d5426870c73d220",
        "Hello, NIP-44!",
        NULL
    },
    {
        "Unicode message",
        "1111111111111111111111111111111111111111111111111111111111111111",
        "2222222222222222222222222222222222222222222222222222222222222222",
        "Hello 🌍 World! 🚀",
        NULL
    },
    {
        "Empty message",
        "3333333333333333333333333333333333333333333333333333333333333333",
        "4444444444444444444444444444444444444444444444444444444444444444",
        "",
        NULL
    }
};

static int hex_to_bytes(const char* hex, unsigned char* bytes, size_t len) {
    if (strlen(hex) != len * 2) return -1;
    
    for (size_t i = 0; i < len; i++) {
        if (sscanf(hex + i * 2, "%2hhx", &bytes[i]) != 1) {
            return -1;
        }
    }
    return 0;
}

static void bytes_to_hex(const unsigned char* bytes, size_t len, char* hex) {
    for (size_t i = 0; i < len; i++) {
        sprintf(hex + i * 2, "%02x", bytes[i]);
    }
    hex[len * 2] = '\0';
}

static int test_nip44_round_trip(const nip44_test_vector_t* tv) {
    printf("Test: %s\n", tv->name);
    
    // Parse keys - both private keys
    unsigned char sender_private_key[32];
    unsigned char recipient_private_key[32];
    
    if (hex_to_bytes(tv->sender_private_key_hex, sender_private_key, 32) != 0) {
        printf("  FAIL: Failed to parse sender private key\n");
        return -1;
    }
    
    if (hex_to_bytes(tv->recipient_private_key_hex, recipient_private_key, 32) != 0) {
        printf("  FAIL: Failed to parse recipient private key\n");
        return -1;
    }
    
    // Generate the public keys from the private keys
    unsigned char sender_public_key[32];
    unsigned char recipient_public_key[32];
    
    if (nostr_ec_public_key_from_private_key(sender_private_key, sender_public_key) != 0) {
        printf("  FAIL: Failed to derive sender public key\n");
        return -1;
    }
    
    if (nostr_ec_public_key_from_private_key(recipient_private_key, recipient_public_key) != 0) {
        printf("  FAIL: Failed to derive recipient public key\n");
        return -1;
    }
    
    // Test encryption
    char encrypted[8192];
    int encrypt_result = nostr_nip44_encrypt(
        sender_private_key,
        recipient_public_key,
        tv->plaintext,
        encrypted,
        sizeof(encrypted)
    );
    
    if (encrypt_result != NOSTR_SUCCESS) {
        printf("  FAIL: Encryption - Expected: %d, Actual: %d\n", NOSTR_SUCCESS, encrypt_result);
        return -1;
    }
    
    // Test decryption - use recipient private key + sender public key
    char decrypted[8192];
    int decrypt_result = nostr_nip44_decrypt(
        recipient_private_key,
        sender_public_key,
        encrypted,
        decrypted,
        sizeof(decrypted)
    );
    
    if (decrypt_result != NOSTR_SUCCESS) {
        printf("  FAIL: Decryption - Expected: %d, Actual: %d\n", NOSTR_SUCCESS, decrypt_result);
        return -1;
    }
    
    // Verify round-trip
    if (strcmp(tv->plaintext, decrypted) != 0) {
        printf("  FAIL: Round-trip mismatch\n");
        printf("  Expected: \"%s\"\n", tv->plaintext);
        printf("  Actual:   \"%s\"\n", decrypted);
        return -1;
    }
    
    printf("  PASS: Expected: \"%s\", Actual: \"%s\"\n", tv->plaintext, decrypted);
    printf("  Encrypted output: %s\n", encrypted);
    
    return 0;
}

static int test_nip44_error_conditions() {
    printf("Test: NIP-44 error conditions\n");
    
    // Use proper valid secp256k1 private keys
    unsigned char valid_sender_key[32];
    unsigned char valid_recipient_key[32];
    unsigned char valid_recipient_pubkey[32];
    
    hex_to_bytes("0000000000000000000000000000000000000000000000000000000000000001", valid_sender_key, 32);
    hex_to_bytes("0000000000000000000000000000000000000000000000000000000000000002", valid_recipient_key, 32);
    
    // Generate the recipient's public key
    if (nostr_ec_public_key_from_private_key(valid_recipient_key, valid_recipient_pubkey) != 0) {
        printf("  FAIL: Failed to generate recipient public key\n");
        return -1;
    }
    
    char output[1024];
    
    // Test NULL parameters
    int result = nostr_nip44_encrypt(NULL, valid_recipient_pubkey, "test", output, sizeof(output));
    if (result != NOSTR_ERROR_INVALID_INPUT) {
        printf("  FAIL: NULL sender key - Expected: %d, Actual: %d\n", NOSTR_ERROR_INVALID_INPUT, result);
        return -1;
    }
    
    result = nostr_nip44_encrypt(valid_sender_key, NULL, "test", output, sizeof(output));
    if (result != NOSTR_ERROR_INVALID_INPUT) {
        printf("  FAIL: NULL recipient key - Expected: %d, Actual: %d\n", NOSTR_ERROR_INVALID_INPUT, result);
        return -1;
    }
    
    result = nostr_nip44_encrypt(valid_sender_key, valid_recipient_pubkey, NULL, output, sizeof(output));
    if (result != NOSTR_ERROR_INVALID_INPUT) {
        printf("  FAIL: NULL plaintext - Expected: %d, Actual: %d\n", NOSTR_ERROR_INVALID_INPUT, result);
        return -1;
    }
    
    result = nostr_nip44_encrypt(valid_sender_key, valid_recipient_pubkey, "test", NULL, sizeof(output));
    if (result != NOSTR_ERROR_INVALID_INPUT) {
        printf("  FAIL: NULL output buffer - Expected: %d, Actual: %d\n", NOSTR_ERROR_INVALID_INPUT, result);
        return -1;
    }
    
    // Test buffer too small
    char small_buffer[10];
    result = nostr_nip44_encrypt(valid_sender_key, valid_recipient_pubkey, "test message", small_buffer, sizeof(small_buffer));
    if (result != NOSTR_ERROR_NIP44_BUFFER_TOO_SMALL) {
        printf("  FAIL: Buffer too small - Expected: %d, Actual: %d\n", NOSTR_ERROR_NIP44_BUFFER_TOO_SMALL, result);
        return -1;
    }
    
    printf("  PASS: All error conditions handled correctly\n");
    return 0;
}

static int test_nip44_decryption_vector(const nip44_test_vector_t* tv) {
    printf("Test: %s\n", tv->name);
    
    // Parse keys
    unsigned char sender_private_key[32];
    unsigned char recipient_private_key[32];
    
    if (hex_to_bytes(tv->sender_private_key_hex, sender_private_key, 32) != 0) {
        printf("  FAIL: Failed to parse sender private key\n");
        return -1;
    }
    
    if (hex_to_bytes(tv->recipient_private_key_hex, recipient_private_key, 32) != 0) {
        printf("  FAIL: Failed to parse recipient private key\n");
        return -1;
    }
    
    // Generate the public keys from the private keys
    unsigned char sender_public_key[32];
    
    if (nostr_ec_public_key_from_private_key(sender_private_key, sender_public_key) != 0) {
        printf("  FAIL: Failed to derive sender public key\n");
        return -1;
    }
    
    // Test decryption of known vector
    char decrypted[8192];
    int decrypt_result = nostr_nip44_decrypt(
        recipient_private_key,
        sender_public_key,
        tv->expected_encrypted,
        decrypted,
        sizeof(decrypted)
    );
    
    if (decrypt_result != NOSTR_SUCCESS) {
        printf("  FAIL: Decryption - Expected: %d, Actual: %d\n", NOSTR_SUCCESS, decrypt_result);
        printf("  Input payload: %s\n", tv->expected_encrypted);
        return -1;
    }
    
    // Verify decrypted plaintext matches expected
    if (strcmp(tv->plaintext, decrypted) != 0) {
        printf("  FAIL: Plaintext mismatch\n");
        printf("  Expected: \"%s\"\n", tv->plaintext);
        printf("  Actual:   \"%s\"\n", decrypted);
        return -1;
    }
    
    printf("  PASS: Expected: \"%s\", Actual: \"%s\"\n", tv->plaintext, decrypted);
    
    return 0;
}

static int test_nip44_encryption_variability() {
    printf("Test: NIP-44 encryption variability (non-deterministic)\n");
    
    const char* test_message = "Test message for variability";
    unsigned char sender_key[32], recipient_key[32];
    
    // Use fixed test keys
    hex_to_bytes("1111111111111111111111111111111111111111111111111111111111111111", sender_key, 32);
    hex_to_bytes("2222222222222222222222222222222222222222222222222222222222222222", recipient_key, 32);
    
    // Generate recipient public key
    unsigned char recipient_pubkey[32];
    if (nostr_ec_public_key_from_private_key(recipient_key, recipient_pubkey) != 0) {
        printf("  FAIL: Failed to generate recipient public key\n");
        return -1;
    }
    
    // Encrypt the same message multiple times
    char encrypted1[8192], encrypted2[8192], encrypted3[8192];
    
    int result1 = nostr_nip44_encrypt(sender_key, recipient_pubkey, test_message, encrypted1, sizeof(encrypted1));
    int result2 = nostr_nip44_encrypt(sender_key, recipient_pubkey, test_message, encrypted2, sizeof(encrypted2));
    int result3 = nostr_nip44_encrypt(sender_key, recipient_pubkey, test_message, encrypted3, sizeof(encrypted3));
    
    if (result1 != NOSTR_SUCCESS || result2 != NOSTR_SUCCESS || result3 != NOSTR_SUCCESS) {
        printf("  FAIL: Encryption failed - Results: %d, %d, %d\n", result1, result2, result3);
        return -1;
    }
    
    // Verify all ciphertexts are different (non-deterministic)
    if (strcmp(encrypted1, encrypted2) == 0 || strcmp(encrypted1, encrypted3) == 0 || strcmp(encrypted2, encrypted3) == 0) {
        printf("  FAIL: NIP-44 encryption should produce different ciphertext each time\n");
        printf("  Encryption 1: %.50s...\n", encrypted1);
        printf("  Encryption 2: %.50s...\n", encrypted2);
        printf("  Encryption 3: %.50s...\n", encrypted3);
        return -1;
    }
    
    // Verify all decrypt to the same plaintext
    unsigned char sender_pubkey[32];
    if (nostr_ec_public_key_from_private_key(sender_key, sender_pubkey) != 0) {
        printf("  FAIL: Failed to generate sender public key\n");
        return -1;
    }
    
    char decrypted1[8192], decrypted2[8192], decrypted3[8192];
    
    int decrypt1 = nostr_nip44_decrypt(recipient_key, sender_pubkey, encrypted1, decrypted1, sizeof(decrypted1));
    int decrypt2 = nostr_nip44_decrypt(recipient_key, sender_pubkey, encrypted2, decrypted2, sizeof(decrypted2));
    int decrypt3 = nostr_nip44_decrypt(recipient_key, sender_pubkey, encrypted3, decrypted3, sizeof(decrypted3));
    
    if (decrypt1 != NOSTR_SUCCESS || decrypt2 != NOSTR_SUCCESS || decrypt3 != NOSTR_SUCCESS) {
        printf("  FAIL: Decryption failed - Results: %d, %d, %d\n", decrypt1, decrypt2, decrypt3);
        return -1;
    }
    
    if (strcmp(decrypted1, test_message) != 0 || strcmp(decrypted2, test_message) != 0 || strcmp(decrypted3, test_message) != 0) {
        printf("  FAIL: Decryption mismatch\n");
        printf("  Expected: \"%s\"\n", test_message);
        printf("  Decrypted1: \"%s\"\n", decrypted1);
        printf("  Decrypted2: \"%s\"\n", decrypted2);
        printf("  Decrypted3: \"%s\"\n", decrypted3);
        return -1;
    }
    
    printf("  PASS: All encryptions different, all decrypt to: \"%s\"\n", test_message);
    printf("  Sample ciphertext lengths: %zu, %zu, %zu bytes\n", strlen(encrypted1), strlen(encrypted2), strlen(encrypted3));
    
    return 0;
}


int main() {
    printf("NIP-44 Encryption Test Suite\n");
    printf("=============================\n");
    
    // Initialize the library
    if (nostr_init() != NOSTR_SUCCESS) {
        printf("FAIL: Failed to initialize NOSTR library\n");
        return 1;
    }
    
    int total_tests = 0;
    int passed_tests = 0;
    
    // Test all vectors
    size_t num_vectors = sizeof(test_vectors) / sizeof(test_vectors[0]);
    for (size_t i = 0; i < num_vectors; i++) {
        total_tests++;
        if (test_nip44_round_trip(&test_vectors[i]) == 0) {
            passed_tests++;
        }
    }
    
    // Test decryption vectors (cross-compatibility)
    size_t num_decryption_vectors = sizeof(decryption_test_vectors) / sizeof(decryption_test_vectors[0]);
    for (size_t i = 0; i < num_decryption_vectors; i++) {
        total_tests++;
        if (test_nip44_decryption_vector(&decryption_test_vectors[i]) == 0) {
            passed_tests++;
        }
    }
    
    // Test encryption variability (NIP-44 non-deterministic behavior)
    total_tests++;
    if (test_nip44_encryption_variability() == 0) {
        passed_tests++;
    }
    
    // Test error conditions
    total_tests++;
    if (test_nip44_error_conditions() == 0) {
        passed_tests++;
    }
    
    // Final results
    printf("\nTest Results: %d/%d passed\n", passed_tests, total_tests);
    
    if (passed_tests == total_tests) {
        printf("All NIP-44 tests PASSED\n");
        nostr_cleanup();
        return 0;
    } else {
        printf("Some tests FAILED\n");
        nostr_cleanup();
        return 1;
    }
}