//! Implementation of symmetric encryption functionality.

use chacha20poly1305::{ChaCha20Poly1305, KeyInit, Nonce};
use chacha20poly1305::aead::Aead;
use rand::{rngs::OsRng, RngCore};
use serde::{Serialize, Deserialize};
use sha2::{Sha256, Digest};

use crate::hero_vault::error::CryptoError;
use crate::hero_vault::keypair::KeySpace;

/// The size of the nonce in bytes.
const NONCE_SIZE: usize = 12;

/// Generates a random 32-byte symmetric key.
///
/// # Returns
///
/// A 32-byte array containing the random key.
pub fn generate_symmetric_key() -> [u8; 32] {
    let mut key = [0u8; 32];
    OsRng.fill_bytes(&mut key);
    key
}

/// Derives a 32-byte key from a password.
///
/// # Arguments
///
/// * `password` - The password to derive the key from.
///
/// # Returns
///
/// A 32-byte array containing the derived key.
pub fn derive_key_from_password(password: &str) -> [u8; 32] {
    let mut hasher = Sha256::default();
    hasher.update(password.as_bytes());
    let result = hasher.finalize();
    
    let mut key = [0u8; 32];
    key.copy_from_slice(&result);
    key
}

/// Encrypts data using ChaCha20Poly1305 with an internally generated nonce.
///
/// The nonce is appended to the ciphertext so it can be extracted during decryption.
///
/// # Arguments
///
/// * `key` - The encryption key (should be 32 bytes).
/// * `message` - The message to encrypt.
///
/// # Returns
///
/// * `Ok(Vec<u8>)` containing the ciphertext with the nonce appended.
/// * `Err(CryptoError::InvalidKeyLength)` if the key length is invalid.
/// * `Err(CryptoError::EncryptionFailed)` if encryption fails.
pub fn encrypt_symmetric(key: &[u8], message: &[u8]) -> Result<Vec<u8>, CryptoError> {
    // Create cipher
    let cipher = ChaCha20Poly1305::new_from_slice(key)
        .map_err(|_| CryptoError::InvalidKeyLength)?;
    
    // Generate random nonce
    let mut nonce_bytes = [0u8; NONCE_SIZE];
    OsRng.fill_bytes(&mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);
    
    // Encrypt message
    let ciphertext = cipher.encrypt(nonce, message)
        .map_err(|e| CryptoError::EncryptionFailed(e.to_string()))?;
    
    // Append nonce to ciphertext
    let mut result = ciphertext;
    result.extend_from_slice(&nonce_bytes);
    
    Ok(result)
}

/// Decrypts data using ChaCha20Poly1305, extracting the nonce from the ciphertext.
///
/// # Arguments
///
/// * `key` - The decryption key (should be 32 bytes).
/// * `ciphertext_with_nonce` - The ciphertext with the nonce appended.
///
/// # Returns
///
/// * `Ok(Vec<u8>)` containing the decrypted message.
/// * `Err(CryptoError::InvalidKeyLength)` if the key length is invalid.
/// * `Err(CryptoError::DecryptionFailed)` if decryption fails or the ciphertext is too short.
pub fn decrypt_symmetric(key: &[u8], ciphertext_with_nonce: &[u8]) -> Result<Vec<u8>, CryptoError> {
    // Check if ciphertext is long enough to contain a nonce
    if ciphertext_with_nonce.len() <= NONCE_SIZE {
        return Err(CryptoError::DecryptionFailed("Ciphertext too short".to_string()));
    }
    
    // Extract nonce from the end of ciphertext
    let ciphertext_len = ciphertext_with_nonce.len() - NONCE_SIZE;
    let ciphertext = &ciphertext_with_nonce[0..ciphertext_len];
    let nonce_bytes = &ciphertext_with_nonce[ciphertext_len..];
    
    // Create cipher
    let cipher = ChaCha20Poly1305::new_from_slice(key)
        .map_err(|_| CryptoError::InvalidKeyLength)?;
    
    let nonce = Nonce::from_slice(nonce_bytes);
    
    // Decrypt message
    cipher.decrypt(nonce, ciphertext)
        .map_err(|e| CryptoError::DecryptionFailed(e.to_string()))
}

/// Encrypts data using a key directly (for internal use).
///
/// # Arguments
///
/// * `key` - The encryption key.
/// * `message` - The message to encrypt.
///
/// # Returns
///
/// * `Ok(Vec<u8>)` containing the ciphertext with the nonce appended.
/// * `Err(CryptoError)` if encryption fails.
pub fn encrypt_with_key(key: &[u8], message: &[u8]) -> Result<Vec<u8>, CryptoError> {
    encrypt_symmetric(key, message)
}

