// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! HPKE-based wrapping of the global SecretSharing private key (SK_priv)
//! for delivery to recipient TEEs.
//!
//! Suite:
//!   - KEM  = DHKEM(X25519, HKDF-SHA256)   (RFC 9180)
//!   - KDF  = HKDF-SHA256
//!   - AEAD = ChaCha20-Poly1305 (preferred; AES-GCM possible with HW accel)
//!
//! ## Security Rationale
//!
//! **ChaCha20-Poly1305 vs AES-GCM**
//! - Constant-time, side-channel resistant
//! - No weak keys or IP concerns
//! - Consistent software performance across CPUs (esp. without AES-NI)
//! - Standardized in RFC 8439, widely deployed (TLS 1.3, WireGuard)
//!
//! **X25519 KEM**
//! - ~128-bit symmetric security
//! - Efficient ECDH with Montgomery ladder (resists timing attacks)
//! - Standard in modern protocols (TLS 1.3, Signal, Noise)
//!
//! ## Threat Model
//! - Adversary may observe ciphertexts but cannot compromise TEE execution
//! - Security relies on HPKE’s IND-CCA2 guarantees
//!
//! Base mode with AAD binding (request_id / key_id), using the
//! HPKE context API with a *detached* tag (we append the 16-byte tag to
//! the ciphertext so the wire format stays `ciphertext = body || tag`).

use std::sync::OnceLock;

use hpke::{
    aead::{AeadTag, ChaCha20Poly1305},
    kdf::HkdfSha256,
    kem::{Kem, X25519HkdfSha256},
    setup_receiver,
    setup_sender,
    Deserializable, // for from_bytes(...)
    OpModeR,
    OpModeS,
    Serializable, // for to_bytes()
};
use rialo_types::PublicKey;
use x25519_dalek::{PublicKey as X25519Pub, StaticSecret as X25519Priv};
use zeroize::Zeroizing;

use crate::{
    key::SecretSharingKey,
    traits::SecretSharingEngine,
    types::{Aad, KeyId, WrappedKeyEnvelope},
    SecretSharingError,
};

// One per process; not persisted across reboots.
static WRAPPING_PRIV_ONCE: OnceLock<[u8; 32]> = OnceLock::new();

/// Wraps/unwraps the SecretSharingKey using HPKE with this TEE’s X25519 wrapping keypair.
pub struct HpkeX25519ChaChaEngine {
    my_wrapping_priv: X25519Priv,
    my_wrapping_pub: X25519Pub,
}

impl Default for HpkeX25519ChaChaEngine {
    fn default() -> Self {
        Self::new()
    }
}

impl HpkeX25519ChaChaEngine {
    /// Construct an engine with a process-lifetime wrapping key.
    ///
    /// - First call in this process generates a new X25519 key (via HW CSPRNG).
    /// - Later calls reuse the same in-process key.
    /// - On reboot/restart, a new key is generated (no sealing/persistence).
    pub fn new() -> Self {
        let bytes = WRAPPING_PRIV_ONCE.get_or_init(|| {
            crate::hw_rng::bytes::<X25519_KEY_SIZE>().expect("HW RNG required in TEE")
        });
        let my_wrapping_priv = X25519Priv::from(*bytes); // dalek clamps internally
        let my_wrapping_pub = X25519Pub::from(&my_wrapping_priv);
        Self {
            my_wrapping_priv,
            my_wrapping_pub,
        }
    }

    /// Deterministic ctor for tests/fixtures.
    /// (Does not touch the global OnceLock.)
    #[cfg(any(test, feature = "testing"))]
    pub fn with_wrapping_priv(my_wrapping_priv: X25519Priv) -> Self {
        let my_wrapping_pub = X25519Pub::from(&my_wrapping_priv);
        Self {
            my_wrapping_priv,
            my_wrapping_pub,
        }
    }

    #[inline]
    fn my_hpke_priv_bytes(&self) -> [u8; 32] {
        self.my_wrapping_priv.to_bytes()
    }

    #[inline]
    fn my_wrapping_pub_bytes(&self) -> PublicKey {
        PublicKey::from(self.my_wrapping_pub)
    }

