/*
 * NIP-44: Encrypted Payloads (Versioned) Implementation
 * https://github.com/nostr-protocol/nips/blob/master/44.md
 */

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

// Include our ChaCha20 implementation
#include "crypto/nostr_chacha20.h"

// Forward declarations for internal functions
static size_t calc_padded_len(size_t unpadded_len);
static unsigned char* pad_plaintext(const char* plaintext, size_t* padded_len);
static char* unpad_plaintext(const unsigned char* padded, size_t padded_len);
static int constant_time_compare(const unsigned char* a, const unsigned char* b, size_t len);

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

// =============================================================================
// NIP-44 UTILITY FUNCTIONS
// =============================================================================

// Constant-time comparison (security critical)
static int constant_time_compare(const unsigned char* a, const unsigned char* b, size_t len) {
    unsigned char result = 0;
    for (size_t i = 0; i < len; i++) {
        result |= (a[i] ^ b[i]);
    }
    return result == 0;
}

// NIP-44 padding calculation (per spec)
static size_t calc_padded_len(size_t unpadded_len) {
    if (unpadded_len <= 32) {
        return 32;
    }
    
    size_t next_power = 1;
    while (next_power < unpadded_len) {
        next_power <<= 1;
    }
    
    size_t chunk = (next_power <= 256) ? 32 : (next_power / 8);
    return chunk * ((unpadded_len - 1) / chunk + 1);
}

// NIP-44 padding (per spec)
static unsigned char* pad_plaintext(const char* plaintext, size_t* padded_len) {
    size_t unpadded_len = strlen(plaintext);
    if (unpadded_len > 65535) {
        return NULL;
    }
    
    size_t padded_content_len = calc_padded_len(unpadded_len);
    *padded_len = padded_content_len + 2; // Add 2 bytes for the length prefix
    unsigned char* padded = malloc(*padded_len);
    if (!padded) return NULL;
    
    // Write length prefix (big-endian u16)
    padded[0] = (unpadded_len >> 8) & 0xFF;
    padded[1] = unpadded_len & 0xFF;
    
    // Copy plaintext and add zero-padding
    memcpy(padded + 2, plaintext, unpadded_len);
    memset(padded + 2 + unpadded_len, 0, padded_content_len - unpadded_len);
    
    return padded;
}

// NIP-44 unpadding (per spec)
static char* unpad_plaintext(const unsigned char* padded, size_t padded_len) {
    if (padded_len < 4) return NULL; 

    size_t unpadded_len = (padded[0] << 8) | padded[1];
    size_t expected_padded_len = calc_padded_len(unpadded_len);

    printf("--- unpad_plaintext DEBUG ---\n");
    printf("padded_len: %zu\n", padded_len);
    printf("unpadded_len: %zu\n", unpadded_len);
    printf("expected_padded_len: %zu\n", expected_padded_len);
    printf("--- end unpad_plaintext DEBUG ---\n");

    if (padded_len != expected_padded_len) {
        return NULL;
    }

    if (unpadded_len > padded_len - 2) {
        return NULL;
    }

    char* plaintext = malloc(unpadded_len + 1);
    if (!plaintext) return NULL;
    
    // Handle empty message case (unpadded_len can be 0)
    if (unpadded_len > 0) {
        memcpy(plaintext, padded + 2, unpadded_len);
    }
    plaintext[unpadded_len] = '\0';
    
    return plaintext;
}

// =============================================================================
// NIP-44 IMPLEMENTATION
// =============================================================================

