//! The ML-DSA-65 implementation this module signs and verifies with, held to FIPS 204's
//! published test vectors before anything else is built on it, and the composite to the
//! test vector `draft-ietf-lamps-pq-composite-sigs-19` publishes for this suite.
//!
//! The vectors are vendored in `pq-conformance/acvp/` (`extract.py` derives them from the
//! ACVP internal projections at a pinned upstream commit; `pq-conformance/README.md` gives
//! the commit and the upstream SHA-256s) and in `pq-conformance/draft-composite-sigs-19.json`.
//! Each file is pinned here by its SHA-256, so a changed byte fails before a vector runs.
//!
//! - keyGen: every ML-DSA-65 seed gives the public key and the expanded private key stated
//!   (FIPS 204 Algorithm 6; Algorithms 22 and 24).
//! - sigGen, deterministic: the external interface with a context (Algorithm 2 with `rnd` =
//!   0³²), the internal interface (Algorithm 7), and the internal interface over an external
//!   μ, each reproduces the stated signature byte for byte.
//! - sigVer: the same three interfaces give every vector its stated verdict, the modified
//!   commitment, response, hint and message ones refused.

#![allow(deprecated)] // `ExpandedSigningKey::{from_expanded, to_expanded}`: the sigGen vectors give expanded keys.

use super::*;
use ml_dsa::{ExpandedSigningKeyBytes, SigningKey};
use sha2::Sha256;

const DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../pq-conformance");

/// The vendored files and their SHA-256s.
const PINNED: [(&str, &str); 4] = [
    ("acvp/ml-dsa-65-keygen.json", "06b637b695150bfc752748e8636858f2e2da9c77547ab9856ddd115af5754a0f"),
    ("acvp/ml-dsa-65-siggen.json", "8c47e6fa05e32abcaa9ae32a6a361771ca5bd8275847db9eac4b5a959eb9c327"),
    ("acvp/ml-dsa-65-sigver.json", "cd48b587b423ec42b62971afcbfc1395491b994101192b0a10ce681a42e78a81"),
    ("draft-composite-sigs-19.json", "d880a8fbb305e520620bff771b1e11bb21a83c97fbb9e286210d7c9e64898f20"),
];

fn load(name: &str) -> Value {
    let (_, digest) = PINNED.iter().find(|(n, _)| *n == name).expect("a pinned file");
    let bytes = std::fs::read(format!("{DIR}/{name}")).expect("the vendored file");
    assert_eq!(crate::claim_receipt::hex(&Sha256::digest(&bytes)), *digest, "{name}: not the pinned bytes");
    serde_json::from_slice(&bytes).expect("JSON")
}

fn h(v: &Value, k: &str) -> Vec<u8> {
    crate::claim_receipt::unhex(&v[k].as_str().unwrap_or_else(|| panic!("{k}")).to_ascii_lowercase()).expect("hex")
}

fn groups(doc: &Value) -> impl Iterator<Item = &Value> {
    doc["testGroups"].as_array().expect("groups").iter()
}

fn tests(g: &Value) -> impl Iterator<Item = &Value> {
    g["tests"].as_array().expect("tests").iter()
}

fn expanded(sk: &[u8]) -> ExpandedSigningKey<MlDsa65> {
    let enc = ExpandedSigningKeyBytes::<MlDsa65>::try_from(sk).expect("4032 bytes");
    ExpandedSigningKey::<MlDsa65>::from_expanded(&enc)
}

fn vk(pk: &[u8]) -> VerifyingKey<MlDsa65> {
    VerifyingKey::<MlDsa65>::decode(&EncodedVerifyingKey::<MlDsa65>::try_from(pk).expect("1952 bytes"))
}

fn sig(s: &[u8]) -> Option<Signature<MlDsa65>> {
    Signature::<MlDsa65>::decode(&EncodedSignature::<MlDsa65>::try_from(s).ok()?)
}

#[test]
fn every_pinned_file_has_its_bytes() {
    for (name, _) in PINNED {
        load(name);
    }
}

#[test]
fn acvp_keygen() {
    let doc = load("acvp/ml-dsa-65-keygen.json");
    let mut n = 0;
    for g in groups(&doc) {
        assert_eq!(g["parameterSet"], "ML-DSA-65");
        for t in tests(g) {
            let seed: [u8; 32] = h(t, "seed").try_into().expect("32 bytes");
            let sk = SigningKey::<MlDsa65>::from_seed(&B32::from(seed));
            assert_eq!(sk.expanded_key().verifying_key().encode().to_vec(), h(t, "pk"), "tcId {}", t["tcId"]);
            assert_eq!(sk.expanded_key().to_expanded().to_vec(), h(t, "sk"), "tcId {}", t["tcId"]);
            n += 1;
        }
    }
    assert_eq!(n, 25);
}

#[test]
fn acvp_siggen_deterministic() {
    let doc = load("acvp/ml-dsa-65-siggen.json");
    let mut seen = [0usize; 3];
    for g in groups(&doc) {
        assert_eq!(g["parameterSet"], "ML-DSA-65");
        assert_eq!(g["deterministic"], true);
        for t in tests(g) {
            let key = expanded(&h(t, "sk"));
            let got = match (g["signatureInterface"].as_str(), g["externalMu"].as_bool()) {
                (Some("external"), Some(false)) => {
                    assert_eq!(g["preHash"], "pure");
                    seen[0] += 1;
                    key.sign_deterministic(&h(t, "message"), &h(t, "context")).expect("a context under 256 bytes")
                }
                (Some("internal"), Some(false)) => {
                    seen[1] += 1;
                    key.sign_internal(&[&h(t, "message")], &B32::default())
                }
                (Some("internal"), Some(true)) => {
                    seen[2] += 1;
                    let mu = h(t, "mu");
                    key.sign_mu_deterministic(&mu.as_slice().try_into().expect("64 bytes"))
                }
                other => panic!("a group this test does not run: {other:?}"),
            };
            assert_eq!(got.encode().to_vec(), h(t, "signature"), "tgId {} tcId {}", g["tgId"], t["tcId"]);
            // The public key the vector states verifies what it signed.
            assert_eq!(key.verifying_key().encode().to_vec(), h(t, "pk"));
        }
    }
    assert_eq!(seen, [15, 15, 15]);
}

