/*
 * NIP-04: Encrypted Direct Message Implementation
 * https://github.com/nostr-protocol/nips/blob/master/04.md
 */

#include "nip004.h"
#include "utils.h"
#include "nostr_common.h"
#include "crypto/nostr_secp256k1.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

// Include our AES implementation
#include "crypto/nostr_aes.h"

// Forward declarations for internal functions
static int aes_cbc_encrypt(const unsigned char* key, const unsigned char* iv,
                          const unsigned char* input, size_t input_len,
                          unsigned char* output);
static int aes_cbc_decrypt(const unsigned char* key, const unsigned char* iv,
                          const unsigned char* input, size_t input_len,
                          unsigned char* output);
static size_t pkcs7_pad(unsigned char* data, size_t data_len, size_t block_size);
static size_t pkcs7_unpad(unsigned char* data, size_t data_len);

// Memory clearing utility
static void memory_clear(const void *p, size_t len) {
    if (p && len) {
        memset((void *)p, 0, len);
    }
}

// =============================================================================
// AES-256-CBC ENCRYPTION/DECRYPTION USING TINYAES
// =============================================================================

static int aes_cbc_encrypt(const unsigned char* key, const unsigned char* iv,
                          const unsigned char* input, size_t input_len,
                          unsigned char* output) {
    if (!key || !iv || !input || !output || input_len % 16 != 0) {
        return -1;
    }
    
    // Initialize AES context with key and IV
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    
    // Copy input to output (tinyAES works in-place)
    memcpy(output, input, input_len);
    
    // Encrypt using AES-256-CBC
    AES_CBC_encrypt_buffer(&ctx, output, input_len);
    
    return 0;
}

static int aes_cbc_decrypt(const unsigned char* key, const unsigned char* iv,
                          const unsigned char* input, size_t input_len,
                          unsigned char* output) {
    if (!key || !iv || !input || !output || input_len % 16 != 0) {
        return -1;
    }
    
    // Initialize AES context with key and IV
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    
    // Copy input to output (tinyAES works in-place)
    memcpy(output, input, input_len);
    
    // Decrypt using AES-256-CBC
    AES_CBC_decrypt_buffer(&ctx, output, input_len);
    
    return 0;
}

// PKCS#7 padding functions
static size_t pkcs7_pad(unsigned char* data, size_t data_len, size_t block_size) {
    size_t padding = block_size - (data_len % block_size);
    for (size_t i = 0; i < padding; i++) {
        data[data_len + i] = (unsigned char)padding;
    }
    return data_len + padding;
}

static size_t pkcs7_unpad(unsigned char* data, size_t data_len) {
    if (data_len == 0) return 0;
    
    unsigned char padding = data[data_len - 1];
    if (padding == 0 || padding > 16) return 0; // Invalid padding
    
    // Verify padding
    for (size_t i = data_len - padding; i < data_len; i++) {
        if (data[i] != padding) return 0; // Invalid padding
    }
    
    return data_len - padding;
}

// =============================================================================
// NIP-04 IMPLEMENTATION
// =============================================================================