    /// HPKE seal: encrypt `sk_priv` for `recipient_wrapping_pub`, with associated AAD.
    ///
    /// Uses the context API with a detached tag. We then append the 16-byte tag
    /// to the 32-byte ciphertext body so callers see `ciphertext.len() == 48`.
    fn hpke_wrap(
        recipient_wrapping_pub: &PublicKey,
        sk_priv: &[u8; 32],
        aad: &[u8],
    ) -> Result<([u8; 32], [u8; 48]), SecretSharingError> {
        let pk_r =
            <X25519HkdfSha256 as Kem>::PublicKey::from_bytes(recipient_wrapping_pub.as_ref())
                .map_err(|_| {
                    SecretSharingError::EnvelopeInvalid("invalid recipient HPKE public key".into())
                })?;

        // HPKE RNG that enforces TEE HW-only policy (RDSEED→RDRAND).
        let mut rng = HpkeHwRng::default();

        // Context setup (sender)
        let (enc, mut sctx) = setup_sender::<ChaCha20Poly1305, HkdfSha256, X25519HkdfSha256, _>(
            &OpModeS::Base,
            &pk_r,
            SECRET_SHARING_HPKE_INFO,
            &mut rng,
        )
        .map_err(|e| SecretSharingError::WrapFailed(e.to_string()))?;

        // Encrypt 32-byte SK_priv; produce detached tag.
        let mut ct_body = *sk_priv; // X25519_KEY_SIZE bytes
        let tag = sctx
            .seal_in_place_detached(&mut ct_body, aad)
            .map_err(|e| SecretSharingError::WrapFailed(e.to_string()))?;

        // Append 16-byte tag so ciphertext = body || tag (48 bytes total).
        let mut ct_out = [0u8; 48];
        ct_out[..32].copy_from_slice(&ct_body);
        let tag_bytes = tag.to_bytes();
        #[allow(deprecated)]
        ct_out[32..].copy_from_slice(tag_bytes.as_slice());

        // Encapsulated key to fixed [u8; 32]
        let mut enc_bytes = [0u8; 32];
        enc_bytes.copy_from_slice(enc.to_bytes().as_ref());

        Ok((enc_bytes, ct_out))
    }

    /// HPKE open: decrypt an envelope and return the raw 32-byte SK_priv.
    ///
    /// Splits `ciphertext` into `[body (32)] || [tag (16)]` and verifies AAD.
    fn hpke_unwrap(
        my_wrapping_priv: &[u8; 32],
        enc: &[u8; 32],
        ct: &[u8; 48],
        aad: &[u8],
    ) -> Result<Zeroizing<[u8; 32]>, SecretSharingError> {
        // Split ciphertext into body(32) + tag(16)
        let (ct_body, tag_bytes) = ct.split_at(32);
        let tag = AeadTag::<ChaCha20Poly1305>::from_bytes(tag_bytes)
            .map_err(|_| SecretSharingError::EnvelopeInvalid("invalid HPKE tag".into()))?;

        let sk_r = <X25519HkdfSha256 as Kem>::PrivateKey::from_bytes(my_wrapping_priv)
            .map_err(|_| SecretSharingError::EnvelopeInvalid("invalid HPKE private key".into()))?;
        let enc = <X25519HkdfSha256 as Kem>::EncappedKey::from_bytes(enc)
            .map_err(|_| SecretSharingError::EnvelopeInvalid("invalid HPKE encapped key".into()))?;

        // Context setup (receiver)
        let mut rctx = setup_receiver::<ChaCha20Poly1305, HkdfSha256, X25519HkdfSha256>(
            &OpModeR::Base,
            &sk_r,
            &enc,
            SECRET_SHARING_HPKE_INFO,
        )
        .map_err(|e| SecretSharingError::UnwrapFailed(e.to_string()))?;

        // Decrypt in-place
        let mut body = [0u8; 32];
        body.copy_from_slice(ct_body);

        rctx.open_in_place_detached(&mut body, aad, &tag)
            .map_err(|e| {
                SecretSharingError::UnwrapFailed(format!(
                    "HPKE decryption failed: invalid ciphertext or key mismatch: {}",
                    e
                ))
            })?;

        Ok(Zeroizing::new(body))
    }