#[test]
fn acvp_sigver() {
    let doc = load("acvp/ml-dsa-65-sigver.json");
    let (mut seen, mut refused) = ([0usize; 3], 0);
    for g in groups(&doc) {
        assert_eq!(g["parameterSet"], "ML-DSA-65");
        for t in tests(g) {
            let k = vk(&h(t, "pk"));
            let s = h(t, "signature");
            let verdict = match (g["signatureInterface"].as_str(), g["externalMu"].as_bool()) {
                (Some("external"), Some(false)) => {
                    seen[0] += 1;
                    // The function every receipt family calls, on the external interface.
                    let via = ml_dsa_verify(&h(t, "pk"), &h(t, "message"), &h(t, "context"), &s);
                    let direct = sig(&s).is_some_and(|s| k.verify_with_context(&h(t, "message"), &h(t, "context"), &s));
                    assert_eq!(via, direct);
                    via
                }
                (Some("internal"), Some(false)) => {
                    seen[1] += 1;
                    sig(&s).is_some_and(|s| k.verify_internal(&h(t, "message"), &s))
                }
                (Some("internal"), Some(true)) => {
                    seen[2] += 1;
                    let mu = h(t, "mu");
                    sig(&s).is_some_and(|s| k.verify_mu(&mu.as_slice().try_into().expect("64 bytes"), &s))
                }
                other => panic!("a group this test does not run: {other:?}"),
            };
            assert_eq!(verdict, t["testPassed"].as_bool().expect("a verdict"), "tgId {} tcId {}: {}", g["tgId"], t["tcId"], t["reason"]);
            refused += usize::from(!verdict);
        }
    }
    assert_eq!(seen, [15, 15, 15]);
    assert_eq!(refused, 36, "four modified vectors of every five");
}

/// The suite against its own test vector: the composite public key `mldsaPK ‖ tradPK`, the
/// private key `mldsaSeed ‖ tradSK`, and the signature value `mldsaSig ‖ tradSig` (draft
/// §4) over the draft's message, with no context and with its context.
#[test]
fn the_composite_test_vector() {
    let v = load("draft-composite-sigs-19.json");
    assert_eq!(v["tcId"], SUITE);
    let (pk, sk, m, ctx) = (h(&v, "pk"), h(&v, "sk"), h(&v, "m"), h(&v, "ctx"));
    assert_eq!((pk.len(), sk.len()), (PQ_KEY_LEN + 32, 64));
    let (pq_pk, ed_pk): (&[u8], [u8; 32]) = (&pk[..PQ_KEY_LEN], pk[PQ_KEY_LEN..].try_into().expect("32"));
    let key = HybridKey::from_seeds(sk[32..].try_into().expect("32"), sk[..32].try_into().expect("32"));
    assert_eq!(key.pq_public(), pq_pk, "the ML-DSA seed gives the vector's key");
    assert_eq!(key.public(), ed_pk, "the Ed25519 seed gives the vector's key");
    // M' with an application context (draft §2.2): WAI signs with none, the vector with both.
    let rep = |ctx: &[u8]| {
        let mut r = PREFIX.to_vec();
        r.extend_from_slice(LABEL);
        r.push(u8::try_from(ctx.len()).expect("under 256"));
        r.extend_from_slice(ctx);
        r.extend_from_slice(&Sha512::digest(&m));
        r
    };
    assert_eq!(rep(&[]), representative(&m));
    for (field, c) in [("s", Vec::new()), ("sWithContext", ctx)] {
        let s = h(&v, field);
        assert_eq!(s.len(), PQ_SIG_LEN + 64);
        let (pq_sig, ed_sig): (&[u8], [u8; 64]) = (&s[..PQ_SIG_LEN], s[PQ_SIG_LEN..].try_into().expect("64"));
        let mp = rep(&c);
        // Draft §3.3: both components over M', the ML-DSA one with the label as context.
        assert!(ml_dsa_verify(pq_pk, &mp, LABEL, pq_sig), "{field}: the ML-DSA half");
        assert_eq!(sig::verify_strict(&ed_pk, &mp, &ed_sig), Ok(()), "{field}: the Ed25519 half, strictly");
        if c.is_empty() {
            // The function every family calls, with the vector's halves.
            let half = PqHalf { key: pq_pk.to_vec(), sig: pq_sig.to_vec() };
            assert_eq!(verify(&ed_pk, &m, &ed_sig, Some(&half), Policy::RequireHybrid), Ok(()));
            // Ed25519 is deterministic: the key signs the vector's classical half again.
            let (ours_ed, ours_pq) = key.sign(&m);
            assert_eq!(ours_ed, ed_sig);
            assert_eq!(verify(&ed_pk, &m, &ours_ed, Some(&ours_pq), Policy::RequireHybrid), Ok(()));
        }
        // Either half over the other context's M' fails.
        let other = rep(if c.is_empty() { b"x" } else { b"" });
        assert!(!ml_dsa_verify(pq_pk, &other, LABEL, pq_sig));
        assert!(sig::verify_strict(&ed_pk, &other, &ed_sig).is_err());
    }
}
