//! ML-DSA-65 (FIPS 204) digital signatures via the `ml-dsa` crate.
//!
//! Uses the seed-based API with entropy from `getrandom`. Key generation
//! is cryptographically equivalent to the RNG-based API — `key_gen(rng)`
//! internally just draws 32 random bytes then calls `from_seed` (FIPS 204 §6.1).
//!
//! Secret keys are stored as the 32-byte seed (FIPS 204 §6.1 `ξ`), not the
//! ~4032-byte expanded form. The signing key is deterministically re-expanded
//! from the seed on each `sign()` call. This minimises key material in memory
//! and on disk, and matches the canonical representation in the standard.
//!
//! ## Hedged Signing via `sign_internal`
//!
//! Signing uses the `sign_internal` API with 32 bytes of fresh `getrandom`
//! entropy as the `rnd` parameter (hedged mode). This provides fault-injection
//! resistance: with deterministic signing, a single induced fault during signing
//! (voltage glitch, Rowhammer) can extract the secret key via differential
//! analysis. Hedged signing mixes fresh randomness into the nonce derivation,
//! making each signing run independent even on the same message.
//!
//! `sign_internal` / `verify_internal` skip the FIPS 204 domain separator and
//! context string, making soliton ML-DSA signatures incompatible with
//! standalone FIPS 204 verifiers. This is intentional — lo signatures are
//! always verified by soliton (or a reimplementation built from Specification.md),
//! never by a generic FIPS 204 implementation. The same approach is used for
//! the X-Wing combiner (§4).

use crate::error::{Error, Result};
use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};

/// ML-DSA-65 public key (1952 bytes).
#[derive(Clone, PartialEq, Eq)]
pub struct PublicKey(pub(crate) Vec<u8>);

/// ML-DSA-65 secret key (32-byte seed, FIPS 204 §6.1 `ξ`).
///
/// The seed is sufficient to deterministically reconstruct the full signing key
/// via `SigningKey::from_seed`. Storing only the seed reduces key material
/// exposure compared to the ~4032-byte expanded form.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct SecretKey(pub(crate) Vec<u8>);

/// ML-DSA-65 signature (3309 bytes, FIPS 204).
#[derive(Clone, PartialEq, Eq)]
pub struct Signature(pub(crate) Vec<u8>);

/// ML-DSA-65 public key length (bytes).
const PK_LEN: usize = 1952;
/// ML-DSA-65 secret key length (bytes, seed form per FIPS 204 §6.1).
const SK_LEN: usize = 32;
/// ML-DSA-65 signature length (bytes).
const SIG_LEN: usize = 3309;

impl PublicKey {
    /// Return the raw byte representation.
    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }

    /// Construct from raw bytes with size validation.
    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self> {
        if bytes.len() != PK_LEN {
            return Err(Error::InvalidLength {
                expected: PK_LEN,
                got: bytes.len(),
            });
        }
        Ok(Self(bytes))
    }

    /// Construct from raw bytes without size validation (internal trusted use).
    pub(crate) fn from_bytes_unchecked(bytes: Vec<u8>) -> Self {
        assert_eq!(
            bytes.len(),
            PK_LEN,
            "mldsa::PublicKey::from_bytes_unchecked: wrong size"
        );
        Self(bytes)
    }
}

impl SecretKey {
    /// Return the raw byte representation.
    pub(crate) fn as_bytes(&self) -> &[u8] {
        &self.0
    }

    /// Construct from raw bytes with size validation.
    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self> {
        // Wrap in Zeroizing so the raw Vec is zeroized on the error path
        // (SecretKey derives ZeroizeOnDrop, but that only fires on success).
        let mut bytes = Zeroizing::new(bytes);
        if bytes.len() != SK_LEN {
            return Err(Error::InvalidLength {
                expected: SK_LEN,
                got: bytes.len(),
            });
        }
        Ok(Self(std::mem::take(&mut *bytes)))
    }
}

impl Signature {
    /// Return the raw byte representation.
    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }

    /// Construct from raw bytes with size validation.
    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self> {
        if bytes.len() != SIG_LEN {
            return Err(Error::InvalidLength {
                expected: SIG_LEN,
                got: bytes.len(),
            });
        }
        Ok(Self(bytes))
    }

    /// Construct from raw bytes without size validation (internal trusted use).
    pub(crate) fn from_bytes_unchecked(bytes: Vec<u8>) -> Self {
        assert_eq!(
            bytes.len(),
            SIG_LEN,
            "mldsa::Signature::from_bytes_unchecked: wrong size"
        );
        Self(bytes)
    }
}

