//! Post-quantum hybrid KEM: X25519 + ML-KEM-768.
//!
//! Wraps MLS payloads in an outer encryption layer using a hybrid key
//! encapsulation mechanism.  The X25519 component provides classical
//! ECDH security; the ML-KEM-768 component (FIPS 203) provides
//! post-quantum security.
//!
//! # Wire format
//!
//! ```text
//! version(1) | x25519_eph_pk(32) | mlkem_ct(1088) | aead_nonce(12) | aead_ct(var)
//! ```
//!
//! # Key derivation
//!
//! ```text
//! ikm  = X25519_shared(32) || ML-KEM_shared(32)
//! key  = HKDF-SHA256(salt=[], ikm, info="quicnprotochat-hybrid-v1", L=32)
//! ```

use chacha20poly1305::{
    aead::{Aead, KeyInit},
    ChaCha20Poly1305, Key, Nonce,
};
use hkdf::Hkdf;
use ml_kem::{
    array::Array,
    kem::{Decapsulate, Encapsulate},
    EncodedSizeUser, KemCore, MlKem768, MlKem768Params,
};
use rand::{rngs::OsRng, rngs::StdRng, CryptoRng, RngCore, SeedableRng};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use x25519_dalek::{EphemeralSecret, PublicKey as X25519Public, StaticSecret};
use zeroize::Zeroizing;

// Re-import the concrete key types from the kem sub-module.
use ml_kem::kem::{DecapsulationKey, EncapsulationKey};

/// Current hybrid envelope version byte.
const HYBRID_VERSION: u8 = 0x01;

/// HKDF info string for domain separation.
/// Frozen at the original project name for backward compatibility with existing
/// encrypted state files and messages. Do not change.
const HKDF_INFO: &[u8] = b"quicnprotochat-hybrid-v1";

/// HKDF salt for domain separation (defence-in-depth; IKM already has 64 bytes of entropy).
/// Frozen — see [`HKDF_INFO`].
const HKDF_SALT: &[u8] = b"quicnprotochat-hybrid-v1-salt";

/// ML-KEM-768 ciphertext size in bytes.
const MLKEM_CT_LEN: usize = 1088;

/// ML-KEM-768 encapsulation key size in bytes.
pub const MLKEM_EK_LEN: usize = 1184;

/// ML-KEM-768 decapsulation key size in bytes.
pub const MLKEM_DK_LEN: usize = 2400;

/// Envelope header: version(1) + x25519 eph pk(32) + mlkem ct(1088) + nonce(12).
const HEADER_LEN: usize = 1 + 32 + MLKEM_CT_LEN + 12;

/// KEM output length (version + x25519 eph pk + mlkem ct) for HPKE adapter.
pub const HYBRID_KEM_OUTPUT_LEN: usize = 1 + 32 + MLKEM_CT_LEN;

/// Hybrid public key length: x25519(32) + mlkem_ek(1184). Used to detect hybrid keys in MLS.
pub const HYBRID_PUBLIC_KEY_LEN: usize = 32 + MLKEM_EK_LEN;

/// Hybrid private key length: x25519(32) + mlkem_dk(2400). Used to detect hybrid keys in MLS.
pub const HYBRID_PRIVATE_KEY_LEN: usize = 32 + MLKEM_DK_LEN;

// ── Error type ──────────────────────────────────────────────────────────────

#[derive(Debug, thiserror::Error)]
pub enum HybridKemError {
    #[error("AEAD encryption failed")]
    EncryptionFailed,

    #[error("AEAD decryption failed (wrong recipient or tampered)")]
    DecryptionFailed,

    #[error("unsupported hybrid envelope version: {0}")]
    UnsupportedVersion(u8),

    #[error("envelope too short ({0} bytes, minimum {HEADER_LEN})")]
    TooShort(usize),

    #[error("invalid ML-KEM encapsulation key")]
    InvalidMlKemKey,

    #[error("ML-KEM decapsulation failed")]
    MlKemDecapsFailed,
}

