//! Core cryptographic operations
//! 
//! This module implements the core encryption and decryption functionality
//! using XChaCha20-Poly1305 for AEAD encryption and Argon2id for password-based
//! key derivation.

use crate::error::{CryptoError, Result};
use argon2::{Argon2, Params};
use chacha20poly1305::{
    aead::{Aead, AeadCore, KeyInit, OsRng},
    XChaCha20Poly1305, XNonce,
};
use rand::RngCore;
use zeroize::ZeroizeOnDrop;

/// Size of encryption keys (256 bits)
pub const KEY_SIZE: usize = 32;

/// Size of nonces for XChaCha20 (192 bits)
pub const NONCE_SIZE: usize = 24;

/// Size of salts for Argon2 (128 bits)
pub const SALT_SIZE: usize = 16;

/// Authentication tag size (128 bits)
pub const TAG_SIZE: usize = 16;

/// Wrapper for key material that zeros on drop
#[derive(Clone, ZeroizeOnDrop)]
pub struct KeyMaterial {
    key: [u8; KEY_SIZE],
}

impl KeyMaterial {
    /// Create a new KeyMaterial from bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() != KEY_SIZE {
            return Err(CryptoError::InvalidKeyLength {
                expected: KEY_SIZE,
                actual: bytes.len(),
            });
        }
        let mut key = [0u8; KEY_SIZE];
        key.copy_from_slice(bytes);
        Ok(Self { key })
    }

    /// Get the key bytes
    pub fn as_bytes(&self) -> &[u8; KEY_SIZE] {
        &self.key
    }

    /// Convert to owned bytes
    pub fn to_vec(&self) -> Vec<u8> {
        self.key.to_vec()
    }
}

/// Generate a secure random encryption key
/// 
/// This generates a cryptographically secure 256-bit key suitable for
/// direct use with XChaCha20-Poly1305 encryption.
/// 
/// # Returns
/// 
/// A 32-byte key as a Vec<u8>
/// 
/// # Example
/// 
/// ```rust
/// use irene_cipher::generate_key;
/// 
/// let key = generate_key();
/// assert_eq!(key.len(), 32);
/// ```
pub fn generate_key() -> Vec<u8> {
    let mut key = vec![0u8; KEY_SIZE];
    OsRng.fill_bytes(&mut key);
    key
}

/// Derive a key from a password using Argon2id
/// 
/// Uses memory-hard Argon2id with the following parameters:
/// - Memory cost: 64 MiB
/// - Time cost: 3 iterations
/// - Parallelism: 4 lanes
/// 
/// # Arguments
/// 
/// * `password` - The password to derive from
/// * `salt` - A unique salt (16 bytes)
/// 
/// # Returns
/// 
/// A 32-byte derived key
pub(crate) fn derive_key_from_password(password: &str, salt: &[u8]) -> Result<KeyMaterial> {
    if salt.len() != SALT_SIZE {
        return Err(CryptoError::InvalidFormat(
            "Salt must be 16 bytes".to_string(),
        ));
    }

    // Argon2id parameters: balanced security vs performance
    // 64 MiB memory, 3 iterations, 4 parallel lanes
    let params = Params::new(65536, 3, 4, Some(KEY_SIZE))
        .map_err(|e| CryptoError::KeyDerivationError(e.to_string()))?;

    let argon2 = Argon2::new(
        argon2::Algorithm::Argon2id,
        argon2::Version::V0x13,
        params,
    );

    let mut output = [0u8; KEY_SIZE];
    argon2
        .hash_password_into(password.as_bytes(), salt, &mut output)
        .map_err(|e| CryptoError::KeyDerivationError(e.to_string()))?;

    KeyMaterial::from_bytes(&output)
}

/// Encrypt bytes using XChaCha20-Poly1305 with password-based key derivation
/// 
/// This function performs authenticated encryption (AEAD) on the input data.
/// The output format is: [salt (16 bytes)][nonce (24 bytes)][ciphertext + tag]
/// 
/// # Arguments
/// 
/// * `plaintext` - The data to encrypt
/// * `password` - Password for key derivation
/// 
/// # Returns
/// 
/// Encrypted bytes with salt and nonce prepended
/// 
/// # Security
/// 
/// - Each encryption uses a unique random salt and nonce
/// - The ciphertext is authenticated (tampering will be detected on decryption)
/// - Password is derived using Argon2id (memory-hard)
pub fn encrypt_bytes(plaintext: &[u8], password: &str) -> Result<Vec<u8>> {
    // Generate random salt for key derivation
    let mut salt = [0u8; SALT_SIZE];
    OsRng.fill_bytes(&mut salt);

    // Derive key from password
    let key_material = derive_key_from_password(password, &salt)?;

    // Create cipher
    let cipher = XChaCha20Poly1305::new(key_material.as_bytes().into());

    // Generate random nonce
    let nonce = XChaCha20Poly1305::generate_nonce(&mut OsRng);

    // Encrypt
    let ciphertext = cipher
        .encrypt(&nonce, plaintext)
        .map_err(|e| CryptoError::EncryptionError(e.to_string()))?;

    // Format: [salt][nonce][ciphertext+tag]
    let mut output = Vec::with_capacity(SALT_SIZE + NONCE_SIZE + ciphertext.len());
    output.extend_from_slice(&salt);
    output.extend_from_slice(&nonce);
    output.extend_from_slice(&ciphertext);

    Ok(output)
}