/// Return the expected public key length in bytes.
pub const fn pk_len() -> usize {
    PK_LEN
}

/// Return the expected secret key length in bytes (32-byte seed).
pub const fn sk_len() -> usize {
    SK_LEN
}

/// Return the expected signature length in bytes.
pub const fn sig_len() -> usize {
    SIG_LEN
}

/// Generate an ML-DSA-65 keypair.
///
/// Uses `getrandom` as the entropy source, passed to the deterministic
/// `from_seed` API. Equivalent to ML-DSA.KeyGen() from FIPS 204 §6.1.
///
/// The secret key is the 32-byte seed `ξ`, not the expanded form.
///
/// # Security
///
/// The 32-byte seed and its `B32` copy are zeroized immediately after key
/// generation. The extracted seed is wrapped in `Zeroizing` before being
/// moved into `SecretKey` (which derives `ZeroizeOnDrop`).
pub fn keygen() -> Result<(PublicKey, SecretKey)> {
    use ml_dsa::{B32, KeyGen, MlDsa65};

    // Generate 32-byte seed from OS CSPRNG.
    let mut seed_bytes = [0u8; 32];
    super::random::random_bytes(&mut seed_bytes);
    let mut seed: B32 = seed_bytes.into();
    // [u8; 32] is Copy — B32::from() received a bitwise copy, so the
    // original stack value must be explicitly zeroized.
    seed_bytes.zeroize();

    let kp = MlDsa65::from_seed(&seed);
    seed.zeroize();

    let pk_bytes = kp.verifying_key().encode().to_vec();
    // B32 (hybrid_array::Array<u8, U32>) implements Zeroize but not ZeroizeOnDrop.
    // Bind the return value to explicitly zeroize the seed copy.
    let mut seed_out = kp.to_seed();
    let mut sk_bytes = zeroize::Zeroizing::new(seed_out.to_vec());
    seed_out.zeroize();

    // Runtime asserts (not debug_assert) because these sizes come from a foreign
    // crate's type-level constants — a crate update could silently change them.
    assert_eq!(
        pk_bytes.len(),
        PK_LEN,
        "ML-DSA-65 VK size mismatch: got {}, expected {PK_LEN}",
        pk_bytes.len()
    );
    assert_eq!(
        sk_bytes.len(),
        SK_LEN,
        "ML-DSA-65 seed size mismatch: got {}, expected {SK_LEN}",
        sk_bytes.len()
    );

    Ok((
        PublicKey(pk_bytes),
        SecretKey(std::mem::take(&mut *sk_bytes)),
    ))
}

/// Sign a message with ML-DSA-65 (hedged mode).
///
/// Re-expands the signing key from the stored 32-byte seed, generates 32
/// bytes of fresh entropy from `getrandom`, and signs via `sign_internal`
/// with the entropy as the `rnd` parameter. See module doc for rationale.
///
/// # Security
///
/// The expanded `SigningKey` is zeroized on drop (`ml-dsa` feature `zeroize`).
/// The 32-byte seed is borrowed from `SecretKey`, not copied. The `rnd`
/// buffer is zeroized after use.
pub fn sign(sk: &SecretKey, message: &[u8]) -> Result<Signature> {
    use ml_dsa::{B32, MlDsa65, Seed, SigningKey};

    if sk.0.len() != SK_LEN {
        return Err(Error::InvalidLength {
            expected: SK_LEN,
            got: sk.0.len(),
        });
    }

    // ml-dsa requires its own typed key — raw &[u8] cannot be passed directly.
    // ZeroizeOnDrop guaranteed by ml-dsa features = ["zeroize"].
    let seed: &Seed = sk.0.as_slice().try_into().map_err(|_| Error::Internal)?;
    let signing_key = SigningKey::<MlDsa65>::from_seed(seed);

    // Hedged signing: 32 bytes of fresh getrandom entropy mixed into the
    // nonce derivation (FIPS 204 Algorithm 2, `rnd` parameter). Provides
    // fault-injection resistance — see module doc.
    let mut rnd_bytes = [0u8; 32];
    super::random::random_bytes(&mut rnd_bytes);
    let mut rnd: B32 = rnd_bytes.into();
    // [u8; 32] is Copy — .into() received a bitwise copy.
    rnd_bytes.zeroize();

    // sign_internal skips the FIPS 204 domain separator / context string.
    // Paired with verify_internal in verify() — both use MuBuilder::internal,
    // so signatures are self-consistent within soliton.
    let sig = signing_key.sign_internal(&[message], &rnd);
    // B32 (hybrid_array::Array<u8, U32>) implements Zeroize but not ZeroizeOnDrop.
    // Explicitly zeroize the hedged randomness after signing — same pattern as
    // seed zeroization in keygen().
    rnd.zeroize();

    let sig_bytes = sig.encode().to_vec();
    // Runtime assert (not debug_assert) because SIG_LEN comes from a foreign
    // crate's type-level constant — a crate update could silently change it.
    assert_eq!(
        sig_bytes.len(),
        SIG_LEN,
        "ML-DSA-65 sig size mismatch: got {}, expected {SIG_LEN}",
        sig_bytes.len()
    );

    Ok(Signature(sig_bytes))
}