// ── Keypair types ───────────────────────────────────────────────────────────

/// A hybrid keypair combining X25519 (classical) + ML-KEM-768 (post-quantum).
///
/// Each peer holds one of these.  The public portion is distributed so
/// senders can encrypt payloads with post-quantum protection.
pub struct HybridKeypair {
    x25519_sk: StaticSecret,
    x25519_pk: X25519Public,
    mlkem_dk: DecapsulationKey<MlKem768Params>,
    mlkem_ek: EncapsulationKey<MlKem768Params>,
}

/// Serialisable form of a [`HybridKeypair`] for persistence.
///
/// Secret fields are wrapped in [`Zeroizing`] so they are securely erased
/// when the struct is dropped.
#[derive(Serialize, Deserialize)]
pub struct HybridKeypairBytes {
    pub x25519_sk: Zeroizing<[u8; 32]>,
    pub mlkem_dk: Zeroizing<Vec<u8>>,
    pub mlkem_ek: Vec<u8>,
}

/// The public portion of a hybrid keypair, sent to peers.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct HybridPublicKey {
    pub x25519_pk: [u8; 32],
    pub mlkem_ek: Vec<u8>,
}

/// HKDF info for deriving HPKE keypair seed from IKM (MLS compatibility).
/// Frozen — see [`HKDF_INFO`].
const HKDF_INFO_HPKE_KEYPAIR: &[u8] = b"quicnprotochat-hybrid-hpke-keypair-v1";

impl HybridKeypair {
    /// Generate a fresh hybrid keypair from OS CSPRNG.
    pub fn generate() -> Self {
        Self::generate_from_rng(&mut OsRng)
    }

    /// Generate a hybrid keypair from a seeded RNG (deterministic).
    pub fn generate_from_rng<R: RngCore + CryptoRng>(rng: &mut R) -> Self {
        let x25519_sk = StaticSecret::random_from_rng(&mut *rng);
        let x25519_pk = X25519Public::from(&x25519_sk);
        let (mlkem_dk, mlkem_ek) = MlKem768::generate(rng);

        Self {
            x25519_sk,
            x25519_pk,
            mlkem_dk,
            mlkem_ek,
        }
    }

    /// Derive a deterministic hybrid keypair from IKM (for MLS HPKE key schedule).
    pub fn derive_from_ikm(ikm: &[u8]) -> Self {
        let mut seed = [0u8; 32];
        let hk = Hkdf::<Sha256>::new(None, ikm);
        hk.expand(HKDF_INFO_HPKE_KEYPAIR, &mut seed)
            .expect("32 bytes is valid HKDF output");
        let mut rng = StdRng::from_seed(seed);
        Self::generate_from_rng(&mut rng)
    }

    /// Serialise private key for MLS key store: x25519_sk(32) || mlkem_dk(2400).
    ///
    /// The returned value is wrapped in [`Zeroizing`] so secret key material
    /// is securely erased when dropped.
    pub fn private_to_bytes(&self) -> Zeroizing<Vec<u8>> {
        let mut out = Vec::with_capacity(HYBRID_PRIVATE_KEY_LEN);
        out.extend_from_slice(self.x25519_sk.as_bytes());
        out.extend_from_slice(self.mlkem_dk.as_bytes().as_slice());
        Zeroizing::new(out)
    }

    /// Reconstruct a hybrid keypair from private key bytes (from MLS key store).
    pub fn from_private_bytes(bytes: &[u8]) -> Result<Self, HybridKemError> {
        if bytes.len() != HYBRID_PRIVATE_KEY_LEN {
            return Err(HybridKemError::TooShort(bytes.len()));
        }
        let x25519_sk = StaticSecret::from(<[u8; 32]>::try_from(&bytes[0..32])
            .expect("slice is exactly 32 bytes (guaranteed by HYBRID_PRIVATE_KEY_LEN check)"));
        let x25519_pk = X25519Public::from(&x25519_sk);

        let mlkem_dk_arr = Array::try_from(&bytes[32..32 + MLKEM_DK_LEN])
            .map_err(|_| HybridKemError::InvalidMlKemKey)?;
        let mlkem_dk = DecapsulationKey::<MlKem768Params>::from_bytes(&mlkem_dk_arr);
        let mlkem_ek = mlkem_dk.encapsulation_key().clone();

        Ok(Self {
            x25519_sk,
            x25519_pk,
            mlkem_dk,
            mlkem_ek,
        })
    }

