//! The test vectors `draft-irtf-cfrg-vdaf-22` Appendix C points to, as published at the
//! commit its [TestVectors] reference names, run against this module: the files under
//! `live-conformance/aggregate/vdaf-22/`.
//!
//! Each file's `operations` run in order (C.1): `shard`, `verify_init`,
//! `verifier_shares_to_message`, `verify_next`, `aggregate` and `unshard`, each with the
//! file's inputs, each required to give the file's bytes or, where `success` is false, to fail.

use std::path::Path;

use serde_json::Value;

use super::field::{encode_vec, F128};
use super::prio3::{State, Vdaf};
use super::xof::Xof;

fn unhex(s: &str) -> Vec<u8> {
    (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).expect("hex")).collect()
}

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

/// The VDAF a file names, from its name and parameters; `None` for a variant this module
/// does not implement.
pub fn vdaf_of(name: &str, doc: &Value) -> Option<Vdaf> {
    let n = |k: &str| doc[k].as_u64().map(|x| x as usize);
    let shares = n("shares")?;
    if name.starts_with("Prio3Count_") {
        Some(Vdaf::count(shares))
    } else if name.starts_with("Prio3SumVec_") {
        Some(Vdaf::sumvec(shares, n("length")?, doc["max_measurement"].as_u64()?, n("chunk_length")?))
    } else if name.starts_with("Prio3Sum_") {
        Some(Vdaf::sum(shares, doc["max_measurement"].as_u64()?))
    } else if name.starts_with("Prio3Histogram_") {
        Some(Vdaf::histogram(shares, n("length")?, n("chunk_length")?))
    } else {
        None
    }
}

fn measurement(name: &str, m: &Value) -> Vec<u128> {
    if name.starts_with("Prio3SumVec_") {
        m.as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u128).collect()
    } else {
        vec![m.as_u64().unwrap() as u128]
    }
}

fn result_ints(v: &Value) -> Vec<u128> {
    match v {
        Value::Array(a) => a.iter().map(|x| x.as_u64().unwrap() as u128).collect(),
        x => vec![x.as_u64().unwrap() as u128],
    }
}

