//! Conformance corpus for hybrid signatures (`pq-conformance/`, extensions/pq-signatures.md §10).
//!
//!   wai_pq_vectors gen      <corpus-dir>   write vectors.json from fixed seeds
//!   wai_pq_vectors verify   <corpus-dir>   vectors.json regenerates byte for byte, no file but
//!                                          README.md, acvp/ and the draft's vector is there,
//!                                          and every case gets the verdict it states
//!   wai_pq_vectors verdicts <vectors.json> print `name<TAB>verdict` for every case, ignoring
//!                                          the stated verdict (the differential's input)
//!   wai_pq_vectors bench    [rounds]       time signing and verifying per suite, and the
//!                                          sizes §9 states (run it built with --release)
//!
//! The generator and this verifier are one implementation. The corpus's second
//! implementation is `pq-conformance/pq_verify.py`, written from the extension's text and
//! sharing no code with this binary.

use std::path::Path;
use std::time::Instant;
use wai::pq_sig::{self, corpus, HybridKey, Policy};

fn main() {
    let a: Vec<String> = std::env::args().collect();
    match (a.get(1).map(String::as_str), a.get(2)) {
        (Some("gen"), Some(d)) => {
            if let Err(e) = std::fs::create_dir_all(d) {
                eprintln!("create {d}: {e}");
                std::process::exit(1);
            }
            for (name, body) in corpus::generate() {
                let p = Path::new(d).join(&name);
                if let Err(e) = std::fs::write(&p, body) {
                    eprintln!("write {}: {e}", p.display());
                    std::process::exit(1);
                }
            }
            println!("wrote {d}");
        }
        (Some("verify"), Some(d)) => std::process::exit(verify(Path::new(d))),
        (Some("verdicts"), Some(f)) => {
            let doc = match std::fs::read_to_string(f).map_err(|e| e.to_string()).and_then(|t| serde_json::from_str(&t).map_err(|e| e.to_string())) {
                Ok(d) => d,
                Err(e) => {
                    eprintln!("read {f}: {e}");
                    std::process::exit(1);
                }
            };
            for (name, v) in corpus::verdicts(&doc) {
                println!("{name}\t{v}");
            }
        }
        (Some("bench"), n) => bench(n.and_then(|s| s.parse().ok()).unwrap_or(7)),
        _ => {
            eprintln!("usage: wai_pq_vectors gen|verify <corpus-dir> | verdicts <vectors.json> | bench [rounds]");
            std::process::exit(2);
        }
    }
}

/// The files the directory holds besides the generated one: its README, the vendored ACVP
/// vectors and their derivation, the draft's vector, and the second implementation, with the
/// cache Python writes beside a module it imports (ignored by git, never committed).
const KEPT: &[&str] = &["README.md", "acvp", "draft-composite-sigs-19.json", "pq_verify.py", "_mldsa.py", "__pycache__"];

fn verify(dir: &Path) -> i32 {
    let mut ok = true;
    let generated = corpus::generate();
    for (name, body) in &generated {
        let same = std::fs::read_to_string(dir.join(name)).is_ok_and(|t| &t == body);
        println!("{} {name} regenerates byte for byte from its seeds", if same { "PASS" } else { "FAIL" });
        ok &= same;
    }
    match std::fs::read_dir(dir) {
        Ok(entries) => {
            for e in entries.flatten() {
                let name = e.file_name().to_string_lossy().into_owned();
                if !KEPT.contains(&name.as_str()) && !generated.iter().any(|(g, _)| *g == name) {
                    println!("FAIL {name} is in the corpus directory and is neither generated nor kept");
                    ok = false;
                }
            }
        }
        Err(e) => {
            eprintln!("read {}: {e}", dir.display());
            return 1;
        }
    }
    let doc: serde_json::Value = match std::fs::read_to_string(dir.join("vectors.json")).ok().and_then(|t| serde_json::from_str(&t).ok()) {
        Some(d) => d,
        None => {
            println!("FAIL vectors.json does not read");
            return 1;
        }
    };
    let fails = corpus::run(&doc);
    for f in &fails {
        println!("FAIL {f}");
    }
    let n = doc["cases"].as_array().map_or(0, Vec::len);
    println!("{} every one of the {n} cases gets the verdict it states", if fails.is_empty() { "PASS" } else { "FAIL" });
    (!ok || !fails.is_empty()) as i32
}