    /// Reconstruct from serialised bytes.
    pub fn from_bytes(bytes: &HybridKeypairBytes) -> Result<Self, HybridKemError> {
        let x25519_sk = StaticSecret::from(*bytes.x25519_sk);
        let x25519_pk = X25519Public::from(&x25519_sk);

        let mlkem_dk_arr = Array::try_from(bytes.mlkem_dk.as_slice())
            .map_err(|_| HybridKemError::InvalidMlKemKey)?;
        let mlkem_dk = DecapsulationKey::<MlKem768Params>::from_bytes(&mlkem_dk_arr);

        let mlkem_ek_arr = Array::try_from(bytes.mlkem_ek.as_slice())
            .map_err(|_| HybridKemError::InvalidMlKemKey)?;
        let mlkem_ek = EncapsulationKey::<MlKem768Params>::from_bytes(&mlkem_ek_arr);

        Ok(Self {
            x25519_sk,
            x25519_pk,
            mlkem_dk,
            mlkem_ek,
        })
    }

    /// Serialise the keypair for persistence.
    pub fn to_bytes(&self) -> HybridKeypairBytes {
        HybridKeypairBytes {
            x25519_sk: Zeroizing::new(self.x25519_sk.to_bytes()),
            mlkem_dk: Zeroizing::new(self.mlkem_dk.as_bytes().to_vec()),
            mlkem_ek: self.mlkem_ek.as_bytes().to_vec(),
        }
    }

    /// Extract the public portion for distribution to peers.
    pub fn public_key(&self) -> HybridPublicKey {
        HybridPublicKey {
            x25519_pk: self.x25519_pk.to_bytes(),
            mlkem_ek: self.mlkem_ek.as_bytes().to_vec(),
        }
    }
}

impl HybridPublicKey {
    /// Serialise to a single byte blob: x25519_pk(32) || mlkem_ek(1184).
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut out = Vec::with_capacity(32 + self.mlkem_ek.len());
        out.extend_from_slice(&self.x25519_pk);
        out.extend_from_slice(&self.mlkem_ek);
        out
    }

    /// Deserialise from a single byte blob.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, HybridKemError> {
        if bytes.len() < 32 + MLKEM_EK_LEN {
            return Err(HybridKemError::TooShort(bytes.len()));
        }
        let mut x25519_pk = [0u8; 32];
        x25519_pk.copy_from_slice(&bytes[..32]);
        let mlkem_ek = bytes[32..32 + MLKEM_EK_LEN].to_vec();
        Ok(Self {
            x25519_pk,
            mlkem_ek,
        })
    }
}

// ── Encrypt / Decrypt ───────────────────────────────────────────────────────

