//! FALCON (FN-DSA) compact lattice-based digital signatures.
//!
//! FALCON is a lattice-based signature scheme selected by NIST as FIPS 206
//! (FN-DSA). It provides significantly smaller signatures than Dilithium,
//! making it ideal for bandwidth-constrained applications like mobile wallets
//! and IoT devices.
//!
//! # Key Properties
//!
//! - **Compact Signatures**: ~690 bytes (vs ~3,300 for Dilithium3)
//! - **Fast Verification**: Faster than Dilithium verification
//! - **NTRU Lattices**: Based on NTRU lattice problems with FFT optimization
//!
//! # Trade-offs
//!
//! | Variant | Security | Signature | Public Key | Private Key |
//! |---------|----------|-----------|------------|-------------|
//! | FALCON-512 | 128-bit | 690 bytes | 897 bytes | 1,281 bytes |
//! | FALCON-1024 | 256-bit | 1,330 bytes | 1,793 bytes | 2,305 bytes |
//!
//! Compared to Dilithium3 (192-bit, 3,293 byte signatures), FALCON-512
//! achieves ~80% signature size reduction at 128-bit security.
//!
//! # Usage
//!
//! ```rust
//! use synor_crypto::falcon::{FalconKeypair, FalconVariant};
//!
//! // Generate keypair (FALCON-512 for mobile devices)
//! let keypair = FalconKeypair::generate(FalconVariant::Falcon512);
//!
//! // Sign a message
//! let message = b"Transaction payload";
//! let signature = keypair.sign(message);
//!
//! // Verify signature
//! assert!(keypair.public_key().verify(message, &signature).is_ok());
//!
//! // Signature is only ~690 bytes!
//! println!("Signature size: {} bytes", signature.len());
//! ```
//!
//! # When to Use FALCON vs Dilithium
//!
//! | Use Case | Recommendation |
//! |----------|----------------|
//! | Desktop wallets | Dilithium3 (higher security) |
//! | Mobile wallets | FALCON-512 (smaller signatures) |
//! | IoT devices | FALCON-512 (bandwidth constrained) |
//! | Layer 2 transactions | FALCON-512 (batch efficiency) |
//! | High-value transactions | Dilithium3 (conservative choice) |

use pqcrypto_falcon::falcon512;
use pqcrypto_falcon::falcon1024;
use pqcrypto_traits::sign::{
    DetachedSignature, PublicKey as PqPublicKey, SecretKey as PqSecretKey,
};
use thiserror::Error;
use zeroize::{Zeroize, ZeroizeOnDrop};

/// FALCON variant selection.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FalconVariant {
    /// 128-bit security, ~690 byte signatures
    Falcon512,
    /// 256-bit security, ~1,330 byte signatures
    Falcon1024,
}

impl FalconVariant {
    /// Returns the expected signature size in bytes.
    pub const fn signature_size(&self) -> usize {
        match self {
            FalconVariant::Falcon512 => 690,
            FalconVariant::Falcon1024 => 1330,
        }
    }

    /// Returns the expected public key size in bytes.
    pub const fn public_key_size(&self) -> usize {
        match self {
            FalconVariant::Falcon512 => 897,
            FalconVariant::Falcon1024 => 1793,
        }
    }

    /// Returns the expected secret key size in bytes.
    pub const fn secret_key_size(&self) -> usize {
        match self {
            FalconVariant::Falcon512 => 1281,
            FalconVariant::Falcon1024 => 2305,
        }
    }

    /// Returns the security level in bits.
    pub const fn security_level(&self) -> u16 {
        match self {
            FalconVariant::Falcon512 => 128,
            FalconVariant::Falcon1024 => 256,
        }
    }

    /// Returns the NIST security category.
    pub const fn nist_category(&self) -> u8 {
        match self {
            FalconVariant::Falcon512 => 1,
            FalconVariant::Falcon1024 => 5,
        }
    }

    /// Returns the n parameter (lattice dimension).
    pub const fn n(&self) -> u16 {
        match self {
            FalconVariant::Falcon512 => 512,
            FalconVariant::Falcon1024 => 1024,
        }
    }
}

impl Default for FalconVariant {
    fn default() -> Self {
        FalconVariant::Falcon512
    }
}

/// FALCON public key.
#[derive(Clone)]
pub struct FalconPublicKey {
    variant: FalconVariant,
    bytes: Vec<u8>,
}

impl FalconPublicKey {
    /// Creates a public key from bytes.
    pub fn from_bytes(variant: FalconVariant, bytes: &[u8]) -> Result<Self, FalconError> {
        if bytes.len() != variant.public_key_size() {
            return Err(FalconError::InvalidKeySize {
                expected: variant.public_key_size(),
                got: bytes.len(),
            });
        }
        Ok(Self {
            variant,
            bytes: bytes.to_vec(),
        })
    }

