//! AES-256-GCM Authenticated Encryption
//!
//! Provides AEAD (Authenticated Encryption with Associated Data) using
//! AES-256 in Galois/Counter Mode (GCM).
//!
//! Features:
//! - 256-bit key security
//! - 128-bit authentication tag
//! - Associated data authentication (metadata not encrypted but authenticated)
//! - Constant-time operations where possible

use aes_gcm::{
    aead::{Aead, KeyInit, Payload},
    Aes256Gcm, Key, Nonce,
};
use zeroize::Zeroizing;

use crate::error::{CryptoError, CryptoResult};

/// Key size for AES-256 (32 bytes)
pub const KEY_SIZE: usize = 32;

/// Nonce size for AES-GCM (12 bytes, 96 bits)
pub const NONCE_SIZE: usize = 12;

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

/// Encrypt data using AES-256-GCM
///
/// # Arguments
/// * `key` - 32-byte (256-bit) encryption key
/// * `nonce` - 12-byte (96-bit) nonce (MUST be unique per key)
/// * `plaintext` - Data to encrypt
/// * `aad` - Associated data to authenticate (but not encrypt)
///
/// # Returns
/// Ciphertext with appended 16-byte authentication tag
///
/// # Security
/// - NEVER reuse a nonce with the same key
/// - Use `crate::random_nonce()` to generate unique nonces
///
/// # Example
/// ```
/// use pqc_py::symmetric::aes_gcm::{encrypt_aes_gcm, decrypt_aes_gcm};
/// use pqc_py::{random_key, random_nonce};
///
/// let key = random_key();
/// let nonce = random_nonce();
/// let plaintext = b"Secret message for pqc-py";
/// let aad = b"pqc-v1";
///
/// let ciphertext = encrypt_aes_gcm(&key, &nonce, plaintext, aad).unwrap();
/// let decrypted = decrypt_aes_gcm(&key, &nonce, &ciphertext, aad).unwrap();
///
/// assert_eq!(decrypted, plaintext);
/// ```
pub fn encrypt_aes_gcm(
    key: &[u8; KEY_SIZE],
    nonce: &[u8; NONCE_SIZE],
    plaintext: &[u8],
    aad: &[u8],
) -> CryptoResult<Vec<u8>> {
    let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(key));
    let nonce = Nonce::from_slice(nonce);

    let payload = Payload {
        msg: plaintext,
        aad,
    };

    cipher
        .encrypt(nonce, payload)
        .map_err(|_| CryptoError::EncryptionFailed("AES-GCM encryption failed".to_string()))
}

/// Decrypt data using AES-256-GCM
///
/// # Arguments
/// * `key` - 32-byte (256-bit) encryption key (same as used for encryption)
/// * `nonce` - 12-byte (96-bit) nonce (same as used for encryption)
/// * `ciphertext` - Encrypted data with appended authentication tag
/// * `aad` - Associated data (same as used for encryption)
///
/// # Returns
/// Decrypted plaintext, or error if authentication fails
///
/// # Security
/// Returns an error if:
/// - The authentication tag doesn't match (tampered data)
/// - The AAD doesn't match what was used during encryption
/// - The ciphertext is too short to contain a valid tag
pub fn decrypt_aes_gcm(
    key: &[u8; KEY_SIZE],
    nonce: &[u8; NONCE_SIZE],
    ciphertext: &[u8],
    aad: &[u8],
) -> CryptoResult<Vec<u8>> {
    if ciphertext.len() < TAG_SIZE {
        return Err(CryptoError::InvalidCiphertext(
            "Ciphertext too short to contain authentication tag".to_string()
        ));
    }

    let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(key));
    let nonce = Nonce::from_slice(nonce);

    let payload = Payload {
        msg: ciphertext,
        aad,
    };

    cipher
        .decrypt(nonce, payload)
        .map_err(|_| CryptoError::DecryptionFailed)
}

