//! SLH-DSA (SPHINCS+) Stateless Hash-Based Signatures
//!
//! NIST FIPS 205 implementation for post-quantum digital signatures.
//! Based on hash functions only - no lattice assumptions.
//!
//! Trade-offs:
//! - Larger signatures than Dilithium
//! - Based solely on hash function security
//! - No secret state management needed (stateless)
//!
//! Variants: SHA2 and SHAKE versions at multiple security levels

use pqcrypto_sphincsplus::{
    sphincssha2128fsimple, sphincssha2128ssimple,
    sphincssha2192fsimple, sphincssha2192ssimple,
    sphincssha2256fsimple, sphincssha2256ssimple,
    sphincsshake128fsimple, sphincsshake128ssimple,
    sphincsshake192fsimple, sphincsshake192ssimple,
    sphincsshake256fsimple, sphincsshake256ssimple,
};
use pqcrypto_traits::sign::{PublicKey as SignPublicKey, SecretKey as SignSecretKey, DetachedSignature};
use zeroize::{Zeroize, ZeroizeOnDrop};

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

/// SPHINCS+ variant selection
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SphincsVariant {
    /// SHA2-128 fast (smaller signatures, faster)
    Sha2_128f,
    /// SHA2-128 small (smaller signatures, slower)
    Sha2_128s,
    /// SHA2-192 fast
    Sha2_192f,
    /// SHA2-192 small
    Sha2_192s,
    /// SHA2-256 fast
    Sha2_256f,
    /// SHA2-256 small
    Sha2_256s,
    /// SHAKE-128 fast
    Shake_128f,
    /// SHAKE-128 small
    Shake_128s,
    /// SHAKE-192 fast
    Shake_192f,
    /// SHAKE-192 small
    Shake_192s,
    /// SHAKE-256 fast
    Shake_256f,
    /// SHAKE-256 small
    Shake_256s,
}

impl Default for SphincsVariant {
    fn default() -> Self {
        SphincsVariant::Sha2_192f
    }
}

impl SphincsVariant {
    /// Get security level in bits
    pub fn security_bits(&self) -> usize {
        match self {
            SphincsVariant::Sha2_128f | SphincsVariant::Sha2_128s |
            SphincsVariant::Shake_128f | SphincsVariant::Shake_128s => 128,
            SphincsVariant::Sha2_192f | SphincsVariant::Sha2_192s |
            SphincsVariant::Shake_192f | SphincsVariant::Shake_192s => 192,
            SphincsVariant::Sha2_256f | SphincsVariant::Sha2_256s |
            SphincsVariant::Shake_256f | SphincsVariant::Shake_256s => 256,
        }
    }

    /// Check if this is a "fast" variant (larger signatures, faster signing)
    pub fn is_fast(&self) -> bool {
        matches!(self,
            SphincsVariant::Sha2_128f | SphincsVariant::Sha2_192f | SphincsVariant::Sha2_256f |
            SphincsVariant::Shake_128f | SphincsVariant::Shake_192f | SphincsVariant::Shake_256f
        )
    }

    /// Check if this is a "small" variant (smaller signatures, slower signing)
    pub fn is_small(&self) -> bool {
        !self.is_fast()
    }
}

/// SPHINCS+ key pair
#[derive(ZeroizeOnDrop)]
pub struct SphincsKeyPair {
    /// Public key for verification
    #[zeroize(skip)]
    pub public_key: Vec<u8>,
    /// Secret key for signing (zeroed on drop)
    secret_key: Vec<u8>,
    /// Variant used
    #[zeroize(skip)]
    pub variant: SphincsVariant,
}

