//! ML-DSA (Dilithium) Digital Signatures
//!
//! NIST FIPS 204 implementation for post-quantum digital signatures.
//! Provides fast signing and verification resistant to quantum attacks.
//!
//! Supported parameter sets:
//! - ML-DSA-44 (Level 2, ~128-bit security)
//! - ML-DSA-65 (Level 3, ~192-bit security) - Default
//! - ML-DSA-87 (Level 5, ~256-bit security)

use pqcrypto_dilithium::{
    dilithium2, dilithium3, dilithium5,
};
use pqcrypto_traits::sign::{PublicKey as SignPublicKey, SecretKey as SignSecretKey, SignedMessage, DetachedSignature};
use zeroize::{Zeroize, ZeroizeOnDrop};

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

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

impl DilithiumKeyPair {
    /// 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 Dilithium key pair
///
/// # Arguments
/// * `level` - Security level (Level1=2, Level3=3, Level5=5)
///
/// # Returns
/// A new Dilithium key pair
///
/// # Example
/// ```
/// use pqc_py::pqc::dilithium::dilithium_keygen;
/// use pqc_py::pqc::SecurityLevel;
///
/// let keypair = dilithium_keygen(SecurityLevel::Level3).unwrap();
/// println!("Public key size: {} bytes", keypair.public_key.len());
/// ```
pub fn dilithium_keygen(level: SecurityLevel) -> CryptoResult<DilithiumKeyPair> {
    match level {
        SecurityLevel::Level1 => {
            let (pk, sk) = dilithium2::keypair();
            Ok(DilithiumKeyPair {
                public_key: pk.as_bytes().to_vec(),
                secret_key: sk.as_bytes().to_vec(),
                level,
            })
        }
        SecurityLevel::Level3 => {
            let (pk, sk) = dilithium3::keypair();
            Ok(DilithiumKeyPair {
                public_key: pk.as_bytes().to_vec(),
                secret_key: sk.as_bytes().to_vec(),
                level,
            })
        }
        SecurityLevel::Level5 => {
            let (pk, sk) = dilithium5::keypair();
            Ok(DilithiumKeyPair {
                public_key: pk.as_bytes().to_vec(),
                secret_key: sk.as_bytes().to_vec(),
                level,
            })
        }
    }
}

/// Sign a message using a Dilithium secret key
///
/// # Arguments
/// * `secret_key` - Signer's secret key
/// * `message` - Message to sign
///
/// # Returns
/// Detached signature
pub fn dilithium_sign(secret_key: &[u8], message: &[u8]) -> CryptoResult<Vec<u8>> {
    // Detect key size to determine level
    match secret_key.len() {
        2560 => {
            let sk = dilithium2::SecretKey::from_bytes(secret_key)
                .map_err(|_| CryptoError::InvalidSecretKey("Invalid Dilithium2 secret key".to_string()))?;
            let sig = dilithium2::detached_sign(message, &sk);
            Ok(sig.as_bytes().to_vec())
        }
        4032 => {
            let sk = dilithium3::SecretKey::from_bytes(secret_key)
                .map_err(|_| CryptoError::InvalidSecretKey("Invalid Dilithium3 secret key".to_string()))?;
            let sig = dilithium3::detached_sign(message, &sk);
            Ok(sig.as_bytes().to_vec())
        }
        4896 => {
            let sk = dilithium5::SecretKey::from_bytes(secret_key)
                .map_err(|_| CryptoError::InvalidSecretKey("Invalid Dilithium5 secret key".to_string()))?;
            let sig = dilithium5::detached_sign(message, &sk);
            Ok(sig.as_bytes().to_vec())
        }
        _ => Err(CryptoError::InvalidSecretKey(
            format!("Unknown Dilithium secret key size: {} bytes", secret_key.len())
        )),
    }
}

/// Verify a Dilithium signature
///
/// # Arguments
/// * `public_key` - Signer's public key
/// * `message` - Original message
/// * `signature` - Signature to verify
///
/// # Returns
/// `Ok(true)` if valid, `Ok(false)` or `Err` if invalid
pub fn dilithium_verify(
    public_key: &[u8],
    message: &[u8],
    signature: &[u8],
) -> CryptoResult<bool> {
    // Detect key size to determine level
    match public_key.len() {
        1312 => {
            let pk = dilithium2::PublicKey::from_bytes(public_key)
                .map_err(|_| CryptoError::InvalidPublicKey("Invalid Dilithium2 public key".to_string()))?;
            let sig = dilithium2::DetachedSignature::from_bytes(signature)
                .map_err(|_| CryptoError::InvalidSignature("Invalid Dilithium2 signature".to_string()))?;
            Ok(dilithium2::verify_detached_signature(&sig, message, &pk).is_ok())
        }
        1952 => {
            let pk = dilithium3::PublicKey::from_bytes(public_key)
                .map_err(|_| CryptoError::InvalidPublicKey("Invalid Dilithium3 public key".to_string()))?;
            let sig = dilithium3::DetachedSignature::from_bytes(signature)
                .map_err(|_| CryptoError::InvalidSignature("Invalid Dilithium3 signature".to_string()))?;
            Ok(dilithium3::verify_detached_signature(&sig, message, &pk).is_ok())
        }
        2592 => {
            let pk = dilithium5::PublicKey::from_bytes(public_key)
                .map_err(|_| CryptoError::InvalidPublicKey("Invalid Dilithium5 public key".to_string()))?;
            let sig = dilithium5::DetachedSignature::from_bytes(signature)
                .map_err(|_| CryptoError::InvalidSignature("Invalid Dilithium5 signature".to_string()))?;
            Ok(dilithium5::verify_detached_signature(&sig, message, &pk).is_ok())
        }
        _ => Err(CryptoError::InvalidPublicKey(
            format!("Unknown Dilithium public key size: {} bytes", public_key.len())
        )),
    }
}

/// Sign a message and return signature + message combined
pub fn dilithium_sign_combined(secret_key: &[u8], message: &[u8]) -> CryptoResult<Vec<u8>> {
    let signature = dilithium_sign(secret_key, message)?;
    let mut combined = signature;
    combined.extend_from_slice(message);
    Ok(combined)
}

/// Key and signature sizes for each Dilithium variant
pub mod sizes {
    /// Dilithium2 (Level 2)
    pub mod dilithium2 {
        pub const PUBLIC_KEY: usize = 1312;
        pub const SECRET_KEY: usize = 2560;
        pub const SIGNATURE: usize = 2420;
    }