    /// Returns the public key bytes.
    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes
    }

    /// Returns the variant used.
    pub fn variant(&self) -> FalconVariant {
        self.variant
    }

    /// Verifies a signature against this public key.
    pub fn verify(&self, message: &[u8], signature: &FalconSignature) -> Result<(), FalconError> {
        if signature.variant != self.variant {
            return Err(FalconError::VariantMismatch);
        }

        match self.variant {
            FalconVariant::Falcon512 => {
                let pk = falcon512::PublicKey::from_bytes(&self.bytes)
                    .map_err(|_| FalconError::InvalidPublicKey)?;
                let sig = falcon512::DetachedSignature::from_bytes(&signature.bytes)
                    .map_err(|_| FalconError::InvalidSignature)?;
                falcon512::verify_detached_signature(&sig, message, &pk)
                    .map_err(|_| FalconError::VerificationFailed)
            }
            FalconVariant::Falcon1024 => {
                let pk = falcon1024::PublicKey::from_bytes(&self.bytes)
                    .map_err(|_| FalconError::InvalidPublicKey)?;
                let sig = falcon1024::DetachedSignature::from_bytes(&signature.bytes)
                    .map_err(|_| FalconError::InvalidSignature)?;
                falcon1024::verify_detached_signature(&sig, message, &pk)
                    .map_err(|_| FalconError::VerificationFailed)
            }
        }
    }
}

impl std::fmt::Debug for FalconPublicKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FalconPublicKey")
            .field("variant", &self.variant)
            .field("bytes", &hex::encode(&self.bytes[..8.min(self.bytes.len())]))
            .finish()
    }
}

/// FALCON secret key (zeroized on drop).
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct FalconSecretKey {
    #[zeroize(skip)]
    variant: FalconVariant,
    bytes: Vec<u8>,
}

impl FalconSecretKey {
    /// Creates a secret key from bytes.
    pub fn from_bytes(variant: FalconVariant, bytes: &[u8]) -> Result<Self, FalconError> {
        if bytes.len() != variant.secret_key_size() {
            return Err(FalconError::InvalidKeySize {
                expected: variant.secret_key_size(),
                got: bytes.len(),
            });
        }
        Ok(Self {
            variant,
            bytes: bytes.to_vec(),
        })
    }

    /// Returns the secret key bytes.
    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes
    }

    /// Returns the variant used.
    pub fn variant(&self) -> FalconVariant {
        self.variant
    }
}

/// FALCON signature.
#[derive(Clone)]
pub struct FalconSignature {
    variant: FalconVariant,
    bytes: Vec<u8>,
}

impl FalconSignature {
    /// Creates a signature from bytes.
    pub fn from_bytes(variant: FalconVariant, bytes: &[u8]) -> Result<Self, FalconError> {
        // FALCON signatures have variable length, so we allow some flexibility
        let expected = variant.signature_size();
        // Allow up to 10% variance in signature size
        if bytes.len() > expected + expected / 10 {
            return Err(FalconError::InvalidSignatureSize {
                expected,
                got: bytes.len(),
            });
        }
        Ok(Self {
            variant,
            bytes: bytes.to_vec(),
        })
    }

    /// Returns the signature bytes.
    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes
    }

    /// Returns the signature length.
    pub fn len(&self) -> usize {
        self.bytes.len()
    }

    /// Returns true if the signature is empty.
    pub fn is_empty(&self) -> bool {
        self.bytes.is_empty()
    }

    /// Returns the variant used.
    pub fn variant(&self) -> FalconVariant {
        self.variant
    }
}

impl std::fmt::Debug for FalconSignature {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FalconSignature")
            .field("variant", &self.variant)
            .field("size", &self.bytes.len())
            .finish()
    }
}

/// FALCON keypair.
pub struct FalconKeypair {
    variant: FalconVariant,
    public_key: FalconPublicKey,
    secret_key: FalconSecretKey,
}

impl FalconKeypair {
    /// Generates a new random keypair.
    pub fn generate(variant: FalconVariant) -> Self {
        match variant {
            FalconVariant::Falcon512 => {
                let (pk, sk) = falcon512::keypair();
                Self {
                    variant,
                    public_key: FalconPublicKey {
                        variant,
                        bytes: pk.as_bytes().to_vec(),
                    },
                    secret_key: FalconSecretKey {
                        variant,
                        bytes: sk.as_bytes().to_vec(),
                    },
                }
            }
            FalconVariant::Falcon1024 => {
                let (pk, sk) = falcon1024::keypair();
                Self {
                    variant,
                    public_key: FalconPublicKey {
                        variant,
                        bytes: pk.as_bytes().to_vec(),
                    },
                    secret_key: FalconSecretKey {
                        variant,
                        bytes: sk.as_bytes().to_vec(),
                    },
                }
            }
        }
    }

    /// Creates a keypair from existing keys.
    pub fn from_keys(
        public_key: FalconPublicKey,
        secret_key: FalconSecretKey,
    ) -> Result<Self, FalconError> {
        if public_key.variant != secret_key.variant {
            return Err(FalconError::VariantMismatch);
        }
        Ok(Self {
            variant: public_key.variant,
            public_key,
            secret_key,
        })
    }

