//! # SWing: Hybrid Digital Signature Scheme
//!
//! SWing is a paranoia-grade, triple-layered hybrid digital signature scheme
//! that provides **NIST Level 5 (256-bit)** post-quantum security.  It follows
//! a **strict redundancy** policy: a composite signature is only accepted if
//! **all three** underlying sub-signatures verify successfully.
//!
//! ## Security Composition
//!
//! | Layer | Algorithm | Security assumption |
//! |-------|-----------|---------------------|
//! | 1 | **Ed25519** | Classical ECDSA — Edwards-curve discrete log |
//! | 2 | **ML-DSA-87** (Dilithium) | Module-lattice MLWE (NIST PQ Level 5) |
//! | 3 | **SLH-DSA-SHAKE256f** (Sphincs+) | Hash-based — conservative PQ |
//!
//! This triple-layer defence ensures that authentication remains secure even
//! after a complete break of two out of three primitives, whether by a quantum
//! adversary, lattice cryptanalysis, or classical side-channel attack.
//!
//! ## Signature and Key Sizes
//!
//! | Item | Bytes |
//! |------|-------|
//! | Composite verification key | 2,688 |
//! | Composite signature | 54,547 |
//! | Master seed | 160 |
//!
//! ## Design Principles
//!
//! * **Domain-separated binding.** Every sub-signature is computed over a
//!   Sha3-512 digest that binds the caller-supplied `context`, the composite
//!   verification key, and the domain constant [`SWING_CONTEXT`].  This
//!   prevents cross-key and cross-context signature reuse.
//! * **Caller-supplied randomness.** The 64-byte `random_seed` passed to
//!   [`SWing::sign`] is the only source of non-determinism.  The API is
//!   suitable for `no_std`, WASM, and embedded targets.
//! * **Heap-cached keys.** [`SWing`] pre-computes the composite verification
//!   key once, amortising key-derivation costs across many sign/verify calls.
//! * **Strict size validation.** [`SWing::verify`] checks both the
//!   verification-key length and the signature length before any cryptographic
//!   work, returning [`Error::InvalidFormat`] on mismatch.
//!
//! ## Usage
//!
//! ```rust,no_run
//! # #[cfg(feature = "sign")] {
//! use b_wing::{SWing, SignError};
//!
//! // --- Key generation (signer) ---------------------------------------
//! // Fill from a CSPRNG in production (e.g. `getrandom::fill`).
//! let master_seed = [0u8; 160];
//! let signer = SWing::from_seed(&master_seed).unwrap();
//! let vk = signer.get_pub_key(); // publish / distribute to verifiers
//!
//! // --- Signing --------------------------------------------------------
//! let message = b"Authenticated payload";
//! let context = b"my-app-v1";
//! // MUST be freshly generated for every signature — never reuse!
//! let random_seed = [1u8; 64];
//! let signature = signer.sign(message, context, random_seed).unwrap();
//!
//! // --- Verification --------------------------------------------------
//! // Succeeds only if all three sub-signatures are valid.
//! let ok = SWing::verify(vk, message, context, &signature).unwrap();
//! assert!(ok);
//! # }
//! ```

use sha3::{Digest, Sha3_512};
use zeroize::Zeroizing;

// RNG Imports for PQ tricking
use rand_chacha::ChaCha20Rng;
use rand_core::SeedableRng;

// Cryptographic Primitives
use ed25519_dalek::{
    Signature as EdSignature, Signer as EdSigner, SigningKey as EdSigningKey,
    VerifyingKey as EdVerifyingKey,
};
use ml_dsa::{
    EncodedVerifyingKey as MlEncodedVerifyingKey, MlDsa87, Signature as MlSignature,
    SigningKey as MlSigningKey, VerifyingKey as MlVerifyingKey,
    signature::SignatureEncoding as MlSignatureEncoding,
};
use slh_dsa::{
    Shake256f, Signature as SlhSignature, SigningKey as SlhSigningKey,
    VerifyingKey as SlhVerifyingKey, signature::Keypair,
};

// ======================================================================
// Constants & Error Types
// ======================================================================

/// Byte size of the Ed25519 verification key (compressed Edwards point).
const ED25519_VK_SIZE: usize = 32;
/// Byte size of an Ed25519 signature.
const ED25519_SIG_SIZE: usize = 64;