/// Encrypt `plaintext` to `recipient_pk` using X25519 + ML-KEM-768 hybrid KEM.
///
/// `info` is optional HPKE context info incorporated into key derivation.
/// `aad` is optional additional authenticated data bound to the AEAD ciphertext.
///
/// Returns the complete hybrid envelope as a byte vector.
pub fn hybrid_encrypt(
    recipient_pk: &HybridPublicKey,
    plaintext: &[u8],
    info: &[u8],
    aad: &[u8],
) -> Result<Vec<u8>, HybridKemError> {
    // 1. Ephemeral X25519 DH
    let eph_secret = EphemeralSecret::random_from_rng(OsRng);
    let eph_public = X25519Public::from(&eph_secret);
    let x25519_recipient = X25519Public::from(recipient_pk.x25519_pk);
    let x25519_ss = eph_secret.diffie_hellman(&x25519_recipient);

    // 2. ML-KEM-768 encapsulation
    let mlkem_ek_arr = Array::try_from(recipient_pk.mlkem_ek.as_slice())
        .map_err(|_| HybridKemError::InvalidMlKemKey)?;
    let mlkem_ek = EncapsulationKey::<MlKem768Params>::from_bytes(&mlkem_ek_arr);
    let (mlkem_ct, mlkem_ss) = mlkem_ek
        .encapsulate(&mut OsRng)
        .map_err(|_| HybridKemError::EncryptionFailed)?;

    // 3. Derive AEAD key from combined shared secrets (with caller info for context binding)
    let aead_key = derive_aead_key(x25519_ss.as_bytes(), mlkem_ss.as_slice(), info);

    // Generate a random 12-byte nonce (not derived from HKDF).
    let mut nonce_bytes = [0u8; 12];
    OsRng.fill_bytes(&mut nonce_bytes);
    let aead_nonce = *Nonce::from_slice(&nonce_bytes);

    // 4. AEAD encrypt with caller-supplied AAD
    let cipher = ChaCha20Poly1305::new(&aead_key);
    let aead_payload = chacha20poly1305::aead::Payload { msg: plaintext, aad };
    let ct = cipher
        .encrypt(&aead_nonce, aead_payload)
        .map_err(|_| HybridKemError::EncryptionFailed)?;

    // 5. Assemble envelope: version || x25519_eph_pk || mlkem_ct || nonce || aead_ct
    let mut out = Vec::with_capacity(HEADER_LEN + ct.len());
    out.push(HYBRID_VERSION);
    out.extend_from_slice(&eph_public.to_bytes());
    out.extend_from_slice(mlkem_ct.as_slice());
    out.extend_from_slice(aead_nonce.as_slice());
    out.extend_from_slice(&ct);

    Ok(out)
}

/// Decrypt a hybrid envelope using the recipient's private key.
///
/// `info` and `aad` must match what was passed to `hybrid_encrypt`.
pub fn hybrid_decrypt(
    keypair: &HybridKeypair,
    envelope: &[u8],
    info: &[u8],
    aad: &[u8],
) -> Result<Vec<u8>, HybridKemError> {
    if envelope.len() < HEADER_LEN + 16 {
        // 16 = minimum AEAD tag
        return Err(HybridKemError::TooShort(envelope.len()));
    }

    let version = envelope[0];
    if version != HYBRID_VERSION {
        return Err(HybridKemError::UnsupportedVersion(version));
    }

    let mut cursor = 1;

    // X25519 ephemeral public key
    let mut eph_pk_bytes = [0u8; 32];
    eph_pk_bytes.copy_from_slice(&envelope[cursor..cursor + 32]);
    cursor += 32;

    // ML-KEM ciphertext
    let mlkem_ct_bytes = &envelope[cursor..cursor + MLKEM_CT_LEN];
    cursor += MLKEM_CT_LEN;

    // AEAD nonce
    let nonce = Nonce::from_slice(&envelope[cursor..cursor + 12]);
    cursor += 12;

    // AEAD ciphertext
    let aead_ct = &envelope[cursor..];

    // 1. X25519 DH with ephemeral public key
    let eph_pk = X25519Public::from(eph_pk_bytes);
    let x25519_ss = keypair.x25519_sk.diffie_hellman(&eph_pk);

    // 2. ML-KEM decapsulation — convert bytes to the ciphertext array type
    //    that `DecapsulationKey::decapsulate` expects.
    let mlkem_ct_arr =
        Array::try_from(mlkem_ct_bytes).map_err(|_| HybridKemError::MlKemDecapsFailed)?;
    let mlkem_ss = keypair
        .mlkem_dk
        .decapsulate(&mlkem_ct_arr)
        .map_err(|_| HybridKemError::MlKemDecapsFailed)?;

    // 3. Derive AEAD key (with caller info for context binding)
    let aead_key = derive_aead_key(x25519_ss.as_bytes(), mlkem_ss.as_slice(), info);

    // 4. Decrypt with caller-supplied AAD
    let cipher = ChaCha20Poly1305::new(&aead_key);
    let aead_payload = chacha20poly1305::aead::Payload { msg: aead_ct, aad };
    let plaintext = cipher
        .decrypt(nonce, aead_payload)
        .map_err(|_| HybridKemError::DecryptionFailed)?;

    Ok(plaintext)
}