/// Decrypt bytes using XChaCha20-Poly1305 with password-based key derivation
/// 
/// This function performs authenticated decryption on data encrypted with
/// `encrypt_bytes`. It will fail if the ciphertext has been tampered with
/// or if the password is incorrect.
/// 
/// # Arguments
/// 
/// * `ciphertext` - The encrypted data (with salt and nonce prepended)
/// * `password` - Password for key derivation
/// 
/// # Returns
/// 
/// Decrypted plaintext bytes
/// 
/// # Errors
/// 
/// Returns `CryptoError::AuthenticationFailed` if:
/// - The password is incorrect
/// - The ciphertext has been modified
/// - The authentication tag is invalid
pub fn decrypt_bytes(ciphertext: &[u8], password: &str) -> Result<Vec<u8>> {
    // Minimum size: salt + nonce + tag
    if ciphertext.len() < SALT_SIZE + NONCE_SIZE + TAG_SIZE {
        return Err(CryptoError::InvalidFormat(
            "Ciphertext too short".to_string(),
        ));
    }

    // Extract components
    let salt = &ciphertext[..SALT_SIZE];
    let nonce = &ciphertext[SALT_SIZE..SALT_SIZE + NONCE_SIZE];
    let encrypted_data = &ciphertext[SALT_SIZE + NONCE_SIZE..];

    // Derive key from password
    let key_material = derive_key_from_password(password, salt)?;

    // Create cipher
    let cipher = XChaCha20Poly1305::new(key_material.as_bytes().into());

    // Decrypt and verify authentication tag
    let nonce = XNonce::from_slice(nonce);
    let plaintext = cipher
        .decrypt(nonce, encrypted_data)
        .map_err(|_| CryptoError::AuthenticationFailed)?;

    Ok(plaintext)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_generate_key() {
        let key1 = generate_key();
        let key2 = generate_key();
        
        assert_eq!(key1.len(), KEY_SIZE);
        assert_eq!(key2.len(), KEY_SIZE);
        assert_ne!(key1, key2); // Should be random
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let plaintext = b"Hello, World! This is a secret message.";
        let password = "super_secure_password_123";

        let ciphertext = encrypt_bytes(plaintext, password).unwrap();
        assert_ne!(ciphertext.as_slice(), plaintext);
        assert!(ciphertext.len() > plaintext.len());

        let decrypted = decrypt_bytes(&ciphertext, password).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_wrong_password_fails() {
        let plaintext = b"Secret data";
        let ciphertext = encrypt_bytes(plaintext, "correct_password").unwrap();

        let result = decrypt_bytes(&ciphertext, "wrong_password");
        assert!(matches!(result, Err(CryptoError::AuthenticationFailed)));
    }

    #[test]
    fn test_tampered_ciphertext_fails() {
        let plaintext = b"Secret data";
        let mut ciphertext = encrypt_bytes(plaintext, "password").unwrap();

        // Tamper with the ciphertext
        let len = ciphertext.len();
        ciphertext[len - 1] ^= 0xFF;

        let result = decrypt_bytes(&ciphertext, "password");
        assert!(matches!(result, Err(CryptoError::AuthenticationFailed)));
    }

    #[test]
    fn test_empty_plaintext() {
        let plaintext = b"";
        let password = "password";

        let ciphertext = encrypt_bytes(plaintext, password).unwrap();
        let decrypted = decrypt_bytes(&ciphertext, password).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_large_plaintext() {
        let plaintext = vec![42u8; 1_000_000]; // 1 MB
        let password = "password";

        let ciphertext = encrypt_bytes(&plaintext, password).unwrap();
        let decrypted = decrypt_bytes(&ciphertext, password).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_nonce_uniqueness() {
        let plaintext = b"Same message";
        let password = "same_password";

        let ct1 = encrypt_bytes(plaintext, password).unwrap();
        let ct2 = encrypt_bytes(plaintext, password).unwrap();

        // Same plaintext and password should produce different ciphertext
        // due to random salt and nonce
        assert_ne!(ct1, ct2);
    }
}