int nostr_nip44_encrypt_with_nonce(const unsigned char* sender_private_key,
                                  const unsigned char* recipient_public_key,
                                  const char* plaintext,
                                  const unsigned char* nonce,
                                  char* output,
                                  size_t output_size) {
    if (!sender_private_key || !recipient_public_key || !plaintext || !nonce || !output) {
        return NOSTR_ERROR_INVALID_INPUT;
    }
    
    size_t plaintext_len = strlen(plaintext);
    if (plaintext_len > NOSTR_NIP44_MAX_PLAINTEXT_SIZE) {
        return NOSTR_ERROR_NIP44_BUFFER_TOO_SMALL;
    }
    
    // Step 1: Compute ECDH shared secret
    unsigned char shared_secret[32];
    if (ecdh_shared_secret(sender_private_key, recipient_public_key, shared_secret) != 0) {
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Step 2: Calculate conversation key (HKDF-extract with "nip44-v2" as salt)
    unsigned char conversation_key[32];
    const char* salt_str = "nip44-v2";
    if (nostr_hkdf_extract((const unsigned char*)salt_str, strlen(salt_str),
                          shared_secret, 32, conversation_key) != 0) {
        memory_clear(shared_secret, 32);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Step 3: Use provided nonce (for testing)
    // Copy nonce for consistency with existing code structure
    unsigned char nonce_copy[32];
    memcpy(nonce_copy, nonce, 32);
    
    // Step 4: Derive message keys (HKDF-expand with nonce as info)
    unsigned char message_keys[76]; // 32 chacha_key + 12 chacha_nonce + 32 hmac_key
    if (nostr_hkdf_expand(conversation_key, 32, nonce_copy, 32, message_keys, 76) != 0) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(nonce_copy, 32);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    unsigned char* chacha_key = message_keys;
    unsigned char* chacha_nonce = message_keys + 32;
    unsigned char* hmac_key = message_keys + 44;
    
    // Step 5: Pad plaintext according to NIP-44 spec
    size_t padded_len;
    unsigned char* padded_plaintext = pad_plaintext(plaintext, &padded_len);
    if (!padded_plaintext) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(nonce, 32);
        memory_clear(message_keys, 76);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Step 6: Encrypt using ChaCha20
    unsigned char* ciphertext = malloc(padded_len);
    if (!ciphertext) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(nonce, 32);
        memory_clear(message_keys, 76);
        memory_clear(padded_plaintext, padded_len);
        free(padded_plaintext);
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    if (chacha20_encrypt(chacha_key, 0, chacha_nonce, padded_plaintext, ciphertext, padded_len) != 0) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(nonce, 32);
        memory_clear(message_keys, 76);
        memory_clear(padded_plaintext, padded_len);
        free(padded_plaintext);
        free(ciphertext);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Step 7: Compute HMAC with AAD (nonce + ciphertext)
    unsigned char* aad_data = malloc(32 + padded_len);
    if (!aad_data) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(nonce_copy, 32);
        memory_clear(message_keys, 76);
        memory_clear(padded_plaintext, padded_len);
        free(padded_plaintext);
        free(ciphertext);
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    memcpy(aad_data, nonce_copy, 32);
    memcpy(aad_data + 32, ciphertext, padded_len);
    
    unsigned char mac[32];
    if (nostr_hmac_sha256(hmac_key, 32, aad_data, 32 + padded_len, mac) != 0) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(nonce, 32);
        memory_clear(message_keys, 76);
        memory_clear(padded_plaintext, padded_len);
        memory_clear(aad_data, 32 + padded_len);
        free(padded_plaintext);
        free(ciphertext);
        free(aad_data);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Step 8: Format as base64(version + nonce + ciphertext + mac)
    size_t payload_len = 1 + 32 + padded_len + 32; // version + nonce + ciphertext + mac
    unsigned char* payload = malloc(payload_len);
    if (!payload) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(nonce, 32);
        memory_clear(message_keys, 76);
        memory_clear(padded_plaintext, padded_len);
        memory_clear(aad_data, 32 + padded_len);
        free(padded_plaintext);
        free(ciphertext);
        free(aad_data);
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    payload[0] = 0x02; // NIP-44 version 2
    memcpy(payload + 1, nonce_copy, 32);
    memcpy(payload + 33, ciphertext, padded_len);
    memcpy(payload + 33 + padded_len, mac, 32);
    
    // Base64 encode
    size_t b64_len = ((payload_len + 2) / 3) * 4 + 1;
    if (b64_len > output_size) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(nonce, 32);
        memory_clear(message_keys, 76);
        memory_clear(padded_plaintext, padded_len);
        memory_clear(aad_data, 32 + padded_len);
        memory_clear(payload, payload_len);
        free(padded_plaintext);
        free(ciphertext);
        free(aad_data);
        free(payload);
        return NOSTR_ERROR_NIP44_BUFFER_TOO_SMALL;
    }
    
    if (base64_encode(payload, payload_len, output, output_size) == 0) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(nonce, 32);
        memory_clear(message_keys, 76);
        memory_clear(padded_plaintext, padded_len);
        memory_clear(aad_data, 32 + padded_len);
        memory_clear(payload, payload_len);
        free(padded_plaintext);
        free(ciphertext);
        free(aad_data);
        free(payload);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Cleanup
    memory_clear(shared_secret, 32);
    memory_clear(conversation_key, 32);
    memory_clear(nonce_copy, 32);
    memory_clear(message_keys, 76);
    memory_clear(padded_plaintext, padded_len);
    memory_clear(aad_data, 32 + padded_len);
    memory_clear(payload, payload_len);
    free(padded_plaintext);
    free(ciphertext);
    free(aad_data);
    free(payload);
    
    return NOSTR_SUCCESS;
}

int nostr_nip44_encrypt(const unsigned char* sender_private_key,
                       const unsigned char* recipient_public_key,
                       const char* plaintext,
                       char* output,
                       size_t output_size) {
    // Generate random nonce and call the _with_nonce version
    unsigned char nonce[32];
    if (nostr_secp256k1_get_random_bytes(nonce, 32) != 1) {
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    return nostr_nip44_encrypt_with_nonce(sender_private_key, recipient_public_key, 
                                         plaintext, nonce, output, output_size);
}

int nostr_nip44_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: Base64 decode the encrypted data
    size_t max_payload_len = ((strlen(encrypted_data) + 3) / 4) * 3;
    unsigned char* payload = malloc(max_payload_len);
    if (!payload) {
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    size_t payload_len = base64_decode(encrypted_data, payload);
    if (payload_len < 66) { // Minimum: version(1) + nonce(32) + mac(32) + 1 byte ciphertext
        free(payload);
        return NOSTR_ERROR_NIP44_INVALID_FORMAT;
    }
    
    // Step 2: Extract components (version + nonce + ciphertext + mac)
    if (payload[0] != 0x02) { // Check NIP-44 version
        free(payload);
        return NOSTR_ERROR_NIP44_INVALID_FORMAT;
    }
    
    unsigned char* nonce = payload + 1;
    unsigned char* ciphertext = payload + 33;
    unsigned char* received_mac = payload + payload_len - 32;
    size_t ciphertext_len = (payload + payload_len - 32) - (payload + 33); // mac_start - ciphertext_start
    
    // Step 3: Compute ECDH shared secret
    unsigned char shared_secret[32];
    if (ecdh_shared_secret(recipient_private_key, sender_public_key, shared_secret) != 0) {
        memory_clear(payload, payload_len);
        free(payload);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Step 4: Calculate conversation key (HKDF-extract with "nip44-v2" as salt)
    unsigned char conversation_key[32];
    const char* salt_str = "nip44-v2";
    if (nostr_hkdf_extract((const unsigned char*)salt_str, strlen(salt_str),
                          shared_secret, 32, conversation_key) != 0) {
        memory_clear(shared_secret, 32);
        memory_clear(payload, payload_len);
        free(payload);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Step 5: Derive message keys (HKDF-expand with nonce as info)
    unsigned char message_keys[76]; // 32 chacha_key + 12 chacha_nonce + 32 hmac_key
    if (nostr_hkdf_expand(conversation_key, 32, nonce, 32, message_keys, 76) != 0) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(payload, payload_len);
        free(payload);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    unsigned char* chacha_key = message_keys;
    unsigned char* chacha_nonce = message_keys + 32;
    unsigned char* hmac_key = message_keys + 44;
    
    // Step 6: Verify HMAC with AAD (nonce + ciphertext)
    unsigned char* aad_data = malloc(32 + ciphertext_len);
    if (!aad_data) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(message_keys, 76);
        memory_clear(payload, payload_len);
        free(payload);
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    memcpy(aad_data, nonce, 32);
    memcpy(aad_data + 32, ciphertext, ciphertext_len);
    
    unsigned char computed_mac[32];
    if (nostr_hmac_sha256(hmac_key, 32, aad_data, 32 + ciphertext_len, computed_mac) != 0) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(message_keys, 76);
        memory_clear(aad_data, 32 + ciphertext_len);
        memory_clear(payload, payload_len);
        free(aad_data);
        free(payload);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Constant-time MAC verification
    // Constant-time MAC verification
    if (!constant_time_compare(received_mac, computed_mac, 32)) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(message_keys, 76);
        memory_clear(aad_data, 32 + ciphertext_len);
        memory_clear(payload, payload_len);
        free(aad_data);
        free(payload);
        return NOSTR_ERROR_NIP44_DECRYPT_FAILED;
    }
    
    // Step 7: Decrypt using ChaCha20
    unsigned char* padded_plaintext = malloc(ciphertext_len);
    if (!padded_plaintext) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(message_keys, 76);
        memory_clear(aad_data, 32 + ciphertext_len);
        memory_clear(payload, payload_len);
        free(aad_data);
        free(payload);
        return NOSTR_ERROR_MEMORY_FAILED;
    }
    
    if (chacha20_encrypt(chacha_key, 0, chacha_nonce, ciphertext, padded_plaintext, ciphertext_len) != 0) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(message_keys, 76);
        memory_clear(aad_data, 32 + ciphertext_len);
        memory_clear(payload, payload_len);
        free(aad_data);
        free(payload);
        free(padded_plaintext);
        return NOSTR_ERROR_CRYPTO_FAILED;
    }
    
    // Step 8: Remove padding according to NIP-44 spec
    // Step 8: Remove padding according to NIP-44 spec
    char* plaintext = unpad_plaintext(padded_plaintext, ciphertext_len);
    if (!plaintext) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(message_keys, 76);
        memory_clear(aad_data, 32 + ciphertext_len);
        memory_clear(payload, payload_len);
        memory_clear(padded_plaintext, ciphertext_len);
        free(aad_data);
        free(payload);
        free(padded_plaintext);
        return NOSTR_ERROR_NIP44_DECRYPT_FAILED;
    }
    
    // Step 9: Copy to output buffer
    size_t plaintext_len = strlen(plaintext);
    if (plaintext_len + 1 > output_size) {
        memory_clear(shared_secret, 32);
        memory_clear(conversation_key, 32);
        memory_clear(message_keys, 76);
        memory_clear(aad_data, 32 + ciphertext_len);
        memory_clear(payload, payload_len);
        memory_clear(padded_plaintext, ciphertext_len);
        memory_clear(plaintext, plaintext_len);
        free(aad_data);
        free(payload);
        free(padded_plaintext);
        free(plaintext);
        return NOSTR_ERROR_NIP44_BUFFER_TOO_SMALL;
    }
    
    strcpy(output, plaintext);
    
    // Cleanup
    memory_clear(shared_secret, 32);
    memory_clear(conversation_key, 32);
    memory_clear(message_keys, 76);
    memory_clear(aad_data, 32 + ciphertext_len);
    memory_clear(payload, payload_len);
    memory_clear(padded_plaintext, ciphertext_len);
    memory_clear(plaintext, plaintext_len);
    free(aad_data);
    free(payload);
    free(padded_plaintext);
    free(plaintext);
    
    return NOSTR_SUCCESS;
}