use aes_gcm::aead::AeadCore;
use aes_gcm::{
    aead::{Aead, KeyInit, OsRng},
    Aes256Gcm, Nonce,
};
use base64::Engine;

/// Load the 32-byte encryption key from PENTEST_ENCRYPTION_KEY env var.
/// Returns None if not set or invalid length.
pub fn load_encryption_key() -> Option<[u8; 32]> {
    let hex_key = std::env::var("PENTEST_ENCRYPTION_KEY").ok()?;
    let bytes = hex::decode(hex_key).ok()?;
    if bytes.len() != 32 {
        return None;
    }
    let mut key = [0u8; 32];
    key.copy_from_slice(&bytes);
    Some(key)
}

/// Encrypt a plaintext string. Returns base64-encoded nonce+ciphertext.
/// Returns the original string if no encryption key is available.
pub fn encrypt(plaintext: &str) -> String {
    let Some(key_bytes) = load_encryption_key() else {
        return plaintext.to_string();
    };
    let Ok(cipher) = Aes256Gcm::new_from_slice(&key_bytes) else {
        return plaintext.to_string();
    };
    let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
    let Ok(ciphertext) = cipher.encrypt(&nonce, plaintext.as_bytes()) else {
        return plaintext.to_string();
    };
    let mut combined = nonce.to_vec();
    combined.extend_from_slice(&ciphertext);
    base64::engine::general_purpose::STANDARD.encode(&combined)
}

/// Decrypt a base64-encoded nonce+ciphertext string.
/// Returns None if decryption fails.
pub fn decrypt(encrypted: &str) -> Option<String> {
    let key_bytes = load_encryption_key()?;
    let cipher = Aes256Gcm::new_from_slice(&key_bytes).ok()?;
    let combined = base64::engine::general_purpose::STANDARD
        .decode(encrypted)
        .ok()?;
    if combined.len() < 12 {
        return None;
    }
    let (nonce_bytes, ciphertext) = combined.split_at(12);
    let nonce = Nonce::from_slice(nonce_bytes);
    let plaintext = cipher.decrypt(nonce, ciphertext).ok()?;
    String::from_utf8(plaintext).ok()
}

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

    // Guard to serialize tests that touch env vars
    static ENV_LOCK: Mutex<()> = Mutex::new(());

    fn with_key<F: FnOnce()>(hex_key: &str, f: F) {
        let _guard = ENV_LOCK.lock();
        unsafe { std::env::set_var("PENTEST_ENCRYPTION_KEY", hex_key) };
        f();
        unsafe { std::env::remove_var("PENTEST_ENCRYPTION_KEY") };
    }

    #[test]
    fn round_trip() {
        let key = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
        with_key(key, || {
            let plaintext = "my_secret_password";
            let encrypted = encrypt(plaintext);
            assert_ne!(encrypted, plaintext);
            let decrypted = decrypt(&encrypted);
            assert_eq!(decrypted, Some(plaintext.to_string()));
        });
    }

    #[test]
    fn wrong_key_fails() {
        let _guard = ENV_LOCK.lock();
        let key1 = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
        let key2 = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
        let encrypted = {
            unsafe { std::env::set_var("PENTEST_ENCRYPTION_KEY", key1) };
            let e = encrypt("secret");
            unsafe { std::env::remove_var("PENTEST_ENCRYPTION_KEY") };
            e
        };
        unsafe { std::env::set_var("PENTEST_ENCRYPTION_KEY", key2) };
        assert!(decrypt(&encrypted).is_none());
        unsafe { std::env::remove_var("PENTEST_ENCRYPTION_KEY") };
    }

    #[test]
    fn no_key_passthrough() {
        let _guard = ENV_LOCK.lock();
        unsafe { std::env::remove_var("PENTEST_ENCRYPTION_KEY") };
        let result = encrypt("plain");
        assert_eq!(result, "plain");
    }

    #[test]
    fn corrupted_ciphertext() {
        let key = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
        with_key(key, || {
            assert!(decrypt("not-valid-base64!!!").is_none());
            // Valid base64 but wrong content
            let garbage = base64::engine::general_purpose::STANDARD.encode(b"tooshort");
            assert!(decrypt(&garbage).is_none());
        });
    }
}