    /// HPKE encrypt for user payloads to be submitted to SecretSharingKey.
    ///
    /// This is designed for client-side encryption of user payloads that will
    /// be submitted to the chain and later decrypted by TEEs during rex execution.
    /// Uses the same HPKE parameters as the corresponding decrypt function.
    ///
    /// # Arguments
    /// * `sk_pub` - The SecretSharingKey public key bytes (from rex-registry)
    /// * `payload` - The plaintext payload to encrypt
    /// * `aad` - Associated authenticated data (should match decrypt)
    ///
    /// # Returns
    /// * `Ok(Vec<u8>)` - Encrypted payload in format: [enc (32 bytes)] || [ciphertext + tag]
    /// * `Err(SecretSharingError)` - If encryption fails
    ///
    /// # Security Notes
    /// - Uses X25519 + ChaCha20-Poly1305 + HKDF-SHA256
    /// - Compatible with `hpke_decrypt_user_payload()`
    /// - Safe for use in client applications (uses OS RNG)
    pub fn hpke_encrypt_user_payload(
        sk_pub: &PublicKey,
        payload: &[u8],
        aad: &[u8],
    ) -> Result<Vec<u8>, SecretSharingError> {
        let pk_r =
            <X25519HkdfSha256 as Kem>::PublicKey::from_bytes(sk_pub.as_ref()).map_err(|_| {
                SecretSharingError::EnvelopeInvalid("invalid public key for HPKE".into())
            })?;

        // Use our getrandom-based RNG
        let mut rng = GetrandomRng;

        // Context setup (sender) - same parameters as decrypt
        let (enc, mut sctx) = setup_sender::<ChaCha20Poly1305, HkdfSha256, X25519HkdfSha256, _>(
            &OpModeS::Base,
            &pk_r,
            SECRET_SHARING_HPKE_INFO,
            &mut rng,
        )
        .map_err(|e| SecretSharingError::WrapFailed(e.to_string()))?;

        // Encrypt payload; produce detached tag
        let mut ct_body = payload.to_vec();
        let tag = sctx
            .seal_in_place_detached(&mut ct_body, aad)
            .map_err(|e| SecretSharingError::WrapFailed(e.to_string()))?;

        // Format: [enc (32 bytes)] || [ciphertext + tag]
        let mut result = Vec::with_capacity(32 + ct_body.len() + 16);
        result.extend_from_slice(enc.to_bytes().as_ref());
        result.extend_from_slice(&ct_body);
        #[allow(deprecated)]
        result.extend_from_slice(tag.to_bytes().as_slice());

        Ok(result)
    }

    /// HPKE decrypt for user-submitted payloads encrypted to SecretSharingKey.
    ///
    /// This is designed for user payloads that are
    /// encrypted directly to the SecretSharingKey (SK_pub).
    ///
    /// # Arguments
    /// * `sk_priv` - The SecretSharingKey private key bytes
    /// * `enc` - HPKE encapsulated key (32 bytes)
    /// * `ciphertext` - Encrypted payload with tag appended
    /// * `aad` - Associated authenticated data
    ///
    /// # Returns
    /// * `Ok(Vec<u8>)` - Decrypted plaintext payload
    /// * `Err(SecretSharingError)` - If decryption fails
    pub fn hpke_decrypt_user_payload(
        sk_priv: &[u8; 32],
        enc: &[u8; 32],
        ciphertext: &[u8],
        aad: &[u8],
    ) -> Result<Zeroizing<Vec<u8>>, SecretSharingError> {
        // Ciphertext should be at least 16 bytes (minimum for tag)
        if ciphertext.len() < 16 {
            return Err(SecretSharingError::EnvelopeInvalid(
                "Ciphertext too short for HPKE tag".into(),
            ));
        }

        // Split ciphertext into body + tag (16 bytes)
        let (ct_body, tag_bytes) = ciphertext.split_at(ciphertext.len() - 16);
        let tag = AeadTag::<ChaCha20Poly1305>::from_bytes(tag_bytes)
            .map_err(|_| SecretSharingError::EnvelopeInvalid("invalid HPKE tag".into()))?;

        let sk_r = <X25519HkdfSha256 as Kem>::PrivateKey::from_bytes(sk_priv)
            .map_err(|_| SecretSharingError::EnvelopeInvalid("invalid HPKE private key".into()))?;
        let enc = <X25519HkdfSha256 as Kem>::EncappedKey::from_bytes(enc)
            .map_err(|_| SecretSharingError::EnvelopeInvalid("invalid HPKE encapped key".into()))?;

        // Context setup (receiver)
        let mut rctx = setup_receiver::<ChaCha20Poly1305, HkdfSha256, X25519HkdfSha256>(
            &OpModeR::Base,
            &sk_r,
            &enc,
            SECRET_SHARING_HPKE_INFO,
        )
        .map_err(|e| SecretSharingError::UnwrapFailed(e.to_string()))?;

        // Decrypt in-place
        let mut body = ct_body.to_vec();
        rctx.open_in_place_detached(&mut body, aad, &tag)
            .map_err(|e| {
                SecretSharingError::UnwrapFailed(format!(
                    "HPKE user payload decryption failed: invalid ciphertext or key mismatch: {}",
                    e
                ))
            })?;

        Ok(Zeroizing::new(body))
    }
}