/// Verify an ML-DSA-65 signature.
///
/// Returns `Ok(())` on valid signature, `Err(VerificationFailed)` otherwise.
pub fn verify(pk: &PublicKey, message: &[u8], signature: &Signature) -> Result<()> {
    use ml_dsa::{EncodedSignature, EncodedVerifyingKey, MlDsa65, VerifyingKey};

    if pk.0.len() != PK_LEN {
        return Err(Error::InvalidLength {
            expected: PK_LEN,
            got: pk.0.len(),
        });
    }
    if signature.0.len() != SIG_LEN {
        return Err(Error::InvalidLength {
            expected: SIG_LEN,
            got: signature.0.len(),
        });
    }

    // ml-dsa requires its own typed key — raw &[u8] cannot be passed directly.
    let vk_enc: &EncodedVerifyingKey<MlDsa65> =
        pk.0.as_slice().try_into().map_err(|_| Error::Internal)?;
    // Infallible: decode() from a correctly-sized EncodedVerifyingKey always
    // succeeds — the encoding is a fixed-length byte array, not a parsed format.
    let vk = VerifyingKey::<MlDsa65>::decode(vk_enc);

    // ml-dsa requires its own typed signature — raw &[u8] cannot be passed directly.
    let sig_enc: &EncodedSignature<MlDsa65> = signature
        .0
        .as_slice()
        .try_into()
        .map_err(|_| Error::Internal)?;
    // VerificationFailed (not InvalidData) — avoids leaking whether the
    // signature was malformed vs mathematically incorrect to the caller.
    let sig = ml_dsa::Signature::<MlDsa65>::decode(sig_enc).ok_or(Error::VerificationFailed)?;

    // verify_internal skips the FIPS 204 domain separator / context string.
    // Paired with sign_internal in sign() — both use MuBuilder::internal.
    // sign_internal takes &[&[u8]] (slice-of-slices); verify_internal takes
    // &[u8] (flat slice). The asymmetry is safe: SHA3's streaming absorb
    // produces identical hashes regardless of chunking boundaries.
    if vk.verify_internal(message, &sig) {
        Ok(())
    } else {
        Err(Error::VerificationFailed)
    }
}

