//! The conformance corpus of `extensions/prior-carriage.md` §14:
//! `live-conformance/prior-forms/`.
//!
//! [`generate`] writes it from native artifacts other corpora commit, each pinned here by its
//! SHA-256 ([`INPUT_NATIVES`]): the two trained `WSR1` priors of the exact-upscale corpus, its
//! seeded one, and four generated `WHM1` models of the integer-operators corpus. Every file is
//! a function of those bytes, so the corpus is the same bytes on every machine. Each vector
//! states its outcome; the generator asserts that this module reaches it before writing it,
//! and that the corpus gives every code of §13.
//!
//! [`verify`] checks a corpus directory against this module.
//! `live-conformance/prior_forms_verify.py` checks the same files from the extension's text,
//! sharing no code with this module.

use std::collections::BTreeMap;

use serde_json::{json, Value};

use super::{binding, encode, gguf, load, nnc, safetensors, DType, Form, Tensor, TensorSet, CODES};

/// The native artifacts the corpus reads: `(name, binding, path from the corpus directory,
/// SHA-256)`.
pub const INPUT_NATIVES: &[(&str, &str, &str, &str)] = &[
    ("generic", "wsr1/1", "../upscale/priors/generic.wsr", "0106ab198efcbc009a32c4fbf85da4a557daf827d1f92d8de9c7939953329915"),
    ("title", "wsr1/1", "../upscale/priors/title.wsr", "9ea90833e37dc07982682799abf95840b32fa4e152eebc1830776a5b0c843a07"),
    ("seeded", "wsr1/1", "../upscale/priors/seeded.wsr", "18b634abf31169c3c5126e258da73ca6cb0c4ba219220fca6f7736f869b94807"),
    ("hyper_g_s", "whm1/1", "../../ops-conformance/sets/hyper/g_s.whm", "fbd6b941d412e7bd0bc73fef90e4131780942aeb06b441abe887b4eb75768d9e"),
    ("hyper_h_s", "whm1/1", "../../ops-conformance/sets/hyper/h_s.whm", "0aad151f7f4a27f9c28fa3452692642802a9146ac102fcf664fdf182c80670a7"),
    ("world_t", "whm1/1", "../../ops-conformance/sets/world/T.whm", "a47ad9f8d12f8c0e8edde8e61cabe31794a8d7cdb6a4768c95db6a5816713c02"),
    ("world_g_s", "whm1/1", "../../ops-conformance/sets/world/g_s.whm", "1c13bee0387fd312580e44ad5f726ada6bd8f6cd06546dd92d5d66a72d9ab3e1"),
];

fn sha_hex(b: &[u8]) -> String {
    crate::pin::sha256_hex(b)
}

/// The inputs whose bytes are missing or do not match their pins.
pub fn check_inputs(natives: &BTreeMap<String, Vec<u8>>) -> Vec<String> {
    INPUT_NATIVES
        .iter()
        .filter(|(n, _, _, d)| natives.get(*n).map(|b| sha_hex(b)) != Some(d.to_string()))
        .map(|(n, ..)| n.to_string())
        .collect()
}

fn ext(form: Form) -> &'static str {
    form.label()
}

/// A safetensors file with header text `h` (padded with spaces to 8 bytes) and `data`.
fn st_raw(h: &str, data: &[u8]) -> Vec<u8> {
    let mut h = h.as_bytes().to_vec();
    while !(8 + h.len()).is_multiple_of(8) {
        h.push(b' ');
    }
    let mut o = (h.len() as u64).to_le_bytes().to_vec();
    o.extend(h);
    o.extend_from_slice(data);
    o
}

fn data_of(set: &TensorSet) -> Vec<u8> {
    set.clone().sorted().tensors.iter().flat_map(|t| t.data.clone()).collect()
}

