//! Known-answer tests validating the primitive crates and our calling
//! conventions against published RFC test vectors.

use chacha20poly1305::aead::{Aead, Payload};
use chacha20poly1305::{ChaCha20Poly1305, KeyInit, Nonce};
use hex_literal::hex;
use hkdf::Hkdf;
use libcrux_ml_kem::mlkem768::{self, MlKem768Ciphertext, MlKem768PrivateKey};
use rmux_web_crypto::ml_kem::{
    encapsulate, KeyPair, CIPHERTEXT_LEN, ENCAPSULATION_KEY_LEN, ENCAPS_RANDOMNESS_LEN,
    KEYGEN_RANDOMNESS_LEN, SHARED_SECRET_LEN,
};
use sha2::{Digest, Sha256};

const DECAPSULATION_KEY_LEN: usize = 2400;

/// RFC 5869, Appendix A, Test Case 1: HKDF-SHA256.
///
/// Validates the `hkdf` + `sha2` usage by checking the exact PRK and the
/// 42-byte OKM.
#[test]
fn rfc5869_hkdf_sha256_test_case_1() {
    let ikm = hex!("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
    let salt = hex!("000102030405060708090a0b0c");
    let info = hex!("f0f1f2f3f4f5f6f7f8f9");

    let expected_prk = hex!("077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5");
    let expected_okm = hex!(
        "3cb25f25faacd57a90434f64d0362f2a"
        "2d2d0a90cf1a5a4c5db02d56ecc4c5bf"
        "34007208d5b887185865"
    );

    let (prk, hk) = Hkdf::<Sha256>::extract(Some(&salt), &ikm);
    assert_eq!(prk.as_slice(), &expected_prk, "PRK mismatch");

    let mut okm = [0u8; 42];
    hk.expand(&info, &mut okm).expect("expand 42 bytes");
    assert_eq!(okm, expected_okm, "OKM mismatch");
}

/// RFC 8439, Section 2.8.2: ChaCha20-Poly1305 AEAD.
///
/// Validates the AEAD primitive and our `(Nonce, Payload { msg, aad })`
/// calling convention by checking the exact ciphertext+tag.
#[test]
fn rfc8439_chacha20poly1305_aead() {
    let plaintext = hex!(
        "4c616469657320616e642047656e746c"
        "656d656e206f662074686520636c6173"
        "73206f66202739393a20496620492063"
        "6f756c64206f6666657220796f75206f"
        "6e6c79206f6e652074697020666f7220"
        "746865206675747572652c2073756e73"
        "637265656e20776f756c642062652069"
        "742e"
    );
    let aad = hex!("50515253c0c1c2c3c4c5c6c7");
    let key = hex!(
        "808182838485868788898a8b8c8d8e8f"
        "909192939495969798999a9b9c9d9e9f"
    );
    let nonce = hex!("070000004041424344454647");

    let expected = hex!(
        "d31a8d34648e60db7b86afbc53ef7ec2"
        "a4aded51296e08fea9e2b5a736ee62d6"
        "3dbea45e8ca9671282fafb69da92728b"
        "1a71de0a9e060b2905d6a5b67ecd3b36"
        "92ddbd7f2d778b8c9803aee328091b58"
        "fab324e4fad675945585808b4831d7bc"
        "3ff4def08e4b7a9de576d26586cec64b"
        "6116"
        // 16-byte Poly1305 tag
        "1ae10b594f09e26a7e902ecbd0600691"
    );

    let cipher = ChaCha20Poly1305::new((&key).into());
    let ct = cipher
        .encrypt(
            Nonce::from_slice(&nonce),
            Payload {
                msg: &plaintext,
                aad: &aad,
            },
        )
        .expect("encrypt");
    assert_eq!(ct.as_slice(), &expected, "ciphertext+tag mismatch");
}

/// NIST ACVP ML-KEM-768 FIPS203 keygen vector.
///
/// The ACVP expected result carries the full 1184-byte encapsulation key. The
/// test pins its SHA-256 digest to keep this fixture small while still checking
/// the exact generated bytes against an external source.
#[test]
fn nist_acvp_ml_kem_768_keygen() {
    let d = hex!("e582b7d75e6c80b05ae392a1fc9f7153b12390fd99930368cc67a768baebc8a0");
    let z = hex!("1cdacb8740c0b87c4a379575f187b367cbfa3b300bf591b109f79816e9cbe8f0");
    let expected_ek_sha256 =
        hex!("4158f6afb5e516c99f1da07da8c651348422b17c1f4e9a08ad73fb1f91249b3e");

    let mut keygen_randomness = [0u8; KEYGEN_RANDOMNESS_LEN];
    keygen_randomness[..32].copy_from_slice(&d);
    keygen_randomness[32..].copy_from_slice(&z);

    let ek = KeyPair::generate(keygen_randomness).encapsulation_key();
    assert_eq!(Sha256::digest(ek).as_slice(), &expected_ek_sha256);
}

/// C2SP CCTV ML-KEM-768 unlucky SampleNTT vector.
///
/// Validates our libcrux ML-KEM-768 call surface against an external vector
/// rather than a same-implementation round trip. The vector is CC0 and targets
/// ML-KEM as specified in FIPS 203.
#[test]
fn cctv_ml_kem_768_encapsulation() {
    let vector = MlKemVector::load();

    let (ciphertext, shared_secret) =
        encapsulate(&vector.ek, vector.m).expect("CCTV encapsulation key is valid");
    assert_eq!(ciphertext, vector.c, "ML-KEM-768 ciphertext mismatch");
    assert_eq!(shared_secret, vector.k, "ML-KEM-768 shared secret mismatch");
}

/// C2SP CCTV ML-KEM-768 decapsulation vector.
///
/// Exercises decapsulation against a published private-key vector. The public
/// rmux wrapper intentionally does not expose private-key import, so this test
/// calls the same libcrux decapsulation surface with the external `dk`.
#[test]
fn cctv_ml_kem_768_decapsulation() {
    let vector = MlKemVector::load();
    let private_key = MlKem768PrivateKey::from(&vector.dk);
    let ciphertext = MlKem768Ciphertext::from(&vector.c);

    assert_eq!(
        mlkem768::decapsulate(&private_key, &ciphertext),
        vector.k,
        "ML-KEM-768 decapsulated shared secret mismatch"
    );
}

struct MlKemVector {
    ek: [u8; ENCAPSULATION_KEY_LEN],
    dk: [u8; DECAPSULATION_KEY_LEN],
    m: [u8; ENCAPS_RANDOMNESS_LEN],
    k: [u8; SHARED_SECRET_LEN],
    c: [u8; CIPHERTEXT_LEN],
}

impl MlKemVector {
    fn load() -> Self {
        let fixture = include_str!("vectors/ml_kem_768_cctv_unlucky.txt");
        Self {
            ek: fixture_hex(fixture, "ek"),
            dk: fixture_hex(fixture, "dk"),
            m: fixture_hex(fixture, "m"),
            k: fixture_hex(fixture, "K"),
            c: fixture_hex(fixture, "c"),
        }
    }
}

fn fixture_hex<const N: usize>(fixture: &str, name: &str) -> [u8; N] {
    let value = fixture
        .lines()
        .find_map(|line| line.strip_prefix(&format!("{name} = ")))
        .unwrap_or_else(|| panic!("missing {name} in ML-KEM vector"));
    decode_hex_array(value)
}

fn decode_hex_array<const N: usize>(hex: &str) -> [u8; N] {
    assert_eq!(hex.len(), N * 2, "unexpected hex length");
    let mut out = [0u8; N];
    for (index, chunk) in hex.as_bytes().chunks_exact(2).enumerate() {
        out[index] = (hex_nibble(chunk[0]) << 4) | hex_nibble(chunk[1]);
    }
    out
}

fn hex_nibble(byte: u8) -> u8 {
    match byte {
        b'0'..=b'9' => byte - b'0',
        b'a'..=b'f' => byte - b'a' + 10,
        b'A'..=b'F' => byte - b'A' + 10,
        _ => panic!("invalid hex digit"),
    }
}