/// Sign with the FIPS 204 §5.2 ML-DSA.Sign envelope (domain separator +
/// context string).
///
/// This is the **standard ML-DSA.Sign** algorithm — equivalent to what
/// any FIPS 204-compliant implementation produces (rustls/aws-lc-rs,
/// OpenSSL once they ship PQ, hardware tokens). Use it whenever a
/// signature must verify under a generic ML-DSA-65 verifier.
///
/// Constructs `M' = 0x00 || len(ctx) || ctx || M` per FIPS 204 §5.2,
/// then signs M' via the underlying primitive. Distinct from
/// [`sign`], which signs M directly (no prefix) — soliton's
/// non-standard mode used by intra-Lo signatures that are always
/// verified by soliton itself.
///
/// `ctx` is the FIPS 204 context string. For TLS 1.3 use it MUST be
/// empty (`b""`) per `draft-ietf-tls-mldsa`. Other protocol-specific
/// contexts (e.g., `b"lo-session-bind"` for application-level domain
/// separation) are supported up to 255 bytes.
///
/// Hedged signing: 32 bytes of fresh OS entropy mixed into the nonce
/// derivation, identical to [`sign`]'s hedged behavior.
///
/// # Errors
/// - [`Error::InvalidData`] if `ctx.len() > 255`
/// - [`Error::InvalidLength`] if the secret-key seed length is wrong
pub fn sign_with_ctx(sk: &SecretKey, message: &[u8], ctx: &[u8]) -> Result<Signature> {
    use ml_dsa::{B32, MlDsa65, Seed, SigningKey};

    if ctx.len() > 255 {
        return Err(Error::InvalidData);
    }
    if sk.0.len() != SK_LEN {
        return Err(Error::InvalidLength {
            expected: SK_LEN,
            got: sk.0.len(),
        });
    }

    let seed: &Seed = sk.0.as_slice().try_into().map_err(|_| Error::Internal)?;
    let signing_key = SigningKey::<MlDsa65>::from_seed(seed);

    // Hedged signing: 32 bytes of fresh getrandom entropy. Same
    // rationale as `sign()` — fault-injection resistance.
    let mut rnd_bytes = [0u8; 32];
    super::random::random_bytes(&mut rnd_bytes);
    let mut rnd: B32 = rnd_bytes.into();
    rnd_bytes.zeroize();

    // FIPS 204 §5.2 envelope: M' = 0x00 || len(ctx) || ctx || M.
    // SHAKE's streaming absorb is associative, so passing M' as a
    // multi-chunk slice (`&[domain_sep, len_byte, ctx, message]`)
    // produces the same μ as concatenating into one buffer first.
    // The `MuBuilder::new(tr, ctx).message(M)` path used by the
    // ml-dsa crate's standard `verify_with_context` absorbs the
    // same bytes in the same order — so signatures from this
    // function verify under the standard verify path, AND the
    // standard sign path verifies under our `verify_with_ctx`.
    let domain_sep = [0x00u8];
    // `ctx.len() > 255` was rejected above, so this cannot fail.
    let len_byte = [u8::try_from(ctx.len()).map_err(|_| Error::Internal)?];
    let sig = signing_key.sign_internal(&[&domain_sep, &len_byte, ctx, message], &rnd);
    rnd.zeroize();

    let sig_bytes = sig.encode().to_vec();
    assert_eq!(
        sig_bytes.len(),
        SIG_LEN,
        "ML-DSA-65 sig size mismatch: got {}, expected {SIG_LEN}",
        sig_bytes.len()
    );

    Ok(Signature(sig_bytes))
}