/// Encapsulate only: compute shared secret and KEM output (no AEAD).
/// Returns `(kem_output, shared_secret)` where `kem_output` is the first
/// `HYBRID_KEM_OUTPUT_LEN` bytes of the hybrid envelope and `shared_secret`
/// is the 32-byte derived key (same as used for AEAD in `hybrid_encrypt`).
/// Used by MLS HPKE exporter (setup_sender_and_export).
pub fn hybrid_encapsulate_only(
    recipient_pk: &HybridPublicKey,
) -> Result<(Vec<u8>, [u8; 32]), HybridKemError> {
    let eph_secret = EphemeralSecret::random_from_rng(OsRng);
    let eph_public = X25519Public::from(&eph_secret);
    let x25519_recipient = X25519Public::from(recipient_pk.x25519_pk);
    let x25519_ss = eph_secret.diffie_hellman(&x25519_recipient);

    let mlkem_ek_arr = Array::try_from(recipient_pk.mlkem_ek.as_slice())
        .map_err(|_| HybridKemError::InvalidMlKemKey)?;
    let mlkem_ek = EncapsulationKey::<MlKem768Params>::from_bytes(&mlkem_ek_arr);
    let (mlkem_ct, mlkem_ss) = mlkem_ek
        .encapsulate(&mut OsRng)
        .map_err(|_| HybridKemError::EncryptionFailed)?;

    let aead_key = derive_aead_key(x25519_ss.as_bytes(), mlkem_ss.as_slice(), b"");
    let shared_secret: [u8; 32] = aead_key.as_slice().try_into()
        .expect("AEAD key is always exactly 32 bytes");

    let mut kem_output = Vec::with_capacity(HYBRID_KEM_OUTPUT_LEN);
    kem_output.push(HYBRID_VERSION);
    kem_output.extend_from_slice(&eph_public.to_bytes());
    kem_output.extend_from_slice(mlkem_ct.as_slice());

    Ok((kem_output, shared_secret))
}

/// Decapsulate only: recover shared secret from KEM output (no AEAD).
/// Used by MLS HPKE exporter (setup_receiver_and_export).
pub fn hybrid_decapsulate_only(
    keypair: &HybridKeypair,
    kem_output: &[u8],
) -> Result<[u8; 32], HybridKemError> {
    if kem_output.len() < HYBRID_KEM_OUTPUT_LEN {
        return Err(HybridKemError::TooShort(kem_output.len()));
    }
    if kem_output[0] != HYBRID_VERSION {
        return Err(HybridKemError::UnsupportedVersion(kem_output[0]));
    }

    let eph_pk_bytes: [u8; 32] = kem_output[1..33].try_into()
        .expect("slice is exactly 32 bytes (guaranteed by HYBRID_KEM_OUTPUT_LEN check)");
    let eph_pk = X25519Public::from(eph_pk_bytes);
    let x25519_ss = keypair.x25519_sk.diffie_hellman(&eph_pk);

    let mlkem_ct_arr = Array::try_from(&kem_output[33..33 + MLKEM_CT_LEN])
        .map_err(|_| HybridKemError::MlKemDecapsFailed)?;
    let mlkem_ss = keypair
        .mlkem_dk
        .decapsulate(&mlkem_ct_arr)
        .map_err(|_| HybridKemError::MlKemDecapsFailed)?;

    let aead_key = derive_aead_key(x25519_ss.as_bytes(), mlkem_ss.as_slice(), b"");
    Ok(aead_key.as_slice().try_into()
        .expect("AEAD key is always exactly 32 bytes"))
}