    /// Dilithium3 (Level 3)
    pub mod dilithium3 {
        pub const PUBLIC_KEY: usize = 1952;
        pub const SECRET_KEY: usize = 4032;
        pub const SIGNATURE: usize = 3309;
    }

    /// Dilithium5 (Level 5)
    pub mod dilithium5 {
        pub const PUBLIC_KEY: usize = 2592;
        pub const SECRET_KEY: usize = 4896;
        pub const SIGNATURE: usize = 4595;
    }
}

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

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

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

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

    #[test]
    fn test_dilithium3_sign_verify() {
        let keypair = dilithium_keygen(SecurityLevel::Level3).unwrap();
        let message = b"Hello, pqc-py!";

        let signature = dilithium_sign(keypair.secret_key(), message).unwrap();
        assert_eq!(signature.len(), sizes::dilithium3::SIGNATURE);

        let valid = dilithium_verify(&keypair.public_key, message, &signature).unwrap();
        assert!(valid);
    }

    #[test]
    fn test_dilithium2_sign_verify() {
        let keypair = dilithium_keygen(SecurityLevel::Level1).unwrap();
        let message = b"Test message";

        let signature = dilithium_sign(keypair.secret_key(), message).unwrap();
        let valid = dilithium_verify(&keypair.public_key, message, &signature).unwrap();
        assert!(valid);
    }

    #[test]
    fn test_dilithium5_sign_verify() {
        let keypair = dilithium_keygen(SecurityLevel::Level5).unwrap();
        let message = b"Test message";

        let signature = dilithium_sign(keypair.secret_key(), message).unwrap();
        let valid = dilithium_verify(&keypair.public_key, message, &signature).unwrap();
        assert!(valid);
    }

    #[test]
    fn test_wrong_message() {
        let keypair = dilithium_keygen(SecurityLevel::Level3).unwrap();
        let message = b"Original message";
        let wrong_message = b"Wrong message";

        let signature = dilithium_sign(keypair.secret_key(), message).unwrap();
        let valid = dilithium_verify(&keypair.public_key, wrong_message, &signature).unwrap();
        assert!(!valid);
    }

    #[test]
    fn test_wrong_key() {
        let keypair1 = dilithium_keygen(SecurityLevel::Level3).unwrap();
        let keypair2 = dilithium_keygen(SecurityLevel::Level3).unwrap();
        let message = b"Test message";

        let signature = dilithium_sign(keypair1.secret_key(), message).unwrap();
        let valid = dilithium_verify(&keypair2.public_key, message, &signature).unwrap();
        assert!(!valid);
    }

    #[test]
    fn test_tampered_signature() {
        let keypair = dilithium_keygen(SecurityLevel::Level3).unwrap();
        let message = b"Test message";

        let mut signature = dilithium_sign(keypair.secret_key(), message).unwrap();
        signature[0] ^= 0xFF;

        // Tampered signature should be invalid format or fail verification
        let result = dilithium_verify(&keypair.public_key, message, &signature);
        // Either parse error or verification failure
        assert!(result.is_err() || result == Ok(false));
    }

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

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

    #[test]
    fn test_empty_message() {
        let keypair = dilithium_keygen(SecurityLevel::Level3).unwrap();
        let message = b"";

        let signature = dilithium_sign(keypair.secret_key(), message).unwrap();
        let valid = dilithium_verify(&keypair.public_key, message, &signature).unwrap();
        assert!(valid);
    }

    #[test]
    fn test_large_message() {
        let keypair = dilithium_keygen(SecurityLevel::Level3).unwrap();
        let message = vec![0xABu8; 1_000_000]; // 1 MB

        let signature = dilithium_sign(keypair.secret_key(), &message).unwrap();
        let valid = dilithium_verify(&keypair.public_key, &message, &signature).unwrap();
        assert!(valid);
    }

    #[test]
    fn test_deterministic_signature() {
        let keypair = dilithium_keygen(SecurityLevel::Level3).unwrap();
        let message = b"Test message";

        let sig1 = dilithium_sign(keypair.secret_key(), message).unwrap();
        let sig2 = dilithium_sign(keypair.secret_key(), message).unwrap();

        // Note: Dilithium signatures may or may not be deterministic depending on implementation
        // The important thing is both verify correctly
        assert!(dilithium_verify(&keypair.public_key, message, &sig1).unwrap());
        assert!(dilithium_verify(&keypair.public_key, message, &sig2).unwrap());
    }
}