/// Byte size of the ML-DSA-87 verification key.
const ML_DSA_87_VK_SIZE: usize = 2592;
/// Byte size of an ML-DSA-87 signature.
const ML_DSA_87_SIG_SIZE: usize = 4627;

/// Byte size of the SLH-DSA-SHAKE256f verification key.
const SLH_DSA_256F_VK_SIZE: usize = 64;
/// Byte size of an SLH-DSA-SHAKE256f signature.
const SLH_DSA_256F_SIG_SIZE: usize = 49856;

/// Required byte length of the master seed passed to [`SWing::from_seed`].
///
/// The 160 bytes are partitioned deterministically as follows:
///
/// | Bytes | Usage |
/// |-------|-------|
/// | `[0..32]`   | Ed25519 signing-key seed |
/// | `[32..64]`  | ML-DSA-87 keygen seed |
/// | `[64..96]`  | SLH-DSA SK seed (`sk_seed`) |
/// | `[96..128]` | SLH-DSA SK PRF (`sk_prf`) |
/// | `[128..160]`| SLH-DSA PK seed (`pk_seed`) |
pub const MASTER_SEED_SIZE: usize = 160;

/// Fixed 64-byte domain-separation tag prepended to every `hash_message` call.
///
/// Binds all three sub-signatures to the SWing construction, preventing
/// cross-protocol confusion.  This constant MUST NOT change across versions;
/// doing so would silently invalidate all previously issued signatures.
const SWING_CONTEXT: &'static [u8; 64] = &[
    43, 229, 51, 244, 223, 83, 139, 19, 204, 237, 234, 148, 235, 203, 225, 75, 40, 130, 210, 79,
    216, 63, 240, 161, 160, 1, 159, 249, 38, 28, 76, 151, 177, 69, 13, 232, 211, 26, 66, 140, 75,
    147, 59, 224, 39, 119, 19, 82, 246, 79, 191, 160, 125, 101, 184, 224, 43, 78, 246, 110, 192,
    156, 169, 14,
];

/// Errors that can occur during SWing signing or verification.
///
/// All variants implement [`core::fmt::Display`] for ergonomic error propagation.
///
/// # Example
///
/// ```rust,no_run
/// # #[cfg(feature = "sign")] {
/// use b_wing::{SWing, SignError};
/// let bad_vk = vec![0u8; 10]; // wrong size
/// assert_eq!(
///     SWing::verify(&bad_vk, b"msg", b"ctx", &vec![0u8; SWing::SIGNATURE_SIZE]),
///     Err(SignError::InvalidFormat),
/// );
/// # }
/// ```
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Error {
    /// One of the three underlying signing primitives returned an error.
    ///
    /// This is unexpected under normal operation; it may indicate an issue
    /// with the caller-supplied `random_seed` or an upstream library fault.
    SigningFailed,
    /// A key, signature, or seed was the wrong length or could not be parsed.
    ///
    /// Expected sizes: [`SWing::VERIFICATION_KEY_SIZE`] and [`SWing::SIGNATURE_SIZE`].
    InvalidFormat,
    /// The composite signature failed cryptographic verification.
    ///
    /// This is returned as soon as *any* sub-signature fails; the remaining
    /// layers are not evaluated (fail-fast behaviour).
    VerificationFailed,
}

impl core::fmt::Display for Error {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Error::SigningFailed => write!(f, "Signing operation failed"),
            Error::InvalidFormat => write!(f, "Invalid format or size"),
            Error::VerificationFailed => write!(f, "Signature verification failed"),
        }
    }
}

// ======================================================================
// Expanded SWing Key (Stateful / High-Throughput)
// ======================================================================