/// Encrypt with automatic nonce generation
///
/// Generates a random 12-byte nonce and prepends it to the ciphertext.
/// The output format is: `[12-byte nonce][ciphertext][16-byte tag]`
///
/// # Arguments
/// * `key` - 32-byte encryption key
/// * `plaintext` - Data to encrypt
/// * `aad` - Associated data to authenticate
///
/// # Returns
/// Nonce + ciphertext + tag combined
pub fn encrypt_aes_gcm_with_nonce(
    key: &[u8; KEY_SIZE],
    plaintext: &[u8],
    aad: &[u8],
) -> CryptoResult<Vec<u8>> {
    let nonce = crate::random_nonce();
    let ciphertext = encrypt_aes_gcm(key, &nonce, plaintext, aad)?;

    let mut result = Vec::with_capacity(NONCE_SIZE + ciphertext.len());
    result.extend_from_slice(&nonce);
    result.extend_from_slice(&ciphertext);
    Ok(result)
}

/// Decrypt data encrypted with `encrypt_aes_gcm_with_nonce`
///
/// Extracts the prepended nonce and decrypts the ciphertext.
///
/// # Arguments
/// * `key` - 32-byte encryption key
/// * `data` - Combined nonce + ciphertext + tag from `encrypt_aes_gcm_with_nonce`
/// * `aad` - Associated data (same as used for encryption)
///
/// # Returns
/// Decrypted plaintext
pub fn decrypt_aes_gcm_with_nonce(
    key: &[u8; KEY_SIZE],
    data: &[u8],
    aad: &[u8],
) -> CryptoResult<Vec<u8>> {
    if data.len() < NONCE_SIZE + TAG_SIZE {
        return Err(CryptoError::InvalidCiphertext(
            "Data too short to contain nonce and tag".to_string()
        ));
    }

    let nonce: [u8; NONCE_SIZE] = data[..NONCE_SIZE]
        .try_into()
        .expect("Slice length verified above");
    let ciphertext = &data[NONCE_SIZE..];

    decrypt_aes_gcm(key, &nonce, ciphertext, aad)
}

/// Encrypt data with key and nonce provided as slices
///
/// Validates key and nonce lengths before encrypting.
pub fn encrypt_aes_gcm_slice(
    key: &[u8],
    nonce: &[u8],
    plaintext: &[u8],
    aad: &[u8],
) -> CryptoResult<Vec<u8>> {
    if key.len() != KEY_SIZE {
        return Err(CryptoError::InvalidKeyLength {
            expected: KEY_SIZE,
            actual: key.len(),
        });
    }
    if nonce.len() != NONCE_SIZE {
        return Err(CryptoError::InvalidNonceLength {
            expected: NONCE_SIZE,
            actual: nonce.len(),
        });
    }

    let key: [u8; KEY_SIZE] = key.try_into().unwrap();
    let nonce: [u8; NONCE_SIZE] = nonce.try_into().unwrap();

    encrypt_aes_gcm(&key, &nonce, plaintext, aad)
}

/// Decrypt data with key and nonce provided as slices
///
/// Validates key and nonce lengths before decrypting.
pub fn decrypt_aes_gcm_slice(
    key: &[u8],
    nonce: &[u8],
    ciphertext: &[u8],
    aad: &[u8],
) -> CryptoResult<Vec<u8>> {
    if key.len() != KEY_SIZE {
        return Err(CryptoError::InvalidKeyLength {
            expected: KEY_SIZE,
            actual: key.len(),
        });
    }
    if nonce.len() != NONCE_SIZE {
        return Err(CryptoError::InvalidNonceLength {
            expected: NONCE_SIZE,
            actual: nonce.len(),
        });
    }

    let key: [u8; KEY_SIZE] = key.try_into().unwrap();
    let nonce: [u8; NONCE_SIZE] = nonce.try_into().unwrap();

    decrypt_aes_gcm(&key, &nonce, ciphertext, aad)
}