/// Run one Prio3 vector file; the failures, empty when every operation gives what it states.
pub fn run_prio3(name: &str, doc: &Value) -> Vec<String> {
    let mut fails = Vec::new();
    let Some(vdaf) = vdaf_of(name, doc) else {
        return vec![format!("{name}: not a variant this module implements")];
    };
    let ctx = unhex(doc["ctx"].as_str().unwrap());
    let vk = unhex(doc["verify_key"].as_str().unwrap());
    let reports = doc["reports"].as_array().unwrap();
    let mut states: Vec<Vec<Option<State>>> = vec![vec![None; vdaf.shares()]; reports.len()];
    let mut check = |what: String, ok: bool| {
        if !ok {
            fails.push(format!("{name}: {what}"));
        }
    };
    for op in doc["operations"].as_array().unwrap() {
        let kind = op["operation"].as_str().unwrap();
        let success = op["success"].as_bool().unwrap();
        let ri = op["report_index"].as_u64().map(|x| x as usize);
        let agg = op["aggregator_id"].as_u64().map(|x| x as usize);
        match kind {
            "shard" => {
                let r = &reports[ri.unwrap()];
                let enc = vdaf.encode(&measurement(name, &r["measurement"]));
                let out = enc.ok_or(()).and_then(|e| {
                    vdaf.shard(&ctx, &e, &unhex(r["nonce"].as_str().unwrap()), &unhex(r["rand"].as_str().unwrap())).map_err(|_| ())
                });
                match out {
                    Ok((ps, is)) => {
                        check(format!("shard {ri:?} public share"), success && hex(&ps) == r["public_share"].as_str().unwrap());
                        for (j, s) in is.iter().enumerate() {
                            check(format!("shard {ri:?} input share {j}"), hex(s) == r["input_shares"][j].as_str().unwrap());
                        }
                    }
                    Err(()) => check(format!("shard {ri:?} failed"), !success),
                }
            }
            "verify_init" => {
                let (i, j) = (ri.unwrap(), agg.unwrap());
                let r = &reports[i];
                let out = vdaf.verify_init(
                    &vk,
                    &ctx,
                    j,
                    &unhex(r["nonce"].as_str().unwrap()),
                    &unhex(r["public_share"].as_str().unwrap()),
                    &unhex(r["input_shares"][j].as_str().unwrap()),
                );
                match out {
                    Ok((st, vs)) => {
                        let want = r["verifier_shares"][0][j].as_str().unwrap_or("");
                        check(format!("verify_init {i}/{j} verifier share"), success && hex(&vs) == want);
                        states[i][j] = Some(st);
                    }
                    Err(_) => check(format!("verify_init {i}/{j} failed"), !success),
                }
            }
            "verifier_shares_to_message" => {
                let i = ri.unwrap();
                let r = &reports[i];
                let shares: Vec<Vec<u8>> = r["verifier_shares"][0].as_array().unwrap().iter().map(|s| unhex(s.as_str().unwrap())).collect();
                match vdaf.verifier_shares_to_message(&ctx, &shares) {
                    Ok(m) => check(format!("message {i}"), success && hex(&m) == r["verifier_messages"][0].as_str().unwrap()),
                    Err(_) => check(format!("message {i} refused"), !success),
                }
            }
            "verify_next" => {
                let (i, j) = (ri.unwrap(), agg.unwrap());
                let r = &reports[i];
                let msg = unhex(r["verifier_messages"][0].as_str().unwrap());
                let out = states[i][j].as_ref().ok_or(()).and_then(|st| vdaf.verify_next(st, &msg).map_err(|_| ()));
                match out {
                    Ok(o) => check(format!("verify_next {i}/{j}"), success && hex(&o) == r["out_shares"][j].as_str().unwrap()),
                    Err(()) => check(format!("verify_next {i}/{j} refused"), !success),
                }
            }
            "aggregate" => {
                let j = agg.unwrap();
                let mut a = vdaf.agg_init();
                for r in reports {
                    a = vdaf.merge(&a, &unhex(r["out_shares"][j].as_str().unwrap())).expect("output share");
                }
                check(format!("aggregate {j}"), success && hex(&a) == doc["agg_shares"][j].as_str().unwrap());
            }
            "unshard" => {
                let shares: Vec<Vec<u8>> = doc["agg_shares"].as_array().unwrap().iter().map(|s| unhex(s.as_str().unwrap())).collect();
                let got = vdaf.unshard(&shares).ok();
                check("unshard".into(), success && got == Some(result_ints(&doc["agg_result"])));
            }
            other => check(format!("unknown operation {other}"), false),
        }
    }
    fails
}

/// The XofTurboShake128 vector: a derived seed and an expansion into Field128.
pub fn run_xof(doc: &Value) -> Vec<String> {
    let (seed, dst, binder) = (unhex(doc["seed"].as_str().unwrap()), unhex(doc["dst"].as_str().unwrap()), unhex(doc["binder"].as_str().unwrap()));
    let mut fails = Vec::new();
    if hex(&Xof::derive_seed(&seed, &dst, &binder)) != doc["derived_seed"].as_str().unwrap() {
        fails.push("XofTurboShake128: derived seed".into());
    }
    let v: Vec<F128> = Xof::expand_into_vec(&seed, &dst, &binder, doc["length"].as_u64().unwrap() as usize);
    if hex(&encode_vec(&v)) != doc["expanded_vec_field128"].as_str().unwrap() {
        fails.push("XofTurboShake128: expanded vector".into());
    }
    fails
}

/// Run every file of the directory; the files run and the failures.
pub fn run_dir(dir: &Path) -> (Vec<String>, Vec<String>) {
    let mut names: Vec<String> = std::fs::read_dir(dir).expect("vdaf-22 directory").flatten().map(|e| e.file_name().to_string_lossy().into_owned()).collect();
    names.sort();
    let mut fails = Vec::new();
    for n in &names {
        let doc: Value = serde_json::from_str(&std::fs::read_to_string(dir.join(n)).unwrap()).unwrap();
        if n == "XofTurboShake128.json" {
            fails.extend(run_xof(&doc));
        } else {
            fails.extend(run_prio3(n.trim_end_matches(".json"), &doc));
        }
    }
    (names, fails)
}