/// A stateful, high-throughput SWing key holder for signing.
///
/// `SWing` derives all three component signing keys from a single 160-byte
/// master seed via [`from_seed`][SWing::from_seed] and caches the composite
/// verification key in heap-allocated memory.  Subsequent calls to
/// [`sign`][SWing::sign] reuse the cached keys without any additional
/// key-derivation overhead.
///
/// # Key Sizes
///
/// | Constant | Bytes | Components |
/// |----------|-------|------------|
/// | [`VERIFICATION_KEY_SIZE`][SWing::VERIFICATION_KEY_SIZE] | 2,688 | `Ed25519(32) ‖ ML-DSA-87(2592) ‖ SLH-DSA(64)` |
/// | [`SIGNATURE_SIZE`][SWing::SIGNATURE_SIZE] | 54,547 | `Ed25519(64) ‖ ML-DSA-87(4627) ‖ SLH-DSA(49856)` |
///
/// # Security Note
///
/// The verification key returned by [`get_pub_key`][SWing::get_pub_key] is
/// safe to distribute publicly.  Secret key material is stored in fields that
/// implement `zeroize::ZeroizeOnDrop` (via the underlying crates) and will be
/// erased from memory when `SWing` is dropped.
///
/// # Example
///
/// ```rust,no_run
/// # #[cfg(feature = "sign")] {
/// use b_wing::SWing;
///
/// let master_seed = [0u8; 160]; // use a real CSPRNG in production
/// let swing = SWing::from_seed(&master_seed).unwrap();
///
/// let vk = swing.get_pub_key();
/// assert_eq!(vk.len(), SWing::VERIFICATION_KEY_SIZE);
/// # }
/// ```
pub struct SWing {
    ed_sk: EdSigningKey,
    ml_sk: Box<MlSigningKey<MlDsa87>>,
    slh_sk: SlhSigningKey<Shake256f>,
    composite_pk: Vec<u8>,
}

impl SWing {
    /// Byte length of the composite verification key: **2,688 bytes**.
    ///
    /// Memory layout:
    /// ```text
    /// [ Ed25519 vk (32) | ML-DSA-87 vk (2592) | SLH-DSA-SHAKE256f vk (64) ]
    /// ```
    pub const VERIFICATION_KEY_SIZE: usize =
        ED25519_VK_SIZE + ML_DSA_87_VK_SIZE + SLH_DSA_256F_VK_SIZE;

    /// Byte length of the composite signature: **54,547 bytes**.
    ///
    /// Memory layout:
    /// ```text
    /// [ Ed25519 sig (64) | ML-DSA-87 sig (4627) | SLH-DSA-SHAKE256f sig (49856) ]
    /// ```
    pub const SIGNATURE_SIZE: usize = ED25519_SIG_SIZE + ML_DSA_87_SIG_SIZE + SLH_DSA_256F_SIG_SIZE;

    /// Expands a 160-byte master seed into a fully initialised `SWing` key holder.
    ///
    /// The seed is partitioned according to [`MASTER_SEED_SIZE`] into the three
    /// component signing keys.  All derivation is deterministic.
    ///
    /// # Security Requirements
    ///
    /// * `master_seed` **must** be generated by a cryptographically secure
    ///   random number generator (CSPRNG) such as `getrandom`.
    /// * Treat the master seed with the same care as a private key — it must
    ///   never be logged, transmitted, or stored in plaintext.
    /// * If the seed is compromised, all signatures produced by this instance
    ///   can be forged.  Rotate by generating a new seed and re-publishing the
    ///   corresponding verification key.
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidFormat`] if an internal slice conversion fails
    /// (impossible for a correctly-sized input).
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "sign")] {
    /// use b_wing::SWing;
    ///
    /// let mut seed = [0u8; 160];
    /// getrandom::fill(&mut seed).expect("CSPRNG failed");
    /// let swing = SWing::from_seed(&seed).expect("key generation failed");
    /// # }
    /// ```
    pub fn from_seed(master_seed: &[u8; MASTER_SEED_SIZE]) -> Result<Self, Error> {
        let ed_sk = EdSigningKey::from_bytes(
            master_seed[0..32]
                .try_into()
                .map_err(|_| Error::InvalidFormat)?,
        );
        let ml_sk = Box::new(MlSigningKey::<MlDsa87>::from_seed(
            master_seed[32..64]
                .try_into()
                .map_err(|_| Error::InvalidFormat)?,
        ));

        let sk_seed = &master_seed[64..96];
        let sk_prf = &master_seed[96..128];
        let pk_seed = &master_seed[128..160];

        // Note: Using an internal method. Upgrading `slh-dsa` might break this API.
        let slh_sk = SlhSigningKey::<Shake256f>::slh_keygen_internal(sk_seed, sk_prf, pk_seed);

        let mut composite_pk = Vec::with_capacity(SWing::VERIFICATION_KEY_SIZE);
        composite_pk.extend_from_slice(EdVerifyingKey::from(&ed_sk).as_bytes());
        composite_pk.extend_from_slice(&ml_sk.verifying_key().encode());
        composite_pk.extend_from_slice(&slh_sk.verifying_key().to_bytes());

        Ok(Self {
            ed_sk,
            ml_sk,
            slh_sk,
            composite_pk,
        })
    }