    /// Signs a message.
    pub fn sign(&self, message: &[u8]) -> FalconSignature {
        let bytes = match self.variant {
            FalconVariant::Falcon512 => {
                let sk = falcon512::SecretKey::from_bytes(self.secret_key.as_bytes())
                    .expect("valid secret key");
                let sig = falcon512::detached_sign(message, &sk);
                sig.as_bytes().to_vec()
            }
            FalconVariant::Falcon1024 => {
                let sk = falcon1024::SecretKey::from_bytes(self.secret_key.as_bytes())
                    .expect("valid secret key");
                let sig = falcon1024::detached_sign(message, &sk);
                sig.as_bytes().to_vec()
            }
        };

        FalconSignature {
            variant: self.variant,
            bytes,
        }
    }

    /// Returns the public key.
    pub fn public_key(&self) -> &FalconPublicKey {
        &self.public_key
    }

    /// Returns the secret key.
    pub fn secret_key(&self) -> &FalconSecretKey {
        &self.secret_key
    }

    /// Returns the variant used.
    pub fn variant(&self) -> FalconVariant {
        self.variant
    }
}

impl std::fmt::Debug for FalconKeypair {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FalconKeypair")
            .field("variant", &self.variant)
            .field("public_key", &self.public_key)
            .finish()
    }
}

/// Errors from FALCON operations.
#[derive(Debug, Error)]
pub enum FalconError {
    #[error("Invalid key size: expected {expected}, got {got}")]
    InvalidKeySize { expected: usize, got: usize },

    #[error("Invalid signature size: expected ~{expected}, got {got}")]
    InvalidSignatureSize { expected: usize, got: usize },

    #[error("Invalid public key")]
    InvalidPublicKey,

    #[error("Invalid secret key")]
    InvalidSecretKey,

    #[error("Invalid signature")]
    InvalidSignature,

    #[error("Signature verification failed")]
    VerificationFailed,

    #[error("Variant mismatch between key and signature")]
    VariantMismatch,
}

/// Comparison of signature sizes for different algorithms.
pub mod comparison {
    /// Ed25519 signature size (classical)
    pub const ED25519_SIG_SIZE: usize = 64;

    /// Dilithium3 signature size (FIPS 204)
    pub const DILITHIUM3_SIG_SIZE: usize = 3293;

    /// FALCON-512 signature size (FIPS 206)
    pub const FALCON512_SIG_SIZE: usize = 690;

    /// SPHINCS+-128s signature size (FIPS 205)
    pub const SPHINCS128S_SIG_SIZE: usize = 7856;

    /// Calculates bandwidth savings using FALCON instead of Dilithium.
    pub fn bandwidth_savings_percent() -> f64 {
        let savings = 1.0 - (FALCON512_SIG_SIZE as f64 / DILITHIUM3_SIG_SIZE as f64);
        savings * 100.0
    }
}

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

    #[test]
    fn test_generate_512() {
        let keypair = FalconKeypair::generate(FalconVariant::Falcon512);
        assert_eq!(keypair.variant(), FalconVariant::Falcon512);
        assert_eq!(keypair.public_key().as_bytes().len(), 897);
    }

    #[test]
    fn test_sign_verify_512() {
        let keypair = FalconKeypair::generate(FalconVariant::Falcon512);
        let message = b"Test message for FALCON signing";

        let signature = keypair.sign(message);
        // FALCON signatures are variable length, but typically ~690 bytes
        assert!(signature.len() <= 800);

        // Verify should succeed
        assert!(keypair.public_key().verify(message, &signature).is_ok());
    }

    #[test]
    fn test_sign_verify_1024() {
        let keypair = FalconKeypair::generate(FalconVariant::Falcon1024);
        let message = b"Test message for FALCON-1024";

        let signature = keypair.sign(message);
        // FALCON-1024 signatures are ~1330 bytes
        assert!(signature.len() <= 1500);

        assert!(keypair.public_key().verify(message, &signature).is_ok());
    }

    #[test]
    fn test_invalid_signature_fails() {
        let keypair = FalconKeypair::generate(FalconVariant::Falcon512);
        let message = b"Original message";
        let signature = keypair.sign(message);

        // Verify with wrong message should fail
        let wrong_message = b"Wrong message";
        assert!(keypair.public_key().verify(wrong_message, &signature).is_err());
    }

    #[test]
    fn test_variant_sizes() {
        assert_eq!(FalconVariant::Falcon512.signature_size(), 690);
        assert_eq!(FalconVariant::Falcon512.public_key_size(), 897);
        assert_eq!(FalconVariant::Falcon512.secret_key_size(), 1281);

        assert_eq!(FalconVariant::Falcon1024.signature_size(), 1330);
        assert_eq!(FalconVariant::Falcon1024.public_key_size(), 1793);
    }

    #[test]
    fn test_security_levels() {
        assert_eq!(FalconVariant::Falcon512.security_level(), 128);
        assert_eq!(FalconVariant::Falcon1024.security_level(), 256);
    }

    #[test]
    fn test_bandwidth_savings() {
        let savings = comparison::bandwidth_savings_percent();
        // FALCON-512 should save ~79% bandwidth vs Dilithium3
        assert!(savings > 75.0);
        assert!(savings < 85.0);
    }
}