int nostr_nip04_encrypt(const unsigned char* sender_private_key,
                       const unsigned char* recipient_public_key, 
                       const char* plaintext,
                       char* output,
                       size_t output_size) {
    printf("[DEBUG] nostr_nip04_encrypt: Starting encryption\n");
    
    if (!sender_private_key || !recipient_public_key || !plaintext || !output) {
        printf("[DEBUG] nostr_nip04_encrypt: Invalid input - null pointer detected\n");
        return NOSTR_ERROR_INVALID_INPUT;
    }
    
    size_t plaintext_len = strlen(plaintext);
    printf("[DEBUG] nostr_nip04_encrypt: Plaintext length = %zu\n", plaintext_len);
    
    if (plaintext_len > NOSTR_NIP04_MAX_PLAINTEXT_SIZE) {
        printf("[DEBUG] nostr_nip04_encrypt: Plaintext too large (%zu > %d)\n", plaintext_len, NOSTR_NIP04_MAX_PLAINTEXT_SIZE);
        return NOSTR_ERROR_NIP04_BUFFER_TOO_SMALL;
    }
    
    // FIX: Calculate final size requirements EARLY before any allocations
    // CRITICAL: Account for PKCS#7 padding which ALWAYS adds 1-16 bytes
    // If plaintext_len is a multiple of 16, PKCS#7 adds a full 16-byte block
    size_t padded_len = ((plaintext_len / 16) + 1) * 16; // Always add one full block for PKCS#7
    size_t ciphertext_b64_max = ((padded_len + 2) / 3) * 4 + 1;
    size_t iv_b64_max = ((16 + 2) / 3) * 4 + 1;  // Always 25 bytes
    size_t estimated_result_len = ciphertext_b64_max + 4 + iv_b64_max;  // +4 for "?iv="
    
    printf("[DEBUG] nostr_nip04_encrypt: Size calculations - padded_len=%zu, ciphertext_b64_max=%zu, estimated_result_len=%zu, output_size=%zu\n", 
           padded_len, ciphertext_b64_max, estimated_result_len, output_size);
    
    // FIX: Check output buffer size BEFORE doing any work
    if (estimated_result_len > output_size) {
        printf("[DEBUG] nostr_nip04_encrypt: Output buffer too small (%zu > %zu)\n", estimated_result_len, output_size);
        return NOSTR_ERROR_NIP04_BUFFER_TOO_SMALL;
    }
    
    // Step 1: Compute ECDH shared secret
    printf("[DEBUG] nostr_nip04_encrypt: Computing ECDH shared secret\n");
    unsigned char shared_secret[32];
    if (ecdh_shared_secret(sender_private_key, recipient_public_key, shared_secret) != 0) {
        printf("[DEBUG] nostr_nip04_encrypt: ECDH shared secret computation failed\n");
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    printf("[DEBUG] nostr_nip04_encrypt: ECDH shared secret computed successfully\n");
    
    // Step 2: Generate random IV (16 bytes)
    printf("[DEBUG] nostr_nip04_encrypt: Generating random IV\n");
    unsigned char iv[16];
        if (nostr_secp256k1_get_random_bytes(iv, 16) != 1) {
        printf("[DEBUG] nostr_nip04_encrypt: Failed to generate random IV\n");
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    printf("[DEBUG] nostr_nip04_encrypt: IV generated successfully\n");
    
    // Step 3: Pad plaintext using PKCS#7
    printf("[DEBUG] nostr_nip04_encrypt: Padding plaintext with PKCS#7\n");
    unsigned char* padded_data = malloc(padded_len);
    if (!padded_data) {
        printf("[DEBUG] nostr_nip04_encrypt: Failed to allocate memory for padded data\n");
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    memcpy(padded_data, plaintext, plaintext_len);
    size_t actual_padded_len = pkcs7_pad(padded_data, plaintext_len, 16);
    printf("[DEBUG] nostr_nip04_encrypt: PKCS#7 padding completed - actual_padded_len=%zu\n", actual_padded_len);
    
    // Step 4: Encrypt using AES-256-CBC
    printf("[DEBUG] nostr_nip04_encrypt: Encrypting with AES-256-CBC\n");
    unsigned char* ciphertext = malloc(padded_len);
    if (!ciphertext) {
        printf("[DEBUG] nostr_nip04_encrypt: Failed to allocate memory for ciphertext\n");
        free(padded_data);
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    if (aes_cbc_encrypt(shared_secret, iv, padded_data, actual_padded_len, ciphertext) != 0) {
        printf("[DEBUG] nostr_nip04_encrypt: AES-256-CBC encryption failed\n");
        free(padded_data);
        free(ciphertext);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    printf("[DEBUG] nostr_nip04_encrypt: AES-256-CBC encryption completed successfully\n");
    
    // Step 5: Base64 encode ciphertext and IV
    printf("[DEBUG] nostr_nip04_encrypt: Base64 encoding ciphertext and IV\n");
    size_t ciphertext_b64_len = ((actual_padded_len + 2) / 3) * 4 + 1;
    size_t iv_b64_len = ((16 + 2) / 3) * 4 + 1;
    
    char* ciphertext_b64 = malloc(ciphertext_b64_len);
    char* iv_b64 = malloc(iv_b64_len);
    
    if (!ciphertext_b64 || !iv_b64) {
        printf("[DEBUG] nostr_nip04_encrypt: Failed to allocate memory for base64 buffers\n");
        free(padded_data);
        free(ciphertext);
        free(ciphertext_b64);
        free(iv_b64);
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    // FIX: Pass buffer sizes to base64_encode and check for success
    size_t ct_b64_len = base64_encode(ciphertext, actual_padded_len, ciphertext_b64, ciphertext_b64_len);
    size_t iv_b64_len_actual = base64_encode(iv, 16, iv_b64, iv_b64_len);
    
    printf("[DEBUG] nostr_nip04_encrypt: Base64 encoding results - ct_b64_len=%zu, iv_b64_len_actual=%zu\n", ct_b64_len, iv_b64_len_actual);
    
    // FIX: Check if encoding succeeded
    if (ct_b64_len == 0 || iv_b64_len_actual == 0) {
        printf("[DEBUG] nostr_nip04_encrypt: Base64 encoding failed\n");
        free(padded_data);
        free(ciphertext);
        free(ciphertext_b64);
        free(iv_b64);
        memory_clear(shared_secret, 32);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    printf("[DEBUG] nostr_nip04_encrypt: Base64 encoding completed successfully\n");
    
    // Step 6: Format as "ciphertext?iv=iv_base64" - size check moved earlier, now guaranteed to fit
    printf("[DEBUG] nostr_nip04_encrypt: Formatting final output string\n");
    size_t result_len = ct_b64_len + 4 + iv_b64_len_actual + 1; // +4 for "?iv=", +1 for null
    
    printf("[DEBUG] nostr_nip04_encrypt: Final size check - result_len=%zu, output_size=%zu\n", result_len, output_size);
    
    if (result_len > output_size) {
        printf("[DEBUG] nostr_nip04_encrypt: Final output buffer too small (%zu > %zu)\n", result_len, output_size);
        free(padded_data);
        free(ciphertext);
        free(ciphertext_b64);
        free(iv_b64);
        memory_clear(shared_secret, 32);
        return NOSTR_ERROR_NIP04_BUFFER_TOO_SMALL;
    }
    
    snprintf(output, output_size, "%s?iv=%s", ciphertext_b64, iv_b64);
    printf("[DEBUG] nostr_nip04_encrypt: Formatted output successfully\n");
    printf("[DEBUG] nostr_nip04_encrypt: Encryption completed successfully - returning NOSTR_SUCCESS\n");
    
    // Cleanup
    memory_clear(shared_secret, 32);
    memory_clear(padded_data, padded_len);
    memory_clear(ciphertext, padded_len);
    free(padded_data);
    free(ciphertext);
    free(ciphertext_b64);
    free(iv_b64);
    
    return NOSTR_SUCCESS;
}

int nostr_nip04_decrypt(const unsigned char* recipient_private_key,
                       const unsigned char* sender_public_key,
                       const char* encrypted_data,
                       char* output,
                       size_t output_size) {
    if (!recipient_private_key || !sender_public_key || !encrypted_data || !output) {
        return NOSTR_ERROR_INVALID_INPUT;
    }
    
    // Step 1: Parse encrypted data format "ciphertext?iv=iv_base64"
    char* separator = strstr(encrypted_data, "?iv=");
    if (!separator) {
        return NOSTR_ERROR_NIP04_INVALID_FORMAT;
    }
    
    size_t ciphertext_b64_len = separator - encrypted_data;
    const char* iv_b64 = separator + 4; // Skip "?iv="
    
    if (ciphertext_b64_len == 0 || strlen(iv_b64) == 0) {
        return NOSTR_ERROR_NIP04_INVALID_FORMAT;
    }
    
    // Step 2: Create null-terminated copy of ciphertext base64
    char* ciphertext_b64 = malloc(ciphertext_b64_len + 1);
    if (!ciphertext_b64) {
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    memcpy(ciphertext_b64, encrypted_data, ciphertext_b64_len);
    ciphertext_b64[ciphertext_b64_len] = '\0';
    
    // Step 3: Calculate proper buffer sizes for decoded data
    // Base64 decoding: 4 chars -> 3 bytes, so max decoded size is (len * 3) / 4
    size_t max_ciphertext_len = ((ciphertext_b64_len + 3) / 4) * 3;
    size_t max_iv_len = ((strlen(iv_b64) + 3) / 4) * 3;
    
    // Allocate buffers with proper sizes
    unsigned char* ciphertext = malloc(max_ciphertext_len);
    unsigned char* iv_buffer = malloc(max_iv_len);
    
    if (!ciphertext || !iv_buffer) {
        free(ciphertext_b64);
        free(ciphertext);
        free(iv_buffer);
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    // Step 4: Base64 decode ciphertext and IV
    size_t ciphertext_len = base64_decode(ciphertext_b64, ciphertext);
    size_t iv_len = base64_decode(iv_b64, iv_buffer);
    
    if (ciphertext_len == 0 || iv_len != 16 || ciphertext_len % 16 != 0) {
        free(ciphertext_b64);
        free(ciphertext);
        free(iv_buffer);
        return NOSTR_ERROR_NIP04_INVALID_FORMAT;
    }
    
    // Copy IV to fixed-size buffer for safety
    unsigned char iv[16];
    memcpy(iv, iv_buffer, 16);
    free(iv_buffer);
    
    // Step 5: Compute ECDH shared secret
    unsigned char shared_secret[32];
    if (ecdh_shared_secret(recipient_private_key, sender_public_key, shared_secret) != 0) {
        free(ciphertext_b64);
        free(ciphertext);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Step 6: Decrypt using AES-256-CBC
    unsigned char* plaintext_padded = malloc(ciphertext_len);
    if (!plaintext_padded) {
        free(ciphertext_b64);
        free(ciphertext);
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    if (aes_cbc_decrypt(shared_secret, iv, ciphertext, ciphertext_len, plaintext_padded) != 0) {
        free(ciphertext_b64);
        free(ciphertext);
        free(plaintext_padded);
        return NOSTR_ERROR_NIP04_DECRYPT_FAILED;
    }
    
    // Step 7: Remove PKCS#7 padding
    size_t plaintext_len = pkcs7_unpad(plaintext_padded, ciphertext_len);
    if (plaintext_len == 0 || plaintext_len > ciphertext_len) {
        free(ciphertext_b64);
        free(ciphertext);
        free(plaintext_padded);
        return NOSTR_ERROR_NIP04_DECRYPT_FAILED;
    }
    
    // Step 8: Copy to output buffer and null-terminate
    if (plaintext_len + 1 > output_size) {
        free(ciphertext_b64);
        free(ciphertext);
        free(plaintext_padded);
        return NOSTR_ERROR_NIP04_BUFFER_TOO_SMALL;
    }
    
    memcpy(output, plaintext_padded, plaintext_len);
    output[plaintext_len] = '\0';
    
    // Cleanup
    memory_clear(shared_secret, 32);
    memory_clear(plaintext_padded, ciphertext_len);
    free(ciphertext_b64);
    free(ciphertext);
    free(plaintext_padded);
    
    return NOSTR_SUCCESS;
}