impl SecretSharingEngine for HpkeX25519ChaChaEngine {
    fn generate_secret_key(&self) -> Result<SecretSharingKey, SecretSharingError> {
        Ok(SecretSharingKey::generate())
    }

    fn secret_key_public_key(&self, sk: &SecretSharingKey) -> PublicKey {
        sk.public_key_bytes()
    }

    fn wrapping_public_key(&self) -> PublicKey {
        self.my_wrapping_pub_bytes()
    }

    fn wrap_for_recipient(
        &self,
        sk: &SecretSharingKey,
        recipient_wrapping_pub_x25519: PublicKey,
        aad: Aad,
    ) -> Result<WrappedKeyEnvelope, SecretSharingError> {
        // AAD = request_id (u64) || key_id (u64)
        let mut aad_buf = [0u8; 8 + 8];
        aad_buf[0..8].copy_from_slice(&aad.request_id.to_le_bytes());
        aad_buf[8..16].copy_from_slice(&aad.key_id.value().to_le_bytes());

        let sk_priv = sk.private_key().to_bytes();
        let (enc, ct) = Self::hpke_wrap(&recipient_wrapping_pub_x25519, &sk_priv, &aad_buf)?;

        Ok(WrappedKeyEnvelope {
            sender_wrapping_pub_x25519: self.my_wrapping_pub_bytes(),
            recipient_wrapping_pub_x25519,
            aad,
            ciphertext: ct,
            enc,
        })
    }

    fn unwrap_and_install(
        &self,
        env: &WrappedKeyEnvelope,
    ) -> Result<SecretSharingKey, SecretSharingError> {
        if env.recipient_wrapping_pub_x25519 != self.my_wrapping_pub_bytes() {
            return Err(SecretSharingError::RecipientMismatch);
        }

        let mut aad_buf = [0u8; 8 + 8];
        aad_buf[0..8].copy_from_slice(&env.aad.request_id.to_le_bytes());
        aad_buf[8..16].copy_from_slice(&env.aad.key_id.value().to_le_bytes());

        let my_priv = self.my_hpke_priv_bytes();
        let sk_priv = Self::hpke_unwrap(&my_priv, &env.enc, &env.ciphertext, &aad_buf)?;
        Ok(SecretSharingKey::from_private_key_bytes(sk_priv))
    }
}

/// HPKE RNG that enforces TEE HW-only policy (RDSEED→RDRAND).
/// We wrap the core's HwRng to satisfy hpke::rand_core traits.
#[derive(Default)]
struct HpkeHwRng(crate::hw_rng::HwRng);

impl hpke::rand_core::CryptoRng for HpkeHwRng {}

impl hpke::rand_core::RngCore for HpkeHwRng {
    fn next_u32(&mut self) -> u32 {
        self.0.next_u32()
    }

    fn next_u64(&mut self) -> u64 {
        self.0.next_u64()
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.0.fill_bytes(dest)
    }
}

/// Getrandom-based RNG that implements hpke::rand_core traits.
/// This avoids version conflicts by using getrandom directly instead of bridging rand versions.
struct GetrandomRng;

impl hpke::rand_core::CryptoRng for GetrandomRng {}

impl hpke::rand_core::RngCore for GetrandomRng {
    fn next_u32(&mut self) -> u32 {
        let mut bytes = [0u8; 4];
        getrandom::getrandom(&mut bytes).expect("getrandom should not fail");
        u32::from_le_bytes(bytes)
    }

    fn next_u64(&mut self) -> u64 {
        let mut bytes = [0u8; 8];
        getrandom::getrandom(&mut bytes).expect("getrandom should not fail");
        u64::from_le_bytes(bytes)
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        getrandom::getrandom(dest).expect("getrandom should not fail");
    }
}

// ----------------- DKG helpers -----------------