    /// Returns a reference to the cached composite verification key.
    ///
    /// The returned slice is [`VERIFICATION_KEY_SIZE`][SWing::VERIFICATION_KEY_SIZE]
    /// bytes long and is safe to distribute publicly.  Pass it to
    /// [`verify`][SWing::verify] on the verifier's side.
    #[must_use]
    pub fn get_pub_key(&self) -> &[u8] {
        &self.composite_pk
    }

    /// Signs `message` using all three underlying schemes and returns the composite signature.
    ///
    /// The composite signature is the concatenation:
    ///
    /// ```text
    /// [ Ed25519 sig (64) | ML-DSA-87 sig (4627) | SLH-DSA-SHAKE256f sig (49856) ]
    /// ```
    ///
    /// Each sub-signature is computed over a Sha3-512 digest that binds
    /// `context` and the composite verification key (see [`hash_message`]),
    /// preventing cross-context and cross-key reuse.
    ///
    /// # Arguments
    ///
    /// * `message`     — Arbitrary-length payload to authenticate.
    /// * `context`     — Application-defined domain tag (e.g. `b"my-app-v1"`).
    ///   Use a distinct context for each logical operation to prevent
    ///   cross-context signature reuse.
    /// * `random_seed` — 64 bytes of fresh CSPRNG entropy.  The value is
    ///   consumed and zeroized after use.  While SWing uses randomized
    ///   signing to defend against side-channel attacks, reusing a seed
    ///   against the same message/key remains secure (falling back to
    ///   deterministic security).
    ///
    /// # Errors
    ///
    /// | Variant | Cause |
    /// |---------|-------|
    /// | [`Error::SigningFailed`]  | An underlying primitive returned an error |
    /// | [`Error::InvalidFormat`]  | Internal slice conversion failed (should not occur) |
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "sign")] {
    /// use b_wing::SWing;
    ///
    /// # let seed = [0u8; 160];
    /// # let swing = SWing::from_seed(&seed).unwrap();
    /// let mut random_seed = [0u8; 64];
    /// getrandom::fill(&mut random_seed).expect("CSPRNG failed");
    ///
    /// let sig = swing.sign(b"hello", b"my-app-v1", random_seed).unwrap();
    /// assert_eq!(sig.len(), SWing::SIGNATURE_SIZE);
    /// # }
    /// ```
    /// # Memory Usage
    ///
    /// For the Ed25519 layer, this method performs a heap allocation
    /// (`.concat()`) proportional to the size of the message.  When signing
    /// gigabyte-scale payloads, ensure the system has sufficient memory or
    /// pre-hash the payload before passing it to SWing.
    pub fn sign(
        &self,
        message: &[u8],
        context: &[u8],
        random_seed: impl Into<Zeroizing<[u8; 64]>>,
    ) -> Result<Vec<u8>, Error> {
        let random_seed: Zeroizing<[u8; 64]> = random_seed.into();

        let context_hash = hash_message(context, &self.get_pub_key());
        let mut combined_sig = Vec::with_capacity(SWing::SIGNATURE_SIZE);

        // 1. Ed25519
        // NOTE: Appends message to context_hash.
        combined_sig.extend_from_slice(
            &self
                .ed_sk
                .sign(&[&context_hash, message].concat())
                .to_bytes(),
        );

        // 2. ML-DSA-87
        let mut ml_seed = Zeroizing::new([0u8; 32]);
        ml_seed.copy_from_slice(&random_seed[32..64]);
        let mut ml_rng = ChaCha20Rng::from_seed(*ml_seed);

        let ml_signature = self
            .ml_sk
            .sign_randomized(&message, &context_hash, &mut ml_rng)
            .map_err(|_| Error::SigningFailed)?;
        combined_sig.extend_from_slice(&ml_signature.to_bytes());

        // 3. SLH-DSA-SHAKE256f
        let slh_rand: Zeroizing<[u8; 32]> = Zeroizing::new(
            random_seed[0..32]
                .try_into()
                .map_err(|_| Error::InvalidFormat)?,
        );
        let slh_signature = Box::new(
            self.slh_sk
                .try_sign_with_context(&message, &context_hash, Some(&*slh_rand))
                .map_err(|_| Error::SigningFailed)?,
        );
        combined_sig.extend_from_slice(&slh_signature.to_vec());

        Ok(combined_sig)
    }

    /// Verifies a composite SWing signature against the provided verification key.
    ///
    /// Verification proceeds sequentially through all three sub-signatures in
    /// the order **Ed25519 → ML-DSA-87 → SLH-DSA-SHAKE256f** and returns
    /// [`Error::VerificationFailed`] as soon as any one of them fails
    /// (fail-fast behaviour).  A return value of `Ok(true)` guarantees that
    /// *all three* sub-signatures are valid.
    ///
    /// # Arguments
    ///
    /// * `vk`        — [`VERIFICATION_KEY_SIZE`][SWing::VERIFICATION_KEY_SIZE]-byte
    ///   composite verification key obtained from [`get_pub_key`][SWing::get_pub_key].
    /// * `message`   — The original payload that was signed.
    /// * `context`   — The context tag used during signing (must match exactly).
    /// * `signature` — [`SIGNATURE_SIZE`][SWing::SIGNATURE_SIZE]-byte composite
    ///   signature produced by [`sign`][SWing::sign].
    ///
    /// # Errors
    ///
    /// | Variant | Cause |
    /// |---------|-------|
    /// | [`Error::InvalidFormat`]      | `vk` or `signature` has wrong length, or a sub-key cannot be parsed |
    /// | [`Error::VerificationFailed`] | Any sub-signature is invalid |
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "sign")] {
    /// use b_wing::SWing;
    ///
    /// # let seed = [0u8; 160];
    /// # let swing = SWing::from_seed(&seed).unwrap();
    /// # let vk = swing.get_pub_key().to_vec();
    /// # let sig = swing.sign(b"hello", b"ctx", [3u8; 64]).unwrap();
    /// let ok = SWing::verify(&vk, b"hello", b"ctx", &sig).unwrap();
    /// assert!(ok);
    ///
    /// // Tampered message must fail:
    /// assert!(SWing::verify(&vk, b"TAMPERED", b"ctx", &sig).is_err());
    /// # }
    /// ```
    #[must_use]
    pub fn verify(
        vk: &[u8],
        message: &[u8],
        context: &[u8],
        signature: &[u8],
    ) -> Result<bool, Error> {
        // Enforce strict size constraints prior to performing slice indexing
        if vk.len() != Self::VERIFICATION_KEY_SIZE || signature.len() != Self::SIGNATURE_SIZE {
            return Err(Error::InvalidFormat);
        }

        let context_hash = hash_message(context, vk);

        let ed_vk_end = ED25519_VK_SIZE;
        let ml_vk_end = ed_vk_end + ML_DSA_87_VK_SIZE;

        let ed_sig_end = ED25519_SIG_SIZE;
        let ml_sig_end = ed_sig_end + ML_DSA_87_SIG_SIZE;

        // 1. Ed25519
        let ed_vk = EdVerifyingKey::from_bytes(
            vk[..ed_vk_end]
                .try_into()
                .map_err(|_| Error::InvalidFormat)?,
        )
        .map_err(|_| Error::InvalidFormat)?;

        let ed_sig = EdSignature::from_bytes(
            signature[..ed_sig_end]
                .try_into()
                .map_err(|_| Error::InvalidFormat)?,
        );

        // NOTE: Appends message to context_hash.
        ed_vk
            .verify_strict(&[&context_hash, message].concat(), &ed_sig)
            .map_err(|_| Error::VerificationFailed)?;

        // 2. ML-DSA-87
        let ml_vk = MlVerifyingKey::<MlDsa87>::decode(
            &MlEncodedVerifyingKey::<MlDsa87>::try_from(&vk[ed_vk_end..ml_vk_end])
                .map_err(|_| Error::InvalidFormat)?,
        );
        let ml_sig = MlSignature::try_from(&signature[ed_sig_end..ml_sig_end])
            .map_err(|_| Error::InvalidFormat)?;

        if !ml_vk.verify_with_context(message, &context_hash, &ml_sig) {
            return Err(Error::VerificationFailed);
        }

        // 3. SLH-DSA-SHAKE256f
        let slh_vk = SlhVerifyingKey::<Shake256f>::try_from(&vk[ml_vk_end..])
            .map_err(|_| Error::InvalidFormat)?;
        let slh_sig = Box::new(
            SlhSignature::try_from(&signature[ml_sig_end..]).map_err(|_| Error::InvalidFormat)?,
        );

        slh_vk
            .try_verify_with_context(message, &context_hash, &slh_sig)
            .map_err(|_| Error::VerificationFailed)?;

        Ok(true)
    }
}