/// Encrypt a covert channel message
///
/// Uses a standardized AAD format for protocol messages.
///
/// # Arguments
/// * `key` - Encryption key derived for this channel
/// * `channel` - Channel name (e.g., "TCOE")
/// * `message` - Message payload to encrypt
///
/// # Returns
/// Nonce + ciphertext + tag
pub fn encrypt_phantom_message(
    key: &[u8; KEY_SIZE],
    channel: &str,
    message: &[u8],
) -> CryptoResult<Vec<u8>> {
    let aad = format!("phantom-protocol:{}", channel);
    encrypt_aes_gcm_with_nonce(key, message, aad.as_bytes())
}

/// Decrypt a covert channel message
///
/// # Arguments
/// * `key` - Encryption key derived for this channel
/// * `channel` - Channel name (must match encryption)
/// * `encrypted` - Output from `encrypt_phantom_message`
pub fn decrypt_phantom_message(
    key: &[u8; KEY_SIZE],
    channel: &str,
    encrypted: &[u8],
) -> CryptoResult<Vec<u8>> {
    let aad = format!("phantom-protocol:{}", channel);
    decrypt_aes_gcm_with_nonce(key, encrypted, aad.as_bytes())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{random_key, random_nonce};

    #[test]
    fn test_encrypt_decrypt_basic() {
        let key = random_key();
        let nonce = random_nonce();
        let plaintext = b"Hello, pqc-py!";
        let aad = b"test-aad";

        let ciphertext = encrypt_aes_gcm(&key, &nonce, plaintext, aad).unwrap();
        let decrypted = decrypt_aes_gcm(&key, &nonce, &ciphertext, aad).unwrap();

        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_ciphertext_length() {
        let key = random_key();
        let nonce = random_nonce();
        let plaintext = b"Test message";
        let aad = b"";

        let ciphertext = encrypt_aes_gcm(&key, &nonce, plaintext, aad).unwrap();

        // Ciphertext = plaintext + 16-byte tag
        assert_eq!(ciphertext.len(), plaintext.len() + TAG_SIZE);
    }

    #[test]
    fn test_empty_plaintext() {
        let key = random_key();
        let nonce = random_nonce();
        let plaintext = b"";
        let aad = b"some-aad";

        let ciphertext = encrypt_aes_gcm(&key, &nonce, plaintext, aad).unwrap();
        let decrypted = decrypt_aes_gcm(&key, &nonce, &ciphertext, aad).unwrap();

        assert_eq!(decrypted.len(), 0);
    }

    #[test]
    fn test_empty_aad() {
        let key = random_key();
        let nonce = random_nonce();
        let plaintext = b"Test message";
        let aad = b"";

        let ciphertext = encrypt_aes_gcm(&key, &nonce, plaintext, aad).unwrap();
        let decrypted = decrypt_aes_gcm(&key, &nonce, &ciphertext, aad).unwrap();

        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_large_plaintext() {
        let key = random_key();
        let nonce = random_nonce();
        let plaintext = vec![0xABu8; 1_000_000]; // 1 MB
        let aad = b"large-file";

        let ciphertext = encrypt_aes_gcm(&key, &nonce, &plaintext, aad).unwrap();
        let decrypted = decrypt_aes_gcm(&key, &nonce, &ciphertext, aad).unwrap();

        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_wrong_key() {
        let key1 = random_key();
        let key2 = random_key();
        let nonce = random_nonce();
        let plaintext = b"Secret message";
        let aad = b"";

        let ciphertext = encrypt_aes_gcm(&key1, &nonce, plaintext, aad).unwrap();
        let result = decrypt_aes_gcm(&key2, &nonce, &ciphertext, aad);

        assert!(matches!(result, Err(CryptoError::DecryptionFailed)));
    }

    #[test]
    fn test_wrong_nonce() {
        let key = random_key();
        let nonce1 = random_nonce();
        let nonce2 = random_nonce();
        let plaintext = b"Secret message";
        let aad = b"";

        let ciphertext = encrypt_aes_gcm(&key, &nonce1, plaintext, aad).unwrap();
        let result = decrypt_aes_gcm(&key, &nonce2, &ciphertext, aad);

        assert!(matches!(result, Err(CryptoError::DecryptionFailed)));
    }

    #[test]
    fn test_wrong_aad() {
        let key = random_key();
        let nonce = random_nonce();
        let plaintext = b"Secret message";

        let ciphertext = encrypt_aes_gcm(&key, &nonce, plaintext, b"aad1").unwrap();
        let result = decrypt_aes_gcm(&key, &nonce, &ciphertext, b"aad2");

        assert!(matches!(result, Err(CryptoError::DecryptionFailed)));
    }

    #[test]
    fn test_tampered_ciphertext() {
        let key = random_key();
        let nonce = random_nonce();
        let plaintext = b"Secret message";
        let aad = b"";

        let mut ciphertext = encrypt_aes_gcm(&key, &nonce, plaintext, aad).unwrap();
        ciphertext[0] ^= 0xFF; // Flip bits in first byte

        let result = decrypt_aes_gcm(&key, &nonce, &ciphertext, aad);
        assert!(matches!(result, Err(CryptoError::DecryptionFailed)));
    }

    #[test]
    fn test_tampered_tag() {
        let key = random_key();
        let nonce = random_nonce();
        let plaintext = b"Secret message";
        let aad = b"";

        let mut ciphertext = encrypt_aes_gcm(&key, &nonce, plaintext, aad).unwrap();
        let len = ciphertext.len();
        ciphertext[len - 1] ^= 0xFF; // Flip bits in last byte (tag)

        let result = decrypt_aes_gcm(&key, &nonce, &ciphertext, aad);
        assert!(matches!(result, Err(CryptoError::DecryptionFailed)));
    }

    #[test]
    fn test_truncated_ciphertext() {
        let key = random_key();
        let nonce = random_nonce();

        let short_ciphertext = vec![0u8; TAG_SIZE - 1];
        let result = decrypt_aes_gcm(&key, &nonce, &short_ciphertext, b"");

        assert!(matches!(result, Err(CryptoError::InvalidCiphertext(_))));
    }

    #[test]
    fn test_encrypt_with_nonce() {
        let key = random_key();
        let plaintext = b"Test message";
        let aad = b"test-aad";

        let encrypted = encrypt_aes_gcm_with_nonce(&key, plaintext, aad).unwrap();
        let decrypted = decrypt_aes_gcm_with_nonce(&key, &encrypted, aad).unwrap();

        assert_eq!(decrypted, plaintext);

        // Verify format: nonce (12) + ciphertext (12) + tag (16) = 40
        assert_eq!(encrypted.len(), NONCE_SIZE + plaintext.len() + TAG_SIZE);
    }

    #[test]
    fn test_with_nonce_different_encryptions() {
        let key = random_key();
        let plaintext = b"Test message";
        let aad = b"test-aad";

        let enc1 = encrypt_aes_gcm_with_nonce(&key, plaintext, aad).unwrap();
        let enc2 = encrypt_aes_gcm_with_nonce(&key, plaintext, aad).unwrap();

        // Should produce different ciphertexts (different random nonces)
        assert_ne!(enc1, enc2);

        // But both should decrypt to same plaintext
        let dec1 = decrypt_aes_gcm_with_nonce(&key, &enc1, aad).unwrap();
        let dec2 = decrypt_aes_gcm_with_nonce(&key, &enc2, aad).unwrap();

        assert_eq!(dec1, plaintext);
        assert_eq!(dec2, plaintext);
    }

    #[test]
    fn test_slice_functions() {
        let key = random_key();
        let nonce = random_nonce();
        let plaintext = b"Test message";
        let aad = b"test";

        let ciphertext = encrypt_aes_gcm_slice(&key, &nonce, plaintext, aad).unwrap();
        let decrypted = decrypt_aes_gcm_slice(&key, &nonce, &ciphertext, aad).unwrap();

        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_slice_wrong_key_length() {
        let short_key = [0u8; 16]; // AES-128 key, not AES-256
        let nonce = random_nonce();

        let result = encrypt_aes_gcm_slice(&short_key, &nonce, b"test", b"");
        assert!(matches!(result, Err(CryptoError::InvalidKeyLength { expected: 32, actual: 16 })));
    }

    #[test]
    fn test_slice_wrong_nonce_length() {
        let key = random_key();
        let short_nonce = [0u8; 8];

        let result = encrypt_aes_gcm_slice(&key, &short_nonce, b"test", b"");
        assert!(matches!(result, Err(CryptoError::InvalidNonceLength { expected: 12, actual: 8 })));
    }

    #[test]
    fn test_phantom_message() {
        let key = random_key();
        let message = b"Covert payload for TCOE channel";

        let encrypted = encrypt_phantom_message(&key, "TCOE", message).unwrap();
        let decrypted = decrypt_phantom_message(&key, "TCOE", &encrypted).unwrap();

        assert_eq!(decrypted, message);
    }

    #[test]
    fn test_phantom_message_wrong_channel() {
        let key = random_key();
        let message = b"Covert payload";

        let encrypted = encrypt_phantom_message(&key, "TCOE", message).unwrap();
        let result = decrypt_phantom_message(&key, "CoTSE", &encrypted);

        // Wrong channel = wrong AAD = decryption failure
        assert!(matches!(result, Err(CryptoError::DecryptionFailed)));
    }

    #[test]
    fn test_deterministic_with_same_nonce() {
        let key = random_key();
        let nonce = random_nonce();
        let plaintext = b"Test message";
        let aad = b"";

        let ct1 = encrypt_aes_gcm(&key, &nonce, plaintext, aad).unwrap();
        let ct2 = encrypt_aes_gcm(&key, &nonce, plaintext, aad).unwrap();

        // Same key + nonce + plaintext + aad = same ciphertext
        assert_eq!(ct1, ct2);
    }

    // NIST GCM Test Vector (from NIST SP 800-38D)
    #[test]
    fn test_nist_gcm_vector() {
        // Test Case 16 from NIST
        let key = hex::decode("feffe9928665731c6d6a8f9467308308feffe9928665731c6d6a8f9467308308")
            .unwrap();
        let nonce = hex::decode("cafebabefacedbaddecaf888").unwrap();
        let plaintext = hex::decode(
            "d9313225f88406e5a55909c5aff5269a86a7a9531534f7da2e4c303d8a318a721c3c0c95956809532fcf0e2449a6b525b16aedf5aa0de657ba637b39"
        ).unwrap();
        let aad = hex::decode(
            "feedfacedeadbeeffeedfacedeadbeefabaddad2"
        ).unwrap();

        let expected_ciphertext = hex::decode(
            "522dc1f099567d07f47f37a32a84427d643a8cdcbfe5c0c97598a2bd2555d1aa8cb08e48590dbb3da7b08b1056828838c5f61e6393ba7a0abcc9f662"
        ).unwrap();
        let expected_tag = hex::decode("76fc6ece0f4e1768cddf8853bb2d551b").unwrap();

        let key_arr: [u8; 32] = key.try_into().unwrap();
        let nonce_arr: [u8; 12] = nonce.try_into().unwrap();

        let ciphertext = encrypt_aes_gcm(&key_arr, &nonce_arr, &plaintext, &aad).unwrap();

        // Compare ciphertext (without tag)
        assert_eq!(&ciphertext[..expected_ciphertext.len()], expected_ciphertext.as_slice());

        // Compare tag
        assert_eq!(&ciphertext[expected_ciphertext.len()..], expected_tag.as_slice());

        // Verify decryption works
        let decrypted = decrypt_aes_gcm(&key_arr, &nonce_arr, &ciphertext, &aad).unwrap();
        assert_eq!(decrypted, plaintext);
    }
}
