use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use anyhow::{Context, Result};
use base64::{engine::general_purpose, Engine as _};

type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;

/// Encrypt a message using AES-256-CBC, compatible with CryptoJS AES.encrypt
pub fn encrypt(message: &str, key: &str) -> Result<String> {
    if message.is_empty() || key.is_empty() {
        anyhow::bail!("Message and key cannot be empty");
    }

    // Generate random salt
    let salt: [u8; 8] = [0; 8]; // CryptoJS uses random salt, but for simplicity we use zero

    // CryptoJS uses EVP_BytesToKey with MD5 for key derivation
    let (key_bytes, iv) = derive_key_and_iv(key.as_bytes(), &salt)?;

    // Prepare buffer for encryption
    let msg_bytes = message.as_bytes();
    let mut buffer = vec![0u8; msg_bytes.len() + 16]; // Extra space for padding
    buffer[..msg_bytes.len()].copy_from_slice(msg_bytes);

    // Encrypt the message
    let ciphertext = Aes256CbcEnc::new(&key_bytes.into(), &iv.into())
        .encrypt_padded_mut::<Pkcs7>(&mut buffer, msg_bytes.len())
        .map_err(|_| anyhow::anyhow!("Encryption failed"))?;

    // CryptoJS format: "Salted__" + salt + ciphertext, then base64
    let mut result = Vec::new();
    result.extend_from_slice(b"Salted__");
    result.extend_from_slice(&salt);
    result.extend_from_slice(ciphertext);

    Ok(general_purpose::STANDARD.encode(&result))
}

/// Decrypt a message encrypted with CryptoJS AES.encrypt
pub fn decrypt(encrypted_message: &str, key: &str) -> Result<String> {
    if encrypted_message.is_empty() || key.is_empty() {
        anyhow::bail!("Encrypted message and key cannot be empty");
    }

    // Decode base64
    let encrypted_data = general_purpose::STANDARD
        .decode(encrypted_message)
        .context("Failed to decode base64")?;

    // Check for "Salted__" prefix
    if encrypted_data.len() < 16 || &encrypted_data[0..8] != b"Salted__" {
        anyhow::bail!("Invalid encrypted data format");
    }

    // Extract salt (8 bytes after "Salted__")
    let salt = &encrypted_data[8..16];
    let ciphertext = &encrypted_data[16..];

    // Derive key and IV
    let (key_bytes, iv) = derive_key_and_iv(key.as_bytes(), salt)?;

    // Decrypt
    let mut buffer = ciphertext.to_vec();
    let plaintext = Aes256CbcDec::new(&key_bytes.into(), &iv.into())
        .decrypt_padded_mut::<Pkcs7>(&mut buffer)
        .map_err(|_| anyhow::anyhow!("Decryption failed - wrong master key?"))?;

    String::from_utf8(plaintext.to_vec()).context("Decrypted data is not valid UTF-8")
}

/// Derive key and IV using EVP_BytesToKey with MD5 (compatible with CryptoJS)
/// This implements the OpenSSL EVP_BytesToKey function with MD5
fn derive_key_and_iv(password: &[u8], salt: &[u8]) -> Result<([u8; 32], [u8; 16])> {
    let mut key = [0u8; 32];
    let mut iv = [0u8; 16];
    let mut derived = Vec::new();

    // First round: MD5(password + salt)
    let mut data = Vec::new();
    data.extend_from_slice(password);
    data.extend_from_slice(salt);
    let hash = md5::compute(&data);
    derived.extend_from_slice(&hash.0);

    // Second round: MD5(previous_hash + password + salt)
    let mut data = Vec::new();
    data.extend_from_slice(&hash.0);
    data.extend_from_slice(password);
    data.extend_from_slice(salt);
    let hash = md5::compute(&data);
    derived.extend_from_slice(&hash.0);

    // Third round: MD5(previous_hash + password + salt)
    let mut data = Vec::new();
    data.extend_from_slice(&hash.0);
    data.extend_from_slice(password);
    data.extend_from_slice(salt);
    let hash = md5::compute(&data);
    derived.extend_from_slice(&hash.0);

    // Extract key (32 bytes) and IV (16 bytes)
    key.copy_from_slice(&derived[0..32]);
    iv.copy_from_slice(&derived[32..48]);

    Ok((key, iv))
}

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

    #[test]
    fn test_encrypt_decrypt() {
        let message = "Hello, World!";
        let key = "my-secret-key";

        let encrypted = encrypt(message, key).unwrap();
        let decrypted = decrypt(&encrypted, key).unwrap();

        assert_eq!(message, decrypted);
    }

    #[test]
    fn test_decrypt_with_wrong_key() {
        let message = "Hello, World!";
        let key = "my-secret-key";
        let wrong_key = "wrong-key";

        let encrypted = encrypt(message, key).unwrap();
        let result = decrypt(&encrypted, wrong_key);

        assert!(result.is_err());
    }
}