/// The profile tensors of the native `n` under `b`, with the metadata of `form`.
fn profile_set(b: &str, native: &[u8], form: Form) -> TensorSet {
    let bd = binding::by_name(b).expect("registered");
    let mut s = (bd.from_native)(native).expect("a native that reads");
    s.metadata = binding::profile_metadata(b, form);
    s.sorted()
}

fn edit(set: &TensorSet, name: &str, f: impl FnOnce(&mut Tensor)) -> TensorSet {
    let mut s = set.clone();
    f(s.tensors.iter_mut().find(|t| t.name == name).expect("a tensor of the set"));
    s
}

/// The units of an NNC bitstream, as `(offset, size)`.
fn nnc_units(b: &[u8]) -> Vec<(usize, usize)> {
    let mut v = Vec::new();
    let mut at = 0;
    while at < b.len() {
        let size = if b[at] & 0x80 == 0 { (u16::from_be_bytes([b[at], b[at + 1]]) & 0x7fff) as usize } else { (u32::from_be_bytes([b[at], b[at + 1], b[at + 2], b[at + 3]]) & 0x7fff_ffff) as usize };
        v.push((at, size));
        at += size;
    }
    v
}

/// One case: a vector entry and the files it adds.
struct Case {
    v: Value,
}

/// Generate the corpus from `natives` (by name, matching [`INPUT_NATIVES`]): file name →
/// bytes, `vectors.json` included.
pub fn generate(natives: &BTreeMap<String, Vec<u8>>) -> BTreeMap<String, Vec<u8>> {
    assert!(check_inputs(natives).is_empty(), "the inputs must match their pins");
    let mut files: BTreeMap<String, Vec<u8>> = BTreeMap::new();
    let mut cases: Vec<Case> = Vec::new();
    let add = |files: &mut BTreeMap<String, Vec<u8>>, path: String, bytes: Vec<u8>| {
        assert!(files.insert(path.clone(), bytes).is_none(), "{path} twice");
        path
    };

    // Every native in every form: the encode, then the load.
    for &(name, b, _, _) in INPUT_NATIVES {
        let native = &natives[name];
        for form in Form::ALL {
            let file = encode(b, native, form).unwrap_or_else(|e| panic!("{name} in {}: {e}", form.label()));
            let (got_b, got) = load(form, &file).expect("loads");
            assert_eq!((got_b, &got), (b, native));
            let path = add(&mut files, format!("files/{name}.{}", ext(form)), file);
            cases.push(Case { v: json!({"name": format!("{name}_{}_encode", ext(form)), "op": "encode", "form": form.capability(), "native": name, "binding": b, "outcome": "accept", "file": path}) });
            cases.push(Case { v: json!({"name": format!("{name}_{}_load", ext(form)), "op": "load", "form": form.capability(), "file": path, "outcome": "accept", "binding": b, "native": name}) });
        }
    }

    // Native artifacts that do not read, and one no NNC integer payload carries.
    let mut trailing = natives["seeded"].clone();
    trailing.push(0);
    let mut badtag = natives["hyper_h_s"].clone();
    badtag[20] = 9; // the first layer's tag
    let mut wide = natives["hyper_h_s"].clone();
    // hyper_h_s: a conv layer first (tag 1, five u32), then its i16 weights, then mult and
    // bias (i64): make the first bias 2^40.
    let (inc, out, k) = (u32::from_le_bytes(wide[21..25].try_into().unwrap()) as usize, u32::from_le_bytes(wide[25..29].try_into().unwrap()) as usize, u32::from_le_bytes(wide[29..33].try_into().unwrap()) as usize);
    let bias0 = 21 + 20 + 2 * inc * out * k * k + 8 * out;
    wide[bias0..bias0 + 8].copy_from_slice(&(1i64 << 40).to_le_bytes());
    let mut gen_natives: Vec<(&str, &str, Vec<u8>)> = vec![("trailing", "wsr1/1", trailing), ("badtag", "whm1/1", badtag), ("wide", "whm1/1", wide)];
    for (n, b, bytes) in gen_natives.drain(..) {
        let ext_n = if b == "wsr1/1" { "wsr" } else { "whm" };
        let path = add(&mut files, format!("files/natives/{n}.{ext_n}"), bytes.clone());
        for form in Form::ALL {
            let r = encode(b, &bytes, form);
            let v = match r {
                Ok(file) => {
                    let fpath = add(&mut files, format!("files/{n}.{}", ext(form)), file);
                    json!({"name": format!("{n}_{}_encode", ext(form)), "op": "encode", "form": form.capability(), "native_file": path, "binding": b, "outcome": "accept", "file": fpath})
                }
                Err(e) => json!({"name": format!("{n}_{}_encode", ext(form)), "op": "encode", "form": form.capability(), "native_file": path, "binding": b, "outcome": "refuse", "code": e.code}),
            };
            cases.push(Case { v });
        }
    }

    // Files refused on load: built from the seeded prior's tensors.
    let seeded = &natives["seeded"];
    let st_set = profile_set("wsr1/1", seeded, Form::Safetensors);
    let st = safetensors::write(&st_set).unwrap();
    let st_h = safetensors::header(&st_set).unwrap();
    let st_d = data_of(&st_set);
    let gg_set = profile_set("wsr1/1", seeded, Form::Gguf);
    let gg = gguf::write(&gg_set).unwrap();
    let nn_set = profile_set("wsr1/1", seeded, Form::Nnc);
    let nn = nnc::write(&nn_set).unwrap();

    let mut neg: Vec<(&str, Form, Vec<u8>)> = Vec::new();
    // safetensors
    neg.push(("st_truncated", Form::Safetensors, st[..6].to_vec()));
    let mut x = st.clone();
    x[..8].copy_from_slice(&(safetensors::MAX_HEADER + 1).to_le_bytes());
    neg.push(("st_header_too_large", Form::Safetensors, x));
    let mut x = st.clone();
    x[..8].copy_from_slice(&(st.len() as u64).to_le_bytes());
    neg.push(("st_header_past_end", Form::Safetensors, x));
    let first = st_h.find("\"act_hi\":").expect("act_hi");
    let act = &st_h[first..st_h[first..].find('}').map(|e| first + e + 1).expect("}")];
    neg.push(("st_duplicate_member", Form::Safetensors, st_raw(&format!("{},{act}}}", &st_h[..st_h.len() - 1]), &st_d)));
    neg.push(("st_not_json", Form::Safetensors, st_raw("{\"act_hi\":", &st_d)));
    neg.push(("st_extra_member", Form::Safetensors, st_raw(&st_h.replacen("\"data_offsets\":[0,2]}", "\"data_offsets\":[0,2],\"x\":0}", 1), &st_d)));
    neg.push(("st_dtype_unknown", Form::Safetensors, st_raw(&st_h.replacen("\"dtype\":\"I16\"", "\"dtype\":\"Q4_0\"", 1), &st_d)));
    let mut gap = st_d.clone();
    gap.insert(0, 0);
    neg.push(("st_offset_gap", Form::Safetensors, st_raw(&st_h, &gap)));
    neg.push(("st_leading_space", Form::Safetensors, st_raw(&format!(" {st_h}"), &st_d)));
    {
        // The tensors' members in descending order of their names; offsets unchanged.
        let doc = crate::strict_json::read(st_h.as_bytes()).unwrap();
        let mut m: Vec<String> = doc.as_obj().unwrap().iter().filter(|(k, _)| k != safetensors::METADATA).map(|(k, _)| {
            let s = st_h.find(&format!("{}:", crate::strict_json::quote(k))).unwrap();
            let e = s + st_h[s..].find('}').unwrap() + 1;
            st_h[s..e].to_string()
        }).collect();
        m.reverse();
        let meta_end = st_h.find("},").unwrap() + 1;
        neg.push(("st_member_order", Form::Safetensors, st_raw(&format!("{},{}}}", &st_h[..meta_end], m.join(",")), &st_d)));
    }
    let mut s = st_set.clone();
    s.metadata.remove(super::KEY_PROFILE);
    neg.push(("st_profile_absent", Form::Safetensors, safetensors::write(&s).unwrap()));
    let mut s = st_set.clone();
    s.metadata.insert(super::KEY_BINDING.into(), "wsr9/1".into());
    neg.push(("st_binding_unknown", Form::Safetensors, safetensors::write(&s).unwrap()));
    let s = edit(&st_set, "act_hi", |t| t.dtype = DType::F16);
    neg.push(("st_float_tensor", Form::Safetensors, safetensors::write(&s).unwrap()));
    let mut s = st_set.clone();
    s.tensors.retain(|t| t.name != "layer.0.shift");
    neg.push(("st_missing_tensor", Form::Safetensors, safetensors::write(&s).unwrap()));
    let s = edit(&st_set, "act_hi", |t| t.shape = vec![1, 1]);
    neg.push(("st_wrong_shape", Form::Safetensors, safetensors::write(&s).unwrap()));
    let s = edit(&st_set, "layer.0.shift", |t| t.data = vec![40]);
    neg.push(("st_shift_range", Form::Safetensors, safetensors::write(&s).unwrap()));
    let s = edit(&st_set, "layer.0.mult", |t| t.data[..4].copy_from_slice(&0i32.to_le_bytes()));
    neg.push(("st_mult_zero", Form::Safetensors, safetensors::write(&s).unwrap()));
    // GGUF
    let mut x = gg.clone();
    x[3] = b'G';
    neg.push(("gg_bad_magic", Form::Gguf, x));
    let mut x = gg.clone();
    x[4..8].copy_from_slice(&2u32.to_le_bytes());
    neg.push(("gg_version_2", Form::Gguf, x));
    neg.push(("gg_truncated_header", Form::Gguf, gg[..40].to_vec()));
    {
        // Keys are written in order from byte 24: general.architecture, wai.binding,
        // wai.profile. Rename the second to the third: a key twice.
        let kv2 = 24 + 8 + "general.architecture".len() + 4 + 8 + "wai".len();
        let mut x = gg.clone();
        x[kv2 + 8..kv2 + 8 + 11].copy_from_slice(b"wai.profile");
        neg.push(("gg_duplicate_key", Form::Gguf, x));
        let mut x = gg.clone();
        x[kv2 + 8] = b'W';
        neg.push(("gg_bad_key", Form::Gguf, x));
    }
    let tinfo = gg.windows(8 + 6).position(|w| w == [&6u64.to_le_bytes()[..], b"act_hi"].concat().as_slice()).expect("act_hi info");
    let mut x = gg.clone();
    x[tinfo + 14..tinfo + 18].copy_from_slice(&5u32.to_le_bytes());
    neg.push(("gg_five_dims", Form::Gguf, x));
    let ty_at = tinfo + 14 + 4 + 8;
    let mut x = gg.clone();
    x[ty_at..ty_at + 4].copy_from_slice(&2u32.to_le_bytes());
    neg.push(("gg_q4_0", Form::Gguf, x));
    let mut x = gg.clone();
    // The second tensor's offset (act_hi is first, at 0; layer.0.bias follows at 32).
    let off2 = gg.windows(8 + 12).position(|w| w == [&12u64.to_le_bytes()[..], b"layer.0.bias"].concat().as_slice()).expect("layer.0.bias info") + 20 + 4 + 8 + 4;
    x[off2..off2 + 8].copy_from_slice(&33u64.to_le_bytes());
    neg.push(("gg_misaligned", Form::Gguf, x));
    {
        // general.alignment = 32 written explicitly, after general.architecture: reads, is
        // not canonical.
        let mut kv = 17u64.to_le_bytes().to_vec();
        kv.extend_from_slice(b"general.alignment");
        kv.extend_from_slice(&4u32.to_le_bytes());
        kv.extend_from_slice(&32u32.to_le_bytes());
        let at = 24 + 8 + "general.architecture".len() + 4 + 8 + 3;
        neg.push(("gg_alignment_explicit", Form::Gguf, gguf_insert_kv(&gg, at, &kv)));
    }
    let mut x = gg.clone();
    x.push(0);
    neg.push(("gg_trailing_byte", Form::Gguf, x));
    let mut s = gg_set.clone();
    s.metadata.remove("general.architecture");
    neg.push(("gg_no_architecture", Form::Gguf, gguf::write(&s).unwrap()));
    {
        let (hend, dstart) = (gguf_header_end(&gg), gguf_data_start(&gg));
        assert!(hend < dstart, "the seeded prior's GGUF header is padded");
        let mut x = gg.clone();
        x[hend] = 1;
        neg.push(("gg_nonzero_padding", Form::Gguf, x));
    }
    // NNC
    let units = nnc_units(&nn);
    neg.push(("nnc_truncated", Form::Nnc, nn[..nn.len() - 1].to_vec()));
    neg.push(("nnc_starts_mps", Form::Nnc, nn[units[1].0..].to_vec()));
    let mut x = nn.clone();
    let mps_last = units[1].0 + units[1].1 - 1;
    assert_eq!(x[mps_last], 0x80);
    x[mps_last] = 0;
    neg.push(("nnc_alignment_bit", Form::Nnc, x));
    let ndu0 = units[3];
    let mut x = nn.clone();
    x[ndu0.0 + 3] = (x[ndu0.0 + 3] & 0x07) | ((nnc::NNR_PT_FLOAT as u8) << 3);
    neg.push(("nnc_pt_float", Form::Nnc, x));
    {
        // dq_flag: the first bit after the topology element name's zero byte.
        let name_end = nn[ndu0.0 + 4..].iter().position(|&c| c == 0).unwrap() + ndu0.0 + 4;
        let mut x = nn.clone();
        x[name_end + 1] |= 0x80;
        neg.push(("nnc_dq_flag", Form::Nnc, x));
    }
    {
        // An NNR_LPS unit after NNR_MPS: size 4, type 2, independently decodable, then
        // lps_self_contained_flag 1 and seven reserved zero bits.
        let mut x = nn[..units[1].0 + units[1].1].to_vec();
        x.extend_from_slice(&[0x00, 0x04, (nnc::NNR_LPS as u8) << 2 | 0x02, 0x80]);
        x.extend_from_slice(&nn[units[2].0..]);
        neg.push(("nnc_lps_unit", Form::Nnc, x));
    }
    {
        // NNR_STR in the 31-bit size form: reads, is not canonical.
        let mut x = vec![0x80, 0x00, 0x00, 0x06];
        x.extend_from_slice(&nn[2..4]);
        x.extend_from_slice(&nn[4..]);
        neg.push(("nnc_long_size", Form::Nnc, x));
    }
    let raw = edit(&nn_set, "act_hi", |t| {
        t.dtype = DType::F32;
        t.data = 200f32.to_le_bytes().to_vec();
    });
    neg.push(("nnc_raw_float", Form::Nnc, nnc::write(&raw).unwrap()));
    {
        // One bit of the entropy-coded payload of the first weight tensor flipped.
        let w0 = units.iter().copied().find(|&(o, s)| nn[o..o + s].windows(15).any(|w| w == b"layer.0.weight\0")).unwrap();
        let mut x = nn.clone();
        x[w0.0 + w0.1 - 5] ^= 0x08;
        neg.push(("nnc_payload_bit", Form::Nnc, x));
    }

    let mut seen_codes = std::collections::BTreeSet::new();
    for c in &cases {
        if let Some(code) = c.v.get("code").and_then(Value::as_str) {
            seen_codes.insert(code.to_string());
        }
    }
    for (name, form, bytes) in neg {
        let code = load(form, &bytes).expect_err(name).code;
        seen_codes.insert(code.to_string());
        let path = add(&mut files, format!("files/neg/{name}.{}", ext(form)), bytes);
        cases.push(Case { v: json!({"name": name, "op": "load", "form": form.capability(), "file": path, "outcome": "refuse", "code": code}) });
    }
    let missing: Vec<&&str> = CODES.iter().filter(|c| !seen_codes.contains(**c)).collect();
    assert!(missing.is_empty(), "the corpus gives no vector the codes {missing:?}");

    let natives_v: Vec<Value> = INPUT_NATIVES.iter().map(|(n, b, p, d)| json!({"name": n, "binding": b, "path": p, "sha256": d})).collect();
    let files_v: Vec<Value> = files.iter().map(|(p, b)| json!({"file": p, "sha256": sha_hex(b)})).collect();
    let doc = json!({
        "about": "the conformance corpus of extensions/prior-carriage.md §14, written by wai_prior_forms_vectors gen",
        "natives": natives_v,
        "files": files_v,
        "vectors": cases.into_iter().map(|c| c.v).collect::<Vec<_>>(),
    });
    let mut text = serde_json::to_string_pretty(&doc).expect("json");
    text.push('\n');
    files.insert("vectors.json".into(), text.into_bytes());
    files
}