// ======================================================================
// Cryptographic Helpers (Internal)
// ======================================================================

/// Derives a deterministic 64-byte domain-separation hash for use in all sub-signatures.
///
/// The hash binds the caller-supplied `context` string and the composite
/// verification key `composite_vk` to the fixed [`SWING_CONTEXT`] domain tag.
/// All three sub-signatures are computed over a message that includes this
/// hash, preventing:
///
/// * **Cross-context reuse:** A signature for `context = "login"` cannot be
///   replayed as a signature for `context = "payment"`.
/// * **Cross-key confusion:** A signature produced under key `A` cannot be
///   presented as valid under key `B` because the key material is mixed into
///   the digest.
///
/// # Construction
///
/// ```text
/// H = SHA3-512( SWING_CONTEXT || composite_vk || len(context) || context )
/// ```
fn hash_message(context: &[u8], composite_vk: &[u8]) -> [u8; 64] {
    let mut hasher = Sha3_512::new();
    hasher.update(SWING_CONTEXT);
    hasher.update(composite_vk);
    hasher.update((context.len() as u64).to_le_bytes());
    hasher.update(context);
    hasher.finalize().into()
}

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

    // ======================================================================
    // Test Constants & Helpers
    // ======================================================================

    const TEST_MESSAGE: &[u8] = b"Paranoia is just a higher level of awareness.";
    const TEST_CONTEXT: &[u8] = b"My-test-Prompt";

    static MASTER_SEED: [u8; 160] = [0x42; 160];
    static RANDOM_SEED: [u8; 64] = [0x84; 64];

    static S_WING: LazyLock<SWing> = LazyLock::new(|| SWing::from_seed(&MASTER_SEED).unwrap());

    static SHARED_SIG: LazyLock<Vec<u8>> = LazyLock::new(|| {
        S_WING
            .sign(TEST_MESSAGE, TEST_CONTEXT, RANDOM_SEED)
            .unwrap()
    });

    // ======================================================================
    // Happy Path & Determinism
    // ======================================================================

    #[test]
    fn test_happy_path_round_trip() {
        // 1. Generate Public Key
        let vk = S_WING.get_pub_key();
        assert_eq!(vk.len(), SWing::VERIFICATION_KEY_SIZE);

        // 2. Sign
        let signature = &*SHARED_SIG;
        assert_eq!(signature.len(), SWing::SIGNATURE_SIZE);

        // 3. Verify
        let is_valid = SWing::verify(vk, TEST_MESSAGE, TEST_CONTEXT, signature)
            .expect("Verification should succeed");
        assert!(is_valid, "Valid signature must pass the gauntlet");
    }

    #[test]
    fn test_strict_determinism() {
        // 1. Key Generation Determinism
        let binding = SWing::from_seed(&MASTER_SEED).unwrap();
        let vk2 = binding.get_pub_key();
        assert_eq!(
            S_WING.get_pub_key(),
            vk2,
            "Verifying keys must be identical for the same master seed"
        );

        // 2. Signature Determinism
        let sig2 = binding
            .sign(TEST_MESSAGE, TEST_CONTEXT, RANDOM_SEED)
            .unwrap();
        assert_eq!(
            *SHARED_SIG, sig2,
            "Signatures must be identical for the same seeds and message"
        );
    }

    // ======================================================================
    // Formatting & Boundary Rejections
    // ======================================================================

    #[test]
    fn test_invalid_lengths_rejected_immediately() {
        let vk = S_WING.get_pub_key();
        let sig = &*SHARED_SIG;

        // 1. Bad VK Length
        let bad_vk = vec![0u8; SWing::VERIFICATION_KEY_SIZE - 1];
        let res_vk = SWing::verify(&bad_vk, TEST_MESSAGE, TEST_CONTEXT, sig);
        assert_eq!(
            res_vk,
            Err(Error::InvalidFormat),
            "Must reject invalid VK lengths"
        );

        // 2. Bad Signature Length
        let bad_sig = vec![0u8; SWing::SIGNATURE_SIZE + 1];
        let res_sig = SWing::verify(vk, TEST_MESSAGE, TEST_CONTEXT, &bad_sig);
        assert_eq!(
            res_sig,
            Err(Error::InvalidFormat),
            "Must reject invalid signature lengths"
        );
    }

    // ======================================================================
    // The Gauntlet (Tampering & Fail-Fast Logic)
    // ======================================================================

    #[test]
    fn test_message_and_context_tampering_fails_fast() {
        let vk = S_WING.get_pub_key();
        let sig = &*SHARED_SIG;

        // Tamper Message
        let res_msg = SWing::verify(vk, b"I am an attacker", TEST_CONTEXT, sig);
        assert_eq!(res_msg, Err(Error::VerificationFailed));

        // Tamper Context
        let res_ctx = SWing::verify(vk, TEST_MESSAGE, b"Wrong-Context", sig);
        assert_eq!(res_ctx, Err(Error::VerificationFailed));
    }

    #[test]
    fn test_wrong_public_key_rejection() {
        let bob_seed = [0x99; 160];
        let binding = SWing::from_seed(&bob_seed).unwrap();
        let bob_vk = binding.get_pub_key();

        let alice_sig = &*SHARED_SIG;

        // Server tries to verify with Bob's public key
        let result = SWing::verify(bob_vk, TEST_MESSAGE, TEST_CONTEXT, alice_sig);
        assert_eq!(
            result,
            Err(Error::VerificationFailed),
            "Signatures verified against the wrong key must fail"
        );
    }

    #[test]
    fn test_gauntlet_layer_1_ed25519_tampering() {
        let vk = S_WING.get_pub_key();
        let mut sig = SHARED_SIG.clone();

        // Tamper with the Ed25519 section of the signature (Bytes 0..64)
        sig[10] ^= 0xFF;

        let result = SWing::verify(vk, TEST_MESSAGE, TEST_CONTEXT, &sig);
        assert_eq!(
            result,
            Err(Error::VerificationFailed),
            "Altering the first 64 bytes must trigger Ed25519 rejection"
        );
    }

    #[test]
    fn test_gauntlet_layer_2_ml_dsa_tampering() {
        let vk = S_WING.get_pub_key();
        let mut sig = SHARED_SIG.clone();

        // Tamper with the ML-DSA section of the signature (Bytes 64..4691)
        sig[1000] ^= 0xFF;

        let result = SWing::verify(vk, TEST_MESSAGE, TEST_CONTEXT, &sig);

        match result {
            Err(Error::VerificationFailed) | Err(Error::InvalidFormat) => {}
            _ => panic!(
                "Expected ML-DSA to reject the tampered signature, got {:?}",
                result
            ),
        }
    }

    #[test]
    fn test_gauntlet_layer_3_slh_dsa_tampering() {
        let vk = S_WING.get_pub_key();
        let mut sig = SHARED_SIG.clone();

        // Tamper with the SLH-DSA section of the signature (Bytes 4691..End)
        sig[5000] ^= 0xFF;

        let result = SWing::verify(vk, TEST_MESSAGE, TEST_CONTEXT, &sig);

        match result {
            Err(Error::VerificationFailed) | Err(Error::InvalidFormat) => {}
            _ => panic!(
                "Expected SLH-DSA to reject the tampered signature, got {:?}",
                result
            ),
        }
    }
}