/// Decrypts data using a key directly (for internal use).
///
/// # Arguments
///
/// * `key` - The decryption key.
/// * `ciphertext_with_nonce` - The ciphertext with the nonce appended.
///
/// # Returns
///
/// * `Ok(Vec<u8>)` containing the decrypted message.
/// * `Err(CryptoError)` if decryption fails.
pub fn decrypt_with_key(key: &[u8], ciphertext_with_nonce: &[u8]) -> Result<Vec<u8>, CryptoError> {
    decrypt_symmetric(key, ciphertext_with_nonce)
}

/// Metadata for an encrypted key space.
#[derive(Serialize, Deserialize, Debug)]
pub struct EncryptedKeySpaceMetadata {
    pub name: String,
    pub created_at: u64,
    pub last_accessed: u64,
}

/// An encrypted key space with metadata.
#[derive(Serialize, Deserialize, Debug)]
pub struct EncryptedKeySpace {
    pub metadata: EncryptedKeySpaceMetadata,
    pub encrypted_data: Vec<u8>,
}

/// Encrypts a key space using a password.
///
/// # Arguments
///
/// * `space` - The key space to encrypt.
/// * `password` - The password to encrypt with.
///
/// # Returns
///
/// * `Ok(EncryptedKeySpace)` containing the encrypted key space.
/// * `Err(CryptoError)` if encryption fails.
pub fn encrypt_key_space(space: &KeySpace, password: &str) -> Result<EncryptedKeySpace, CryptoError> {
    // Serialize the key space
    let serialized = match serde_json::to_vec(space) {
        Ok(data) => data,
        Err(e) => {
            log::error!("Serialization error during encryption: {}", e);
            return Err(CryptoError::SerializationError(e.to_string()));
        }
    };
    
    // Derive key from password
    let key = derive_key_from_password(password);
    
    // Encrypt the serialized data
    let encrypted_data = encrypt_symmetric(&key, &serialized)?;
    
    // Create metadata
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64;
    let metadata = EncryptedKeySpaceMetadata {
        name: space.name.clone(),
        created_at: now,
        last_accessed: now,
    };
    
    Ok(EncryptedKeySpace {
        metadata,
        encrypted_data,
    })
}

/// Decrypts a key space using a password.
///
/// # Arguments
///
/// * `encrypted_space` - The encrypted key space.
/// * `password` - The password to decrypt with.
///
/// # Returns
///
/// * `Ok(KeySpace)` containing the decrypted key space.
/// * `Err(CryptoError)` if decryption fails.
pub fn decrypt_key_space(encrypted_space: &EncryptedKeySpace, password: &str) -> Result<KeySpace, CryptoError> {
    // Derive key from password
    let key = derive_key_from_password(password);
    
    // Decrypt the data
    let decrypted_data = decrypt_symmetric(&key, &encrypted_space.encrypted_data)?;
    
    // Deserialize the key space
    let space: KeySpace = match serde_json::from_slice(&decrypted_data) {
        Ok(space) => space,
        Err(e) => {
            log::error!("Deserialization error: {}", e);
            return Err(CryptoError::SerializationError(e.to_string()));
        }
    };
    
    Ok(space)
}

/// Serializes an encrypted key space to a JSON string.
///
/// # Arguments
///
/// * `encrypted_space` - The encrypted key space to serialize.
///
/// # Returns
///
/// * `Ok(String)` containing the serialized encrypted key space.
/// * `Err(CryptoError)` if serialization fails.
pub fn serialize_encrypted_space(encrypted_space: &EncryptedKeySpace) -> Result<String, CryptoError> {
    serde_json::to_string(encrypted_space)
        .map_err(|e| CryptoError::SerializationError(e.to_string()))
}

/// Deserializes an encrypted key space from a JSON string.
///
/// # Arguments
///
/// * `serialized` - The serialized encrypted key space.
///
/// # Returns
///
/// * `Ok(EncryptedKeySpace)` containing the deserialized encrypted key space.
/// * `Err(CryptoError)` if deserialization fails.
pub fn deserialize_encrypted_space(serialized: &str) -> Result<EncryptedKeySpace, CryptoError> {
    match serde_json::from_str(serialized) {
        Ok(space) => Ok(space),
        Err(e) => {
            log::error!("Error deserializing encrypted space: {}", e);
            Err(CryptoError::SerializationError(e.to_string()))
        }
    }
}