impl SphincsKeyPair {
    /// 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 SPHINCS+ key pair
///
/// # Arguments
/// * `variant` - The SPHINCS+ variant to use
///
/// # Returns
/// A new SPHINCS+ key pair
pub fn sphincs_keygen(variant: SphincsVariant) -> CryptoResult<SphincsKeyPair> {
    macro_rules! keygen_variant {
        ($module:ident) => {{
            let (pk, sk) = $module::keypair();
            Ok(SphincsKeyPair {
                public_key: pk.as_bytes().to_vec(),
                secret_key: sk.as_bytes().to_vec(),
                variant,
            })
        }};
    }

    match variant {
        SphincsVariant::Sha2_128f => keygen_variant!(sphincssha2128fsimple),
        SphincsVariant::Sha2_128s => keygen_variant!(sphincssha2128ssimple),
        SphincsVariant::Sha2_192f => keygen_variant!(sphincssha2192fsimple),
        SphincsVariant::Sha2_192s => keygen_variant!(sphincssha2192ssimple),
        SphincsVariant::Sha2_256f => keygen_variant!(sphincssha2256fsimple),
        SphincsVariant::Sha2_256s => keygen_variant!(sphincssha2256ssimple),
        SphincsVariant::Shake_128f => keygen_variant!(sphincsshake128fsimple),
        SphincsVariant::Shake_128s => keygen_variant!(sphincsshake128ssimple),
        SphincsVariant::Shake_192f => keygen_variant!(sphincsshake192fsimple),
        SphincsVariant::Shake_192s => keygen_variant!(sphincsshake192ssimple),
        SphincsVariant::Shake_256f => keygen_variant!(sphincsshake256fsimple),
        SphincsVariant::Shake_256s => keygen_variant!(sphincsshake256ssimple),
    }
}

/// Sign a message using SPHINCS+
///
/// Note: SPHINCS+ signatures are quite large (8KB-50KB depending on variant)
pub fn sphincs_sign(secret_key: &[u8], message: &[u8]) -> CryptoResult<Vec<u8>> {
    // Detect variant from secret key size
    // This is a simplified approach - in production you'd want to store the variant
    macro_rules! try_sign {
        ($module:ident) => {
            if let Ok(sk) = $module::SecretKey::from_bytes(secret_key) {
                let sig = $module::detached_sign(message, &sk);
                return Ok(sig.as_bytes().to_vec());
            }
        };
    }

    // Try each variant
    try_sign!(sphincssha2128fsimple);
    try_sign!(sphincssha2128ssimple);
    try_sign!(sphincssha2192fsimple);
    try_sign!(sphincssha2192ssimple);
    try_sign!(sphincssha2256fsimple);
    try_sign!(sphincssha2256ssimple);
    try_sign!(sphincsshake128fsimple);
    try_sign!(sphincsshake128ssimple);
    try_sign!(sphincsshake192fsimple);
    try_sign!(sphincsshake192ssimple);
    try_sign!(sphincsshake256fsimple);
    try_sign!(sphincsshake256ssimple);

    Err(CryptoError::InvalidSecretKey(
        format!("Unknown SPHINCS+ secret key size: {} bytes", secret_key.len())
    ))
}

/// Verify a SPHINCS+ signature
pub fn sphincs_verify(
    public_key: &[u8],
    message: &[u8],
    signature: &[u8],
) -> CryptoResult<bool> {
    macro_rules! try_verify {
        ($module:ident) => {
            if let (Ok(pk), Ok(sig)) = (
                $module::PublicKey::from_bytes(public_key),
                $module::DetachedSignature::from_bytes(signature),
            ) {
                return Ok($module::verify_detached_signature(&sig, message, &pk).is_ok());
            }
        };
    }

    // Try each variant
    try_verify!(sphincssha2128fsimple);
    try_verify!(sphincssha2128ssimple);
    try_verify!(sphincssha2192fsimple);
    try_verify!(sphincssha2192ssimple);
    try_verify!(sphincssha2256fsimple);
    try_verify!(sphincssha2256ssimple);
    try_verify!(sphincsshake128fsimple);
    try_verify!(sphincsshake128ssimple);
    try_verify!(sphincsshake192fsimple);
    try_verify!(sphincsshake192ssimple);
    try_verify!(sphincsshake256fsimple);
    try_verify!(sphincsshake256ssimple);

    Err(CryptoError::InvalidPublicKey(
        format!("Unknown SPHINCS+ public key size: {} bytes", public_key.len())
    ))
}

/// Sign with a known variant (more efficient)
pub fn sphincs_sign_with_variant(
    secret_key: &[u8],
    message: &[u8],
    variant: SphincsVariant,
) -> CryptoResult<Vec<u8>> {
    macro_rules! sign_variant {
        ($module:ident) => {{
            let sk = $module::SecretKey::from_bytes(secret_key)
                .map_err(|_| CryptoError::InvalidSecretKey("Invalid SPHINCS+ secret key".to_string()))?;
            let sig = $module::detached_sign(message, &sk);
            Ok(sig.as_bytes().to_vec())
        }};
    }

    match variant {
        SphincsVariant::Sha2_128f => sign_variant!(sphincssha2128fsimple),
        SphincsVariant::Sha2_128s => sign_variant!(sphincssha2128ssimple),
        SphincsVariant::Sha2_192f => sign_variant!(sphincssha2192fsimple),
        SphincsVariant::Sha2_192s => sign_variant!(sphincssha2192ssimple),
        SphincsVariant::Sha2_256f => sign_variant!(sphincssha2256fsimple),
        SphincsVariant::Sha2_256s => sign_variant!(sphincssha2256ssimple),
        SphincsVariant::Shake_128f => sign_variant!(sphincsshake128fsimple),
        SphincsVariant::Shake_128s => sign_variant!(sphincsshake128ssimple),
        SphincsVariant::Shake_192f => sign_variant!(sphincsshake192fsimple),
        SphincsVariant::Shake_192s => sign_variant!(sphincsshake192ssimple),
        SphincsVariant::Shake_256f => sign_variant!(sphincsshake256fsimple),
        SphincsVariant::Shake_256s => sign_variant!(sphincsshake256ssimple),
    }
}

/// Verify with a known variant (more efficient)
pub fn sphincs_verify_with_variant(
    public_key: &[u8],
    message: &[u8],
    signature: &[u8],
    variant: SphincsVariant,
) -> CryptoResult<bool> {
    macro_rules! verify_variant {
        ($module:ident) => {{
            let pk = $module::PublicKey::from_bytes(public_key)
                .map_err(|_| CryptoError::InvalidPublicKey("Invalid SPHINCS+ public key".to_string()))?;
            let sig = $module::DetachedSignature::from_bytes(signature)
                .map_err(|_| CryptoError::InvalidSignature("Invalid SPHINCS+ signature".to_string()))?;
            Ok($module::verify_detached_signature(&sig, message, &pk).is_ok())
        }};
    }

    match variant {
        SphincsVariant::Sha2_128f => verify_variant!(sphincssha2128fsimple),
        SphincsVariant::Sha2_128s => verify_variant!(sphincssha2128ssimple),
        SphincsVariant::Sha2_192f => verify_variant!(sphincssha2192fsimple),
        SphincsVariant::Sha2_192s => verify_variant!(sphincssha2192ssimple),
        SphincsVariant::Sha2_256f => verify_variant!(sphincssha2256fsimple),
        SphincsVariant::Sha2_256s => verify_variant!(sphincssha2256ssimple),
        SphincsVariant::Shake_128f => verify_variant!(sphincsshake128fsimple),
        SphincsVariant::Shake_128s => verify_variant!(sphincsshake128ssimple),
        SphincsVariant::Shake_192f => verify_variant!(sphincsshake192fsimple),
        SphincsVariant::Shake_192s => verify_variant!(sphincsshake192ssimple),
        SphincsVariant::Shake_256f => verify_variant!(sphincsshake256fsimple),
        SphincsVariant::Shake_256s => verify_variant!(sphincsshake256ssimple),
    }
}

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

