//! ML-KEM (Kyber) Key Encapsulation Mechanism
//!
//! NIST FIPS 203 implementation for post-quantum key encapsulation.
//! Provides shared secret establishment resistant to quantum attacks.
//!
//! Supported parameter sets:
//! - ML-KEM-512 (Level 1, 128-bit security)
//! - ML-KEM-768 (Level 3, 192-bit security) - Default
//! - ML-KEM-1024 (Level 5, 256-bit security)

use pqcrypto_kyber::{
    kyber512, kyber768, kyber1024,
};
use pqcrypto_traits::kem::{PublicKey as KemPublicKey, SecretKey as KemSecretKey, SharedSecret, Ciphertext};
use zeroize::{Zeroize, ZeroizeOnDrop};

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

/// Kyber key pair
#[derive(ZeroizeOnDrop)]
pub struct KyberKeyPair {
    /// Public key for encapsulation
    #[zeroize(skip)]
    pub public_key: Vec<u8>,
    /// Secret key for decapsulation (zeroed on drop)
    secret_key: Vec<u8>,
    /// Security level
    #[zeroize(skip)]
    pub level: SecurityLevel,
}

impl KyberKeyPair {
    /// Get the secret key bytes
    pub fn secret_key(&self) -> &[u8] {
        &self.secret_key
    }

    /// Get public key size in bytes
    pub fn public_key_size(&self) -> usize {
        self.public_key.len()
    }

    /// Get secret key size in bytes
    pub fn secret_key_size(&self) -> usize {
        self.secret_key.len()
    }
}

/// Generate a Kyber key pair
///
/// # Arguments
/// * `level` - Security level (Level1=512, Level3=768, Level5=1024)
///
/// # Returns
/// A new Kyber key pair
///
/// # Example
/// ```
/// use pqc_py::pqc::kyber::kyber_keygen;
/// use pqc_py::pqc::SecurityLevel;
///
/// let keypair = kyber_keygen(SecurityLevel::Level3).unwrap();
/// println!("Public key size: {} bytes", keypair.public_key.len());
/// ```
pub fn kyber_keygen(level: SecurityLevel) -> CryptoResult<KyberKeyPair> {
    match level {
        SecurityLevel::Level1 => {
            let (pk, sk) = kyber512::keypair();
            Ok(KyberKeyPair {
                public_key: pk.as_bytes().to_vec(),
                secret_key: sk.as_bytes().to_vec(),
                level,
            })
        }
        SecurityLevel::Level3 => {
            let (pk, sk) = kyber768::keypair();
            Ok(KyberKeyPair {
                public_key: pk.as_bytes().to_vec(),
                secret_key: sk.as_bytes().to_vec(),
                level,
            })
        }
        SecurityLevel::Level5 => {
            let (pk, sk) = kyber1024::keypair();
            Ok(KyberKeyPair {
                public_key: pk.as_bytes().to_vec(),
                secret_key: sk.as_bytes().to_vec(),
                level,
            })
        }
    }
}

/// Encapsulate a shared secret using a Kyber public key
///
/// # Arguments
/// * `public_key` - Recipient's public key
///
/// # Returns
/// Tuple of (ciphertext, shared_secret)
/// - ciphertext: Send to recipient
/// - shared_secret: 32-byte shared key
pub fn kyber_encapsulate(public_key: &[u8]) -> CryptoResult<(Vec<u8>, [u8; 32])> {
    // Detect key size to determine level
    match public_key.len() {
        800 => {
            let pk = kyber512::PublicKey::from_bytes(public_key)
                .map_err(|_| CryptoError::InvalidPublicKey("Invalid Kyber-512 public key".to_string()))?;
            let (ss, ct) = kyber512::encapsulate(&pk);
            let mut shared = [0u8; 32];
            shared.copy_from_slice(ss.as_bytes());
            Ok((ct.as_bytes().to_vec(), shared))
        }
        1184 => {
            let pk = kyber768::PublicKey::from_bytes(public_key)
                .map_err(|_| CryptoError::InvalidPublicKey("Invalid Kyber-768 public key".to_string()))?;
            let (ss, ct) = kyber768::encapsulate(&pk);
            let mut shared = [0u8; 32];
            shared.copy_from_slice(ss.as_bytes());
            Ok((ct.as_bytes().to_vec(), shared))
        }
        1568 => {
            let pk = kyber1024::PublicKey::from_bytes(public_key)
                .map_err(|_| CryptoError::InvalidPublicKey("Invalid Kyber-1024 public key".to_string()))?;
            let (ss, ct) = kyber1024::encapsulate(&pk);
            let mut shared = [0u8; 32];
            shared.copy_from_slice(ss.as_bytes());
            Ok((ct.as_bytes().to_vec(), shared))
        }
        _ => Err(CryptoError::InvalidPublicKey(
            format!("Unknown Kyber public key size: {} bytes", public_key.len())
        )),
    }
}