/// Verify with the FIPS 204 §5.2 ML-DSA.Verify envelope.
///
/// Counterpart of [`sign_with_ctx`]. Accepts signatures from any
/// FIPS 204-compliant ML-DSA-65 implementation (when called with
/// the same `ctx`).
///
/// `ctx` MUST match the context string used by the signer. For TLS
/// 1.3 use that is `b""`.
///
/// Returns `Ok(())` on a valid signature, `Err(VerificationFailed)`
/// for an invalid one or for `ctx.len() > 255`.
pub fn verify_with_ctx(
    pk: &PublicKey,
    message: &[u8],
    ctx: &[u8],
    signature: &Signature,
) -> Result<()> {
    use ml_dsa::{EncodedSignature, EncodedVerifyingKey, MlDsa65, VerifyingKey};

    if ctx.len() > 255 {
        return Err(Error::VerificationFailed);
    }
    if pk.0.len() != PK_LEN {
        return Err(Error::InvalidLength {
            expected: PK_LEN,
            got: pk.0.len(),
        });
    }
    if signature.0.len() != SIG_LEN {
        return Err(Error::InvalidLength {
            expected: SIG_LEN,
            got: signature.0.len(),
        });
    }

    let vk_enc: &EncodedVerifyingKey<MlDsa65> =
        pk.0.as_slice().try_into().map_err(|_| Error::Internal)?;
    let vk = VerifyingKey::<MlDsa65>::decode(vk_enc);

    let sig_enc: &EncodedSignature<MlDsa65> = signature
        .0
        .as_slice()
        .try_into()
        .map_err(|_| Error::Internal)?;
    let sig = ml_dsa::Signature::<MlDsa65>::decode(sig_enc).ok_or(Error::VerificationFailed)?;

    // `verify_with_context` applies the FIPS 204 §5.2 envelope
    // (`0x00 || len(ctx) || ctx || M`) before verifying, matching
    // the construction in `sign_with_ctx`.
    if vk.verify_with_context(message, ctx, &sig) {
        Ok(())
    } else {
        Err(Error::VerificationFailed)
    }
}

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

    #[test]
    fn keygen_sizes() {
        let (pk, sk) = keygen().unwrap();
        assert_eq!(pk.as_bytes().len(), 1952);
        assert_eq!(sk.as_bytes().len(), 32);
    }

    #[test]
    fn sign_verify_round_trip() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign(&sk, b"hello mldsa").unwrap();
        assert!(verify(&pk, b"hello mldsa", &sig).is_ok());
    }

    #[test]
    fn sign_size() {
        let (_, sk) = keygen().unwrap();
        let sig = sign(&sk, b"test").unwrap();
        assert_eq!(sig.as_bytes().len(), 3309);
    }

    #[test]
    fn sign_verify_empty_message() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign(&sk, b"").unwrap();
        assert!(verify(&pk, b"", &sig).is_ok());
    }

    /// `sign_with_ctx` round-trip with empty context (the form used
    /// for TLS 1.3 cert verify per draft-ietf-tls-mldsa).
    #[test]
    fn sign_with_ctx_empty_round_trip() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign_with_ctx(&sk, b"hello fips204", b"").unwrap();
        assert!(verify_with_ctx(&pk, b"hello fips204", b"", &sig).is_ok());
    }

    /// Round-trip with a non-empty context. Confirms the FIPS 204
    /// envelope `0x00 || len(ctx) || ctx || M` round-trips through
    /// the sign/verify pair.
    #[test]
    fn sign_with_ctx_nonempty_round_trip() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign_with_ctx(&sk, b"payload", b"app-domain-v1").unwrap();
        assert!(verify_with_ctx(&pk, b"payload", b"app-domain-v1", &sig).is_ok());
    }

    /// Domain-separation: a signature produced with one context
    /// MUST NOT verify under a different context. This is the
    /// security property that makes per-protocol context strings
    /// useful — a TLS-bound signature can't be replayed to an
    /// application-layer verifier and vice versa.
    #[test]
    fn sign_with_ctx_different_ctx_fails() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign_with_ctx(&sk, b"payload", b"ctx-a").unwrap();
        assert!(matches!(
            verify_with_ctx(&pk, b"payload", b"ctx-b", &sig),
            Err(Error::VerificationFailed)
        ));
    }

    /// `sign` (no ctx) and `sign_with_ctx` produce DIFFERENT
    /// signatures, and the signatures don't cross-verify. This is
    /// the deliberate distinction: `sign` skips the FIPS 204
    /// envelope, `sign_with_ctx` applies it. If a regression ever
    /// made them equivalent, soliton's intra-Lo signatures would
    /// silently become FIPS 204-compatible (or vice versa) — both
    /// outcomes need to be intentional, not accidental.
    #[test]
    fn sign_and_sign_with_ctx_are_distinct() {
        let (pk, sk) = keygen().unwrap();
        let sig_internal = sign(&sk, b"payload").unwrap();
        let sig_standard = sign_with_ctx(&sk, b"payload", b"").unwrap();
        // Bytes differ (the rnd hedging makes equality unreliable
        // even for identical inputs, so this isn't a strict check).
        // The cross-verify check is the load-bearing one:
        assert!(
            matches!(
                verify_with_ctx(&pk, b"payload", b"", &sig_internal),
                Err(Error::VerificationFailed)
            ),
            "internal-mode signature must NOT verify under FIPS 204 verify"
        );
        assert!(
            matches!(
                verify(&pk, b"payload", &sig_standard),
                Err(Error::VerificationFailed)
            ),
            "FIPS 204 signature must NOT verify under internal-mode verify"
        );
    }

    /// 255-byte context is the spec maximum. Confirm it's accepted.
    #[test]
    fn sign_with_ctx_max_length_ok() {
        let (pk, sk) = keygen().unwrap();
        let ctx = vec![0xAB; 255];
        let sig = sign_with_ctx(&sk, b"payload", &ctx).unwrap();
        assert!(verify_with_ctx(&pk, b"payload", &ctx, &sig).is_ok());
    }

    /// 256-byte context exceeds the FIPS 204 §5.2 limit and MUST
    /// be rejected at the API boundary rather than producing an
    /// invalid signature.
    #[test]
    fn sign_with_ctx_oversized_ctx_rejected() {
        let (_, sk) = keygen().unwrap();
        let ctx = vec![0u8; 256];
        assert!(matches!(
            sign_with_ctx(&sk, b"payload", &ctx),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn verify_with_ctx_oversized_ctx_rejected() {
        let (pk, sk) = keygen().unwrap();
        let valid_sig = sign_with_ctx(&sk, b"payload", b"").unwrap();
        let too_long_ctx = vec![0u8; 256];
        assert!(matches!(
            verify_with_ctx(&pk, b"payload", &too_long_ctx, &valid_sig),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn verify_wrong_message() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign(&sk, b"message one").unwrap();
        assert!(matches!(
            verify(&pk, b"message two", &sig),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn verify_wrong_key() {
        let (_, sk) = keygen().unwrap();
        let (pk2, _) = keygen().unwrap();
        let sig = sign(&sk, b"test").unwrap();
        assert!(matches!(
            verify(&pk2, b"test", &sig),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn verify_tampered_signature() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign(&sk, b"test").unwrap();
        let mut bad = sig.as_bytes().to_vec();
        bad[0] ^= 0xFF;
        let bad_sig = Signature::from_bytes(bad).unwrap();
        assert!(matches!(
            verify(&pk, b"test", &bad_sig),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn verify_all_zeros_signature() {
        let (pk, _) = keygen().unwrap();
        let zeros_sig = Signature::from_bytes(vec![0u8; 3309]).unwrap();
        assert!(matches!(
            verify(&pk, b"test", &zeros_sig),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn pk_from_bytes_wrong_size() {
        assert!(matches!(
            PublicKey::from_bytes(vec![0u8; 100]),
            Err(Error::InvalidLength {
                expected: 1952,
                got: 100
            })
        ));
    }

    #[test]
    fn sig_from_bytes_wrong_size() {
        assert!(matches!(
            Signature::from_bytes(vec![0u8; 100]),
            Err(Error::InvalidLength {
                expected: 3309,
                got: 100
            })
        ));
    }

    #[test]
    fn sk_from_bytes_wrong_size() {
        assert!(matches!(
            SecretKey::from_bytes(vec![0u8; 100]),
            Err(Error::InvalidLength {
                expected: 32,
                got: 100
            })
        ));
    }

    #[test]
    fn sign_hedged_both_verify() {
        let (pk, sk) = keygen().unwrap();
        let msg = b"hedged signing test";
        let sig1 = sign(&sk, msg).unwrap();
        let sig2 = sign(&sk, msg).unwrap();
        // Hedged signing: same message produces different signatures (fresh
        // rnd each time), but both must verify.
        assert_ne!(
            sig1.as_bytes(),
            sig2.as_bytes(),
            "hedged signatures on the same message must differ"
        );
        assert!(verify(&pk, msg, &sig1).is_ok());
        assert!(verify(&pk, msg, &sig2).is_ok());
    }

    proptest::proptest! {
        // Default case count (256) — ML-DSA keygen is computationally expensive;
        // 256 iterations provide adequate coverage without excessive test time.
        #[test]
        fn proptest_round_trip(
            msg in proptest::collection::vec(proptest::prelude::any::<u8>(), 0..1024),
            flip_byte in 0..3309usize,
        ) {
            let kg = keygen();
            proptest::prop_assert!(kg.is_ok());
            let (pk, sk) = kg.unwrap();
            let sign_result = sign(&sk, &msg);
            proptest::prop_assert!(sign_result.is_ok());
            let sig = sign_result.unwrap();
            proptest::prop_assert!(verify(&pk, &msg, &sig).is_ok());

            // Forgery rejection: mutated signature must fail verification.
            let mut bad_sig_bytes = sig.as_bytes().to_vec();
            bad_sig_bytes[flip_byte] ^= 0x01;
            let bad_sig = Signature::from_bytes(bad_sig_bytes).unwrap();
            proptest::prop_assert!(verify(&pk, &msg, &bad_sig).is_err());

            // Forgery rejection: mutated message must fail verification.
            if !msg.is_empty() {
                let mut bad_msg = msg.clone();
                bad_msg[0] ^= 0x01;
                proptest::prop_assert!(verify(&pk, &bad_msg, &sig).is_err());
            }
        }
    }
}