/// The per-operation time of `f`, in microseconds: the median of `rounds` rounds of `n`
/// calls, with the least and the most beside it.
fn time(rounds: usize, n: usize, mut f: impl FnMut()) -> (f64, f64, f64) {
    let mut per: Vec<f64> = (0..rounds)
        .map(|_| {
            let t = Instant::now();
            for _ in 0..n {
                f();
            }
            t.elapsed().as_secs_f64() * 1e6 / n as f64
        })
        .collect();
    per.sort_by(f64::total_cmp);
    (per[0], per[per.len() / 2], per[per.len() - 1])
}

fn bench(rounds: usize) {
    use ed25519_dalek::Signer;
    let k = HybridKey::from_seeds([0x71; 32], [0x72; 32]);
    let doc = corpus::document();
    // 64 distinct messages of 512 bytes, about a claim's signing input: ML-DSA's signing
    // loop repeats a number of times that depends on the message (FIPS 204 §4, Table 1:
    // 5.1 on average for ML-DSA-65), so one message would time one draw of that number.
    let msgs: Vec<Vec<u8>> = (0..64u8).map(|i| [vec![i; 8], vec![0x5a; 504]].concat()).collect();
    let ed: Vec<[u8; 64]> = msgs.iter().map(|m| k.ed25519().sign(m).to_bytes()).collect();
    let hy: Vec<([u8; 64], wai::sig::PqHalf)> = msgs.iter().map(|m| k.sign(m)).collect();
    let mut i = 0usize;
    let mut next = || {
        i = (i + 1) % msgs.len();
        i
    };
    let rows: Vec<(&str, (f64, f64, f64))> = vec![
        ("Ed25519 sign", time(rounds, 2048, || {
            let j = next();
            std::hint::black_box(k.ed25519().sign(&msgs[j]));
        })),
        ("Ed25519 verify (SPEC §7 strict)", time(rounds, 2048, || {
            let j = next();
            std::hint::black_box(wai::sig::verify_strict(&k.public(), &msgs[j], &ed[j])).ok();
        })),
        ("ML-DSA-65 sign (deterministic)", time(rounds, 256, || {
            let j = next();
            std::hint::black_box(k.ml_dsa_sign(&msgs[j], pq_sig::LABEL));
        })),
        ("ML-DSA-65 verify", time(rounds, 512, || {
            let j = next();
            std::hint::black_box(pq_sig::ml_dsa_verify(k.pq_public(), &pq_sig::representative(&msgs[j]), pq_sig::LABEL, &hy[j].1.sig));
        })),
        ("hybrid sign (both halves)", time(rounds, 256, || {
            let j = next();
            std::hint::black_box(k.sign(&msgs[j]));
        })),
        ("hybrid verify (both halves)", time(rounds, 512, || {
            let j = next();
            std::hint::black_box(pq_sig::verify(&k.public(), &msgs[j], &hy[j].0, Some(&hy[j].1), Policy::RequireHybrid)).ok();
        })),
    ];
    println!("rounds={rounds} (median, least, most per operation, microseconds)");
    for (name, (lo, med, hi)) in &rows {
        println!("{name:<36} {med:>10.1} µs   [{lo:.1} .. {hi:.1}]");
    }
    let s = &doc["sizes"];
    println!("sizes (bytes): {}", serde_json::to_string_pretty(s).expect("JSON"));
    let g = |k: &str, f: &str| s[k][f].as_u64().unwrap_or(0);
    // The corpus's object body is of a three-object group: two siblings in its proof. An
    // object of an n-object group has ⌈log₂ n⌉ (jwp-receipts §2.7.2, §3 step 5).
    let obj_base = s["object_receipt_binary"].as_u64().unwrap_or(0) - 2 * 32;
    // A group's receipts record is the group receipt, a count, and one object body per object
    // (jwp-receipts §2.7.3): what a receipts track carries per media group
    // (moq-streaming-format §4.1).
    for n in [1u64, 30, 60] {
        let obj = obj_base + 32 * u64::from(n.next_power_of_two().trailing_zeros());
        let classical = 4 + g("group_receipt_binary", "classical") + 4 + n * obj;
        let hybrid = 4 + g("group_receipt_binary", "hybrid") + 4 + n * obj;
        println!(
            "receipts record, {n:>2} objects of {obj} bytes: classical {classical} B, hybrid {hybrid} B (+{} B, {:.2}x); at one group a second {:.1} kbit/s -> {:.1} kbit/s",
            hybrid - classical,
            hybrid as f64 / classical as f64,
            classical as f64 * 8.0 / 1000.0,
            hybrid as f64 * 8.0 / 1000.0
        );
    }
}