/// Build AAD for DKG share transport.
///
/// This binds `{epoch, dealer, recipient}` into the HPKE AEAD context so that
/// shares cannot be replayed across participants or epochs. The commit set hash
/// should also be bound by the higher-level protocol (e.g., included in the
/// message envelope alongside this AAD) to disambiguate runs.
pub fn aad_for_dkg(epoch: u64, dealer: u16, recipient: u16, commit_set_hash: [u8; 32]) -> Aad {
    // Reuse the KeyId field to carry epoch; request_id carries (dealer<<16)|recipient for tracing.
    let req = ((dealer as u64) << 32) | (recipient as u64);
    let _ = commit_set_hash; // included implicitly via domain separation at higher layers if needed
    Aad {
        request_id: req,
        key_id: KeyId::new(epoch).expect("valid epoch"),
    }
}

use curve25519_dalek::scalar::Scalar;
use rialo_types::X25519_KEY_SIZE;

use crate::constants::SECRET_SHARING_HPKE_INFO;

/// HPKE-wrap a Curve25519 scalar for a recipient using this engine.
///
/// The scalar is encoded in canonical 32-byte little-endian form before
/// encryption. The returned envelope carries the provided AAD and the sender
/// and recipient wrapping keys.
pub fn wrap_scalar_for_recipient(
    engine: &HpkeX25519ChaChaEngine,
    s: Scalar,
    recipient_wrapping_pub_x25519: PublicKey,
    aad: Aad,
) -> Result<WrappedKeyEnvelope, SecretSharingError> {
    let sk_priv = s.to_bytes();
    let (enc, ct) = HpkeX25519ChaChaEngine::hpke_wrap(
        &recipient_wrapping_pub_x25519,
        &sk_priv,
        &bincode::serialize(&aad).map_err(|e| SecretSharingError::WrapFailed(e.to_string()))?,
    )?;
    Ok(WrappedKeyEnvelope {
        sender_wrapping_pub_x25519: engine.my_wrapping_pub_bytes(),
        recipient_wrapping_pub_x25519,
        aad,
        ciphertext: ct,
        enc,
    })
}

// Note: test-only RNG injection helper removed to keep this module minimal.

/// HPKE-unwrap a Curve25519 scalar from an envelope using this engine.
///
/// Verifies the provided `aad` matches the envelope's AAD, then decrypts and
/// parses a canonical Scalar. Returns an error if decoding is non-canonical.
pub fn unwrap_scalar(
    engine: &HpkeX25519ChaChaEngine,
    env: &WrappedKeyEnvelope,
    aad: &Aad,
) -> Result<Scalar, SecretSharingError> {
    if &env.aad != aad {
        return Err(SecretSharingError::EnvelopeInvalid("AAD mismatch".into()));
    }
    let my_priv = engine.my_hpke_priv_bytes();
    let body = HpkeX25519ChaChaEngine::hpke_unwrap(
        &my_priv,
        &env.enc,
        &env.ciphertext,
        &bincode::serialize(aad).map_err(|e| SecretSharingError::UnwrapFailed(e.to_string()))?,
    )?;
    let ct = Scalar::from_canonical_bytes(*body);
    match Option::<Scalar>::from(ct) {
        Some(s) => Ok(s),
        None => Err(SecretSharingError::EnvelopeInvalid(
            "non-canonical scalar bytes".into(),
        )),
    }
}

#[cfg(test)]
mod user_payload_tests {
    use super::*;
    use crate::USER_SECRET_AAD;

    fn create_test_secret_key() -> SecretSharingKey {
        SecretSharingKey::from_test_seed(b"test-secret-key-payload")
    }

    #[test]
    fn test_hpke_encrypt_user_payload_different_sizes() {
        let sk = create_test_secret_key();
        let sk_pub = sk.public_key_bytes();
        let sk_priv = sk.private_key().to_bytes();
        let aad = USER_SECRET_AAD;

        // Test various payload sizes
        let medium_payload = vec![0u8; 100];
        let large_payload = vec![42u8; 1000];
        let test_cases = [
            b"".as_slice(),              // Empty payload
            b"a".as_slice(),             // Single byte
            b"Hello, world!".as_slice(), // Short message
            medium_payload.as_slice(),   // Medium payload
            large_payload.as_slice(),    // Large payload
        ];

        for (i, payload) in test_cases.iter().enumerate() {
            let ciphertext =
                HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(&sk_pub, payload, aad)
                    .unwrap_or_else(|_| panic!("Encryption should succeed for test case {}", i));

            // Verify format
            assert_eq!(
                ciphertext.len(),
                32 + payload.len() + 16,
                "Ciphertext length should be enc + payload + tag for test case {}",
                i
            );

            // Decrypt and verify
            let (enc_bytes, ct_with_tag) = ciphertext.split_at(32);
            let mut enc_array = [0u8; 32];
            enc_array.copy_from_slice(enc_bytes);

            let decrypted = HpkeX25519ChaChaEngine::hpke_decrypt_user_payload(
                &sk_priv,
                &enc_array,
                ct_with_tag,
                aad,
            )
            .unwrap_or_else(|_| panic!("Decryption should succeed for test case {}", i));

            assert_eq!(
                *payload, &**decrypted,
                "Roundtrip should preserve payload for test case {}",
                i
            );
        }
    }