/// Export a secret from shared secret (MLS HPKE exporter compatibility).
/// Uses HKDF-Expand(prk, exporter_context, length) with prk = HKDF-Extract(0, shared_secret).
pub fn hybrid_export(
    shared_secret: &[u8; 32],
    exporter_context: &[u8],
    length: usize,
) -> Vec<u8> {
    let hk = Hkdf::<Sha256>::new(Some(HKDF_SALT), shared_secret);
    let mut out = vec![0u8; length];
    hk.expand(exporter_context, &mut out).expect("valid length");
    out
}

/// Derive AEAD key from the combined X25519 + ML-KEM shared secrets.
///
/// `extra_info` is optional caller-supplied context (e.g. HPKE `info`) that is
/// appended to the domain-separation label for additional binding.
///
/// The nonce is generated randomly per-encryption rather than derived from
/// HKDF, preventing nonce reuse when the same shared secret is (accidentally)
/// used more than once.
fn derive_aead_key(x25519_ss: &[u8], mlkem_ss: &[u8], extra_info: &[u8]) -> Key {
    let mut ikm = Zeroizing::new(vec![0u8; x25519_ss.len() + mlkem_ss.len()]);
    ikm[..x25519_ss.len()].copy_from_slice(x25519_ss);
    ikm[x25519_ss.len()..].copy_from_slice(mlkem_ss);

    let hk = Hkdf::<Sha256>::new(Some(HKDF_SALT), &ikm);

    // Combine domain-separation label with caller-supplied context.
    let mut info = Vec::with_capacity(HKDF_INFO.len() + extra_info.len());
    info.extend_from_slice(HKDF_INFO);
    info.extend_from_slice(extra_info);

    let mut key_bytes = Zeroizing::new([0u8; 32]);
    hk.expand(&info, &mut *key_bytes)
        .expect("32 bytes is valid HKDF-SHA256 output length");

    *Key::from_slice(&*key_bytes)
}