    #[test]
    fn test_sphincs_sha2_128f_sign_verify() {
        let keypair = sphincs_keygen(SphincsVariant::Sha2_128f).unwrap();
        let message = b"Hello, SPHINCS+!";

        let signature = sphincs_sign_with_variant(
            keypair.secret_key(),
            message,
            SphincsVariant::Sha2_128f
        ).unwrap();

        let valid = sphincs_verify_with_variant(
            &keypair.public_key,
            message,
            &signature,
            SphincsVariant::Sha2_128f
        ).unwrap();

        assert!(valid);
    }

    #[test]
    fn test_sphincs_sha2_192f_sign_verify() {
        let keypair = sphincs_keygen(SphincsVariant::Sha2_192f).unwrap();
        let message = b"Test message";

        let signature = sphincs_sign_with_variant(
            keypair.secret_key(),
            message,
            SphincsVariant::Sha2_192f
        ).unwrap();

        let valid = sphincs_verify_with_variant(
            &keypair.public_key,
            message,
            &signature,
            SphincsVariant::Sha2_192f
        ).unwrap();

        assert!(valid);
    }

    #[test]
    fn test_sphincs_shake_256s_sign_verify() {
        let keypair = sphincs_keygen(SphincsVariant::Shake_256s).unwrap();
        let message = b"High security message";

        let signature = sphincs_sign_with_variant(
            keypair.secret_key(),
            message,
            SphincsVariant::Shake_256s
        ).unwrap();

        let valid = sphincs_verify_with_variant(
            &keypair.public_key,
            message,
            &signature,
            SphincsVariant::Shake_256s
        ).unwrap();

        assert!(valid);
    }

    #[test]
    fn test_wrong_message() {
        let keypair = sphincs_keygen(SphincsVariant::Sha2_128f).unwrap();
        let message = b"Original message";
        let wrong_message = b"Wrong message";

        let signature = sphincs_sign_with_variant(
            keypair.secret_key(),
            message,
            SphincsVariant::Sha2_128f
        ).unwrap();

        let valid = sphincs_verify_with_variant(
            &keypair.public_key,
            wrong_message,
            &signature,
            SphincsVariant::Sha2_128f
        ).unwrap();

        assert!(!valid);
    }

    #[test]
    fn test_auto_detect_sign() {
        let keypair = sphincs_keygen(SphincsVariant::Sha2_128f).unwrap();
        let message = b"Test";

        // Sign with auto-detection
        let signature = sphincs_sign(keypair.secret_key(), message).unwrap();

        // Verify with auto-detection
        let valid = sphincs_verify(&keypair.public_key, message, &signature).unwrap();
        assert!(valid);
    }

    #[test]
    fn test_variant_properties() {
        assert_eq!(SphincsVariant::Sha2_128f.security_bits(), 128);
        assert_eq!(SphincsVariant::Sha2_192s.security_bits(), 192);
        assert_eq!(SphincsVariant::Shake_256f.security_bits(), 256);

        assert!(SphincsVariant::Sha2_128f.is_fast());
        assert!(SphincsVariant::Sha2_128s.is_small());
    }

    #[test]
    fn test_default_variant() {
        let variant = SphincsVariant::default();
        assert_eq!(variant, SphincsVariant::Sha2_192f);
    }
}