/// Decapsulate a shared secret using a Kyber secret key
///
/// # Arguments
/// * `ciphertext` - Ciphertext from encapsulation
/// * `secret_key` - Recipient's secret key
///
/// # Returns
/// 32-byte shared secret
pub fn kyber_decapsulate(ciphertext: &[u8], secret_key: &[u8]) -> CryptoResult<[u8; 32]> {
    // Detect key size to determine level
    match secret_key.len() {
        1632 => {
            let sk = kyber512::SecretKey::from_bytes(secret_key)
                .map_err(|_| CryptoError::InvalidSecretKey("Invalid Kyber-512 secret key".to_string()))?;
            let ct = kyber512::Ciphertext::from_bytes(ciphertext)
                .map_err(|_| CryptoError::InvalidCiphertext("Invalid Kyber-512 ciphertext".to_string()))?;
            let ss = kyber512::decapsulate(&ct, &sk);
            let mut shared = [0u8; 32];
            shared.copy_from_slice(ss.as_bytes());
            Ok(shared)
        }
        2400 => {
            let sk = kyber768::SecretKey::from_bytes(secret_key)
                .map_err(|_| CryptoError::InvalidSecretKey("Invalid Kyber-768 secret key".to_string()))?;
            let ct = kyber768::Ciphertext::from_bytes(ciphertext)
                .map_err(|_| CryptoError::InvalidCiphertext("Invalid Kyber-768 ciphertext".to_string()))?;
            let ss = kyber768::decapsulate(&ct, &sk);
            let mut shared = [0u8; 32];
            shared.copy_from_slice(ss.as_bytes());
            Ok(shared)
        }
        3168 => {
            let sk = kyber1024::SecretKey::from_bytes(secret_key)
                .map_err(|_| CryptoError::InvalidSecretKey("Invalid Kyber-1024 secret key".to_string()))?;
            let ct = kyber1024::Ciphertext::from_bytes(ciphertext)
                .map_err(|_| CryptoError::InvalidCiphertext("Invalid Kyber-1024 ciphertext".to_string()))?;
            let ss = kyber1024::decapsulate(&ct, &sk);
            let mut shared = [0u8; 32];
            shared.copy_from_slice(ss.as_bytes());
            Ok(shared)
        }
        _ => Err(CryptoError::InvalidSecretKey(
            format!("Unknown Kyber secret key size: {} bytes", secret_key.len())
        )),
    }
}

/// Key sizes for each Kyber variant
pub mod sizes {
    /// Kyber-512 (Level 1)
    pub mod kyber512 {
        pub const PUBLIC_KEY: usize = 800;
        pub const SECRET_KEY: usize = 1632;
        pub const CIPHERTEXT: usize = 768;
        pub const SHARED_SECRET: usize = 32;
    }

    /// Kyber-768 (Level 3)
    pub mod kyber768 {
        pub const PUBLIC_KEY: usize = 1184;
        pub const SECRET_KEY: usize = 2400;
        pub const CIPHERTEXT: usize = 1088;
        pub const SHARED_SECRET: usize = 32;
    }

    /// Kyber-1024 (Level 5)
    pub mod kyber1024 {
        pub const PUBLIC_KEY: usize = 1568;
        pub const SECRET_KEY: usize = 3168;
        pub const CIPHERTEXT: usize = 1568;
        pub const SHARED_SECRET: usize = 32;
    }
}

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

    #[test]
    fn test_kyber512_keygen() {
        let keypair = kyber_keygen(SecurityLevel::Level1).unwrap();
        assert_eq!(keypair.public_key.len(), sizes::kyber512::PUBLIC_KEY);
        assert_eq!(keypair.secret_key.len(), sizes::kyber512::SECRET_KEY);
    }

    #[test]
    fn test_kyber768_keygen() {
        let keypair = kyber_keygen(SecurityLevel::Level3).unwrap();
        assert_eq!(keypair.public_key.len(), sizes::kyber768::PUBLIC_KEY);
        assert_eq!(keypair.secret_key.len(), sizes::kyber768::SECRET_KEY);
    }

    #[test]
    fn test_kyber1024_keygen() {
        let keypair = kyber_keygen(SecurityLevel::Level5).unwrap();
        assert_eq!(keypair.public_key.len(), sizes::kyber1024::PUBLIC_KEY);
        assert_eq!(keypair.secret_key.len(), sizes::kyber1024::SECRET_KEY);
    }

    #[test]
    fn test_kyber768_encap_decap() {
        let keypair = kyber_keygen(SecurityLevel::Level3).unwrap();

        let (ciphertext, shared1) = kyber_encapsulate(&keypair.public_key).unwrap();
        let shared2 = kyber_decapsulate(&ciphertext, keypair.secret_key()).unwrap();

        assert_eq!(shared1, shared2);
        assert_eq!(shared1.len(), 32);
    }

    #[test]
    fn test_kyber512_encap_decap() {
        let keypair = kyber_keygen(SecurityLevel::Level1).unwrap();

        let (ciphertext, shared1) = kyber_encapsulate(&keypair.public_key).unwrap();
        let shared2 = kyber_decapsulate(&ciphertext, keypair.secret_key()).unwrap();

        assert_eq!(shared1, shared2);
    }

    #[test]
    fn test_kyber1024_encap_decap() {
        let keypair = kyber_keygen(SecurityLevel::Level5).unwrap();

        let (ciphertext, shared1) = kyber_encapsulate(&keypair.public_key).unwrap();
        let shared2 = kyber_decapsulate(&ciphertext, keypair.secret_key()).unwrap();

        assert_eq!(shared1, shared2);
    }

    #[test]
    fn test_different_keypairs_different_secrets() {
        let keypair1 = kyber_keygen(SecurityLevel::Level3).unwrap();
        let keypair2 = kyber_keygen(SecurityLevel::Level3).unwrap();

        let (_, shared1) = kyber_encapsulate(&keypair1.public_key).unwrap();
        let (_, shared2) = kyber_encapsulate(&keypair2.public_key).unwrap();

        assert_ne!(shared1, shared2);
    }

    #[test]
    fn test_invalid_public_key() {
        let result = kyber_encapsulate(&[0u8; 100]);
        assert!(matches!(result, Err(CryptoError::InvalidPublicKey(_))));
    }

    #[test]
    fn test_invalid_secret_key() {
        let result = kyber_decapsulate(&[0u8; 1088], &[0u8; 100]);
        assert!(matches!(result, Err(CryptoError::InvalidSecretKey(_))));
    }
}