    #[test]
    fn test_hpke_encrypt_user_payload_weak_public_key() {
        // Note: All zeros is technically a valid X25519 public key, just a weak/identity element
        let weak_x25519_pubkey = x25519_dalek::PublicKey::from([0u8; 32]);
        let weak_pubkey = PublicKey::from(weak_x25519_pubkey);
        let payload = b"test payload";
        let aad = USER_SECRET_AAD;

        // This should succeed even with a weak key (HPKE allows it)
        let result = HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(&weak_pubkey, payload, aad);

        // The encryption should work, but decryption with the corresponding private key would fail
        // because the all-zeros private key doesn't correspond to the all-zeros public key in X25519
        match result {
            Ok(ciphertext) => {
                // Verify it has the expected format even with weak key
                assert!(
                    ciphertext.len() >= 48,
                    "Should produce valid ciphertext format"
                );
            }
            Err(_) => {
                // Some HPKE implementations might reject weak keys, which is also acceptable
                // This depends on the specific implementation behavior
            }
        }
    }

    #[test]
    fn test_hpke_encrypt_user_payload_deterministic_with_different_keys() {
        let sk1 = SecretSharingKey::from_test_seed(b"seed1");
        let sk2 = SecretSharingKey::from_test_seed(b"seed2");
        let payload = b"same payload";
        let aad = USER_SECRET_AAD;

        let ciphertext1 = HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(
            &sk1.public_key_bytes(),
            payload,
            aad,
        )
        .expect("Encryption should succeed");

        let ciphertext2 = HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(
            &sk2.public_key_bytes(),
            payload,
            aad,
        )
        .expect("Encryption should succeed");

        // Different keys should produce different ciphertexts
        assert_ne!(
            ciphertext1, ciphertext2,
            "Different keys should produce different ciphertexts"
        );

        // Each should decrypt correctly with its own key
        let (enc1, ct1) = ciphertext1.split_at(32);
        let mut enc1_array = [0u8; 32];
        enc1_array.copy_from_slice(enc1);

        let decrypted1 = HpkeX25519ChaChaEngine::hpke_decrypt_user_payload(
            &sk1.private_key().to_bytes(),
            &enc1_array,
            ct1,
            aad,
        )
        .expect("Decryption with correct key should succeed");

        assert_eq!(payload, &**decrypted1);
    }

    #[test]
    fn test_hpke_encrypt_user_payload_nondeterministic() {
        let sk = create_test_secret_key();
        let sk_pub = sk.public_key_bytes();
        let payload = b"nondeterministic test";
        let aad = USER_SECRET_AAD;

        // Same parameters should produce different ciphertexts due to randomness
        let ciphertext1 = HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(&sk_pub, payload, aad)
            .expect("Encryption should succeed");

        let ciphertext2 = HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(&sk_pub, payload, aad)
            .expect("Encryption should succeed");

        assert_ne!(
            ciphertext1, ciphertext2,
            "Same parameters should produce different ciphertexts due to randomness"
        );

        // Both should decrypt to the same payload
        let sk_priv = sk.private_key().to_bytes();

        for (i, ciphertext) in [&ciphertext1, &ciphertext2].iter().enumerate() {
            let (enc, ct) = ciphertext.split_at(32);
            let mut enc_array = [0u8; 32];
            enc_array.copy_from_slice(enc);

            let decrypted =
                HpkeX25519ChaChaEngine::hpke_decrypt_user_payload(&sk_priv, &enc_array, ct, aad)
                    .unwrap_or_else(|_| {
                        panic!("Decryption should succeed for ciphertext {}", i + 1)
                    });

            assert_eq!(
                payload, &**decrypted,
                "Both ciphertexts should decrypt to same payload"
            );
        }
    }
}