/// The end of a GGUF file's tensor descriptions.
fn gguf_header_end(gg: &[u8]) -> usize {
    let rd = |p: usize| u64::from_le_bytes(gg[p..p + 8].try_into().unwrap()) as usize;
    let (n_t, n_kv) = (rd(8), rd(16));
    let mut p = 24usize;
    for _ in 0..n_kv {
        p += 8 + rd(p) + 4; // the key and the value type (every value here is a string)
        p += 8 + rd(p);
    }
    for _ in 0..n_t {
        p += 8 + rd(p);
        let nd = u32::from_le_bytes(gg[p..p + 4].try_into().unwrap()) as usize;
        p += 4 + 8 * nd + 4 + 8;
    }
    p
}

/// The offset of a canonical GGUF file's tensor data: the header's end, rounded up to 32.
fn gguf_data_start(gg: &[u8]) -> usize {
    gguf_header_end(gg).div_ceil(32) * 32
}

/// `gg` with the key-value pair `kv` inserted at byte `at` of its header, the pair count
/// raised by one and the header re-padded so the data still begins at a multiple of 32.
fn gguf_insert_kv(gg: &[u8], at: usize, kv: &[u8]) -> Vec<u8> {
    let mut hdr = gg[..gguf_header_end(gg)].to_vec();
    hdr.splice(at..at, kv.iter().copied());
    let n = u64::from_le_bytes(hdr[16..24].try_into().unwrap()) + 1;
    hdr[16..24].copy_from_slice(&n.to_le_bytes());
    while !hdr.len().is_multiple_of(32) {
        hdr.push(0);
    }
    hdr.extend_from_slice(&gg[gguf_data_start(gg)..]);
    hdr
}

/// Check a corpus directory's files (`files`, by path, with the natives under their
/// [`INPUT_NATIVES`] names) against this module: every file regenerates byte for byte, and
/// every vector reaches its outcome. Returns the failures.
pub fn verify(corpus: &BTreeMap<String, Vec<u8>>, natives: &BTreeMap<String, Vec<u8>>) -> Vec<String> {
    let mut bad = Vec::new();
    let missing = check_inputs(natives);
    if !missing.is_empty() {
        return vec![format!("natives missing or not matching their pins: {missing:?}")];
    }
    let want = generate(natives);
    for (p, b) in &want {
        match corpus.get(p) {
            Some(x) if x == b => {}
            Some(_) => bad.push(format!("{p} does not regenerate byte for byte")),
            None => bad.push(format!("{p} is missing")),
        }
    }
    for p in corpus.keys() {
        if !want.contains_key(p) {
            bad.push(format!("{p} is not a file of the corpus"));
        }
    }
    bad
}