// ── Tests ───────────────────────────────────────────────────────────────────

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

    #[test]
    fn keygen_produces_valid_public_key() {
        let kp = HybridKeypair::generate();
        let pk = kp.public_key();
        assert_eq!(pk.x25519_pk.len(), 32);
        assert_eq!(pk.mlkem_ek.len(), MLKEM_EK_LEN);
    }

    #[test]
    fn encrypt_decrypt_round_trip() {
        let kp = HybridKeypair::generate();
        let pk = kp.public_key();
        let plaintext = b"hello post-quantum world!";

        let envelope = hybrid_encrypt(&pk, plaintext, b"", b"").unwrap();
        let recovered = hybrid_decrypt(&kp, &envelope, b"", b"").unwrap();

        assert_eq!(recovered, plaintext);
    }

    #[test]
    fn encrypt_decrypt_with_info_aad() {
        let kp = HybridKeypair::generate();
        let pk = kp.public_key();
        let plaintext = b"context-bound payload";
        let info = b"mls epoch 42";
        let aad = b"group-id-abc";

        let envelope = hybrid_encrypt(&pk, plaintext, info, aad).unwrap();
        let recovered = hybrid_decrypt(&kp, &envelope, info, aad).unwrap();
        assert_eq!(recovered, plaintext);

        // Mismatched info must fail
        assert!(hybrid_decrypt(&kp, &envelope, b"wrong info", aad).is_err());
        // Mismatched aad must fail
        assert!(hybrid_decrypt(&kp, &envelope, info, b"wrong aad").is_err());
    }

    #[test]
    fn wrong_key_decryption_fails() {
        let kp_sender_target = HybridKeypair::generate();
        let kp_wrong = HybridKeypair::generate();

        let pk = kp_sender_target.public_key();
        let envelope = hybrid_encrypt(&pk, b"secret", b"", b"").unwrap();

        let result = hybrid_decrypt(&kp_wrong, &envelope, b"", b"");
        assert!(result.is_err());
    }

    #[test]
    fn tampered_aead_ciphertext_fails() {
        let kp = HybridKeypair::generate();
        let pk = kp.public_key();

        let mut envelope = hybrid_encrypt(&pk, b"payload", b"", b"").unwrap();
        let last = envelope.len() - 1;
        envelope[last] ^= 0x01;

        assert!(matches!(
            hybrid_decrypt(&kp, &envelope, b"", b""),
            Err(HybridKemError::DecryptionFailed)
        ));
    }

    #[test]
    fn tampered_mlkem_ct_fails() {
        let kp = HybridKeypair::generate();
        let pk = kp.public_key();

        let mut envelope = hybrid_encrypt(&pk, b"payload", b"", b"").unwrap();
        // Flip a byte in the ML-KEM ciphertext region (starts at offset 33)
        envelope[40] ^= 0xFF;

        assert!(hybrid_decrypt(&kp, &envelope, b"", b"").is_err());
    }

    #[test]
    fn tampered_x25519_eph_pk_fails() {
        let kp = HybridKeypair::generate();
        let pk = kp.public_key();

        let mut envelope = hybrid_encrypt(&pk, b"payload", b"", b"").unwrap();
        // Flip a byte in the X25519 ephemeral pk region (offset 1..33)
        envelope[5] ^= 0xFF;

        assert!(hybrid_decrypt(&kp, &envelope, b"", b"").is_err());
    }

    #[test]
    fn unsupported_version_rejected() {
        let kp = HybridKeypair::generate();
        let pk = kp.public_key();

        let mut envelope = hybrid_encrypt(&pk, b"payload", b"", b"").unwrap();
        envelope[0] = 0xFF;

        assert!(matches!(
            hybrid_decrypt(&kp, &envelope, b"", b""),
            Err(HybridKemError::UnsupportedVersion(0xFF))
        ));
    }

    #[test]
    fn envelope_too_short_rejected() {
        let kp = HybridKeypair::generate();
        assert!(matches!(
            hybrid_decrypt(&kp, &[0x01; 10], b"", b""),
            Err(HybridKemError::TooShort(10))
        ));
    }

    #[test]
    fn keypair_serialisation_round_trip() {
        let kp = HybridKeypair::generate();
        let bytes = kp.to_bytes();
        let restored = HybridKeypair::from_bytes(&bytes).unwrap();

        assert_eq!(kp.x25519_pk.to_bytes(), restored.x25519_pk.to_bytes());
        assert_eq!(kp.public_key().mlkem_ek, restored.public_key().mlkem_ek);

        // Verify restored keypair can decrypt
        let pk = kp.public_key();
        let ct = hybrid_encrypt(&pk, b"test", b"", b"").unwrap();
        let pt = hybrid_decrypt(&restored, &ct, b"", b"").unwrap();
        assert_eq!(pt, b"test");
    }

    #[test]
    fn public_key_serialisation_round_trip() {
        let kp = HybridKeypair::generate();
        let pk = kp.public_key();
        let bytes = pk.to_bytes();
        let restored = HybridPublicKey::from_bytes(&bytes).unwrap();

        assert_eq!(pk.x25519_pk, restored.x25519_pk);
        assert_eq!(pk.mlkem_ek, restored.mlkem_ek);
    }

    #[test]
    fn large_payload_round_trip() {
        let kp = HybridKeypair::generate();
        let pk = kp.public_key();
        let plaintext = vec![0xAB; 50_000]; // 50 KB

        let envelope = hybrid_encrypt(&pk, &plaintext, b"", b"").unwrap();
        let recovered = hybrid_decrypt(&kp, &envelope, b"", b"").unwrap();

        assert_eq!(recovered, plaintext);
    }
}