//! Memory Encryption Module
//!
//! Provides AES-256-GCM encryption for sensitive memory content.
//!
//! NOTE: Some constants and types are defined for future use.

#![allow(dead_code)] // Crypto utilities reserved for future encryption features

use aes_gcm::{
    aead::{Aead, KeyInit, OsRng},
    Aes256Gcm, Nonce,
};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use rand::RngCore;
use sha2::{Digest, Sha256};

/// Encryption key size (256 bits = 32 bytes)
pub const KEY_SIZE: usize = 32;
/// Nonce size for AES-GCM (96 bits = 12 bytes)
const NONCE_SIZE: usize = 12;

/// Encryption error type
#[derive(Debug)]
pub enum CryptoError {
    InvalidKeyLength,
    EncryptionFailed(String),
    DecryptionFailed(String),
    InvalidBase64(String),
    InvalidNonce,
    InvalidUtf8(String),
}

impl std::fmt::Display for CryptoError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CryptoError::InvalidKeyLength => write!(f, "Invalid encryption key length"),
            CryptoError::EncryptionFailed(e) => write!(f, "Encryption failed: {}", e),
            CryptoError::DecryptionFailed(e) => write!(f, "Decryption failed: {}", e),
            CryptoError::InvalidBase64(e) => write!(f, "Invalid base64: {}", e),
            CryptoError::InvalidNonce => write!(f, "Invalid nonce"),
            CryptoError::InvalidUtf8(e) => write!(f, "Invalid UTF-8: {}", e),
        }
    }
}

impl std::error::Error for CryptoError {}

/// Derive a 256-bit key from a password using SHA-256
pub fn derive_key(password: &str) -> [u8; KEY_SIZE] {
    let mut hasher = Sha256::new();
    hasher.update(password.as_bytes());
    let result = hasher.finalize();
    let mut key = [0u8; KEY_SIZE];
    key.copy_from_slice(&result);
    key
}

/// Generate a random encryption key
pub fn generate_key() -> [u8; KEY_SIZE] {
    let mut key = [0u8; KEY_SIZE];
    OsRng.fill_bytes(&mut key);
    key
}

/// Generate a random nonce
fn generate_nonce() -> [u8; NONCE_SIZE] {
    let mut nonce = [0u8; NONCE_SIZE];
    OsRng.fill_bytes(&mut nonce);
    nonce
}

/// Encrypt plaintext using AES-256-GCM
/// Returns base64-encoded ciphertext (nonce + encrypted data)
pub fn encrypt(plaintext: &str, key: &[u8; KEY_SIZE]) -> Result<String, CryptoError> {
    let cipher = Aes256Gcm::new_from_slice(key)
        .map_err(|e| CryptoError::EncryptionFailed(e.to_string()))?;
    
    let nonce_bytes = generate_nonce();
    let nonce = Nonce::from_slice(&nonce_bytes);
    
    let ciphertext = cipher
        .encrypt(nonce, plaintext.as_bytes())
        .map_err(|e| CryptoError::EncryptionFailed(e.to_string()))?;
    
    let mut combined = nonce_bytes.to_vec();
    combined.extend(ciphertext);
    
    Ok(BASE64.encode(&combined))
}

/// Decrypt ciphertext using AES-256-GCM
/// Expects base64-encoded ciphertext (nonce + encrypted data)
pub fn decrypt(ciphertext_b64: &str, key: &[u8; KEY_SIZE]) -> Result<String, CryptoError> {
    let combined = BASE64
        .decode(ciphertext_b64)
        .map_err(|e| CryptoError::InvalidBase64(e.to_string()))?;
    
    if combined.len() < NONCE_SIZE {
        return Err(CryptoError::InvalidNonce);
    }
    
    let (nonce_bytes, ciphertext) = combined.split_at(NONCE_SIZE);
    let nonce = Nonce::from_slice(nonce_bytes);
    
    let cipher = Aes256Gcm::new_from_slice(key)
        .map_err(|e| CryptoError::DecryptionFailed(e.to_string()))?;
    
    let plaintext = cipher
        .decrypt(nonce, ciphertext)
        .map_err(|e| CryptoError::DecryptionFailed(e.to_string()))?;
    
    String::from_utf8(plaintext)
        .map_err(|e| CryptoError::InvalidUtf8(e.to_string()))
}

/// Key storage key name in OS keyring
pub const MEMORY_ENCRYPTION_KEY_NAME: &str = "zclaw_memory_encryption_key";

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_encrypt_decrypt() {
        let key = generate_key();
        let plaintext = "Hello, ZCLAW!";
        
        let encrypted = encrypt(plaintext, &key).unwrap();
        let decrypted = decrypt(&encrypted, &key).unwrap();
        
        assert_eq!(plaintext, decrypted);
    }
    
    #[test]
    fn test_derive_key() {
        let key1 = derive_key("password123");
        let key2 = derive_key("password123");
        let key3 = derive_key("different");
        
        assert_eq!(key1, key2);
        assert_ne!(key1, key3);
    }
    
    #[test]
    fn test_encrypt_produces_different_ciphertext() {
        let key = generate_key();
        let plaintext = "Same message";
        
        let encrypted1 = encrypt(plaintext, &key).unwrap();
        let encrypted2 = encrypt(plaintext, &key).unwrap();
        
        // Different nonces should produce different ciphertext
        assert_ne!(encrypted1, encrypted2);
        
        // But both should decrypt to the same plaintext
        assert_eq!(plaintext, decrypt(&encrypted1, &key).unwrap());
        assert_eq!(plaintext, decrypt(&encrypted2, &key).unwrap());
    }
}