//! Conformance corpus for the integer-synthesis neural decoder
//! (`neural-conformance/`). This is the one neural decoder that can make
//! the WAI promise the reconstruct classes make: **bit-identical output on
//! every machine, no tolerance** — because the synthesis is pure `i64`
//! arithmetic, not float ONNX.
//!
//! The float decoders (`bmshj2018`, `mbt2018_mean`, `elic`, `mlicpp`, …)
//! are deliberately *not* here: their ONNX synthesis is float, so two
//! conforming sinks on different hardware may differ in the low bits — a
//! byte-equality corpus would be a false promise for them. The integer
//! decoder is the member where `intent = replicate` is sound, so it is the
//! member with a cross-arch byte-equality corpus.
//!
//! Each vector is a realistic multi-layer integer decoder: a committed
//! `<name>.int_synth.json` (the model) + a `<name>.bin` payload (header +
//! int8 latents) + a frozen digest of the decoded RGB: the SHA-256 in
//! `sha256.json` (the conformance digest, SPEC §7 *Digests*) and a domain-separated
//! BLAKE3 in `expected_hashes.json` (a diagnostic). `verify` re-decodes and checks
//! both — on every platform leg of the conformance workflow.
//!
//!   wai_int_synth_vectors gen <dir>  |  verify <dir>

use std::path::Path;
use wai::int_synth::decode;

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)) => generate(Path::new(d)),
        (Some("verify"), Some(d)) => verify(Path::new(d)),
        _ => {
            eprintln!("usage: wai_int_synth_vectors <gen|verify> <corpus-dir>");
            std::process::exit(2);
        }
    }
}

const RGB8: &str = "RGB8: rows top to bottom, pixels left to right, R, G, B bytes";

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

/// A small deterministic integer PRNG (SplitMix64) — used only to author
/// fixture weights/latents reproducibly. No floats, no platform dependence.
struct Rng(u64);
impl Rng {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E3779B97F4A7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
        z ^ (z >> 31)
    }
    /// Signed value in `[-range, range]`.
    fn signed(&mut self, range: i64) -> i32 {
        (self.next() as i64 % (2 * range + 1) - range) as i32
    }
}

/// One transposed-conv layer description as JSON.
struct Layer {
    in_ch: usize,
    out_ch: usize,
    k: usize,
    stride: usize,
    shift: u32,
    relu: bool,
}

fn layer_json(l: &Layer, rng: &mut Rng) -> serde_json::Value {
    let n = l.in_ch * l.out_ch * l.k * l.k;
    // Small weights keep the i64 accumulators well inside range and the
    // output spread across the byte gamut after the requantize shift.
    let weight: Vec<i32> = (0..n).map(|_| rng.signed(6)).collect();
    let bias: Vec<i32> = (0..l.out_ch).map(|_| rng.signed(64)).collect();
    serde_json::json!({
        "in_ch": l.in_ch, "out_ch": l.out_ch, "k": l.k, "stride": l.stride,
        "pad": 0, "out_pad": 0, "weight": weight, "bias": bias,
        "shift": l.shift, "relu": l.relu,
    })
}

/// `(name, seed, S_log2, in_channels, layer stack)`. Each stack must end at
/// 3 output channels and upsample S→H,W consistently with the header.
fn vectors() -> Vec<(&'static str, u64, u8, usize, Vec<Layer>)> {
    vec![
        // 8×8 latent, 4 channels → ×2 → ×2 = 32×32 RGB, two real layers with
        // an integer ReLU between them. The canonical decoder shape.
        (
            "gradient",
            0x1234_5678,
            3, // S = 8
            4,
            vec![
                Layer { in_ch: 4, out_ch: 6, k: 4, stride: 2, shift: 4, relu: true },
                Layer { in_ch: 6, out_ch: 3, k: 4, stride: 2, shift: 5, relu: false },
            ],
        ),
        // 4×4 latent, 8 channels → ×4 (single deep upsample) = 16×16 RGB.
        (
            "bloom",
            0xC0FF_EE42,
            2, // S = 4
            8,
            vec![Layer { in_ch: 8, out_ch: 3, k: 6, stride: 4, shift: 6, relu: false }],
        ),
        // 16×16 latent, 3 channels → ×2 → ×2 → ×1 (refine) = 64×64 RGB, three
        // layers, two ReLUs — exercises a deeper integer stack.
        (
            "facade",
            0x5EED_1010,
            4, // S = 16
            3,
            vec![
                Layer { in_ch: 3, out_ch: 8, k: 4, stride: 2, shift: 4, relu: true },
                Layer { in_ch: 8, out_ch: 6, k: 4, stride: 2, shift: 5, relu: true },
                Layer { in_ch: 6, out_ch: 3, k: 3, stride: 1, shift: 4, relu: false },
            ],
        ),
    ]
}

/// Build the int8 latent payload (header + channel-major int8 latents).
fn payload(s_log2: u8, c: usize, h: u32, w: u32, rng: &mut Rng) -> Vec<u8> {
    let s = 1usize << s_log2;
    let mut p = Vec::new();
    p.extend_from_slice(&h.to_le_bytes());
    p.extend_from_slice(&w.to_le_bytes());
    p.extend_from_slice(&(c as u16).to_le_bytes());
    p.push(s_log2);
    p.push(0);
    for _ in 0..c * s * s {
        p.push(rng.signed(120) as i8 as u8);
    }
    p
}

/// Compute the (H, W) a layer stack yields from an S×S latent (no pad/out_pad).
fn output_hw(s: usize, layers: &[Layer]) -> (u32, u32) {
    let mut d = s;
    for l in layers {
        d = (d - 1) * l.stride + l.k;
    }
    (d as u32, d as u32)
}

fn generate(dir: &Path) {
    let v = dir.join("vectors");
    std::fs::create_dir_all(&v).expect("mkdir");
    let mut manifest = serde_json::Map::new();
    // SPEC §7 *Digests*: the SHA-256 of each decoded RGB8 buffer, in a sidecar so
    // expected_hashes.json keeps its BLAKE3 values unchanged.
    let mut sha = serde_json::Map::new();
    for (name, seed, s_log2, c, layers) in vectors() {
        let (h, w) = output_hw(1 << s_log2, &layers);
        // One stream authors weights then the latent, deterministically.
        let mut rng = Rng(seed);
        let layers_json: Vec<serde_json::Value> = layers.iter().map(|l| layer_json(l, &mut rng)).collect();
        let model = serde_json::json!({ "layers": layers_json });
        let pay = payload(s_log2, c, h, w, &mut rng);

        // Write the model next to a `model` stub so `decode` finds the sibling
        // `int_synth.json`; mirror the per-vector names into the manifest.
        let model_dir = v.join(name);
        std::fs::create_dir_all(&model_dir).expect("mkdir vec");
        std::fs::write(model_dir.join("int_synth.json"), serde_json::to_vec(&model).unwrap()).expect("write model");
        std::fs::write(model_dir.join("latent.bin"), &pay).expect("write payload");

        let img = decode(&pay, &model_dir.join("model")).expect("decode fixture");
        let mut hasher = blake3::Hasher::new();
        hasher.update(b"wai:int-synth-rgb\x01");
        hasher.update(&img.rgb);
        let rgb_hash = hex(hasher.finalize().as_bytes());
        sha.insert(name.to_string(), serde_json::json!({ "buffer": RGB8, "sha256": wai::pin::sha256_hex(&img.rgb) }));

        manifest.insert(
            name.to_string(),
            serde_json::json!({
                "s_log2": s_log2,
                "in_channels": c,
                "layers": layers.len(),
                "width": img.width,
                "height": img.height,
                "rgb_hash": rgb_hash,
            }),
        );
    }
    // The committed trained vector: a model trained in float on pictures that
    // tools/wai_int_synth_real.py generates from a seed, and quantised after
    // training; that tool writes it, not this one. Decode it through the
    // same integer path and CROSS-CHECK against the python integer-forward
    // reference (ref_rgb.bin) — Rust must match it byte-for-byte (max|Δ|=0),
    // proving the trained model decodes bit-identically across implementations.
    let photo = v.join("photo");
    if photo.join("int_synth.json").exists() && photo.join("latent.bin").exists() {
        let pay = std::fs::read(photo.join("latent.bin")).expect("photo latent");
        let img = decode(&pay, &photo.join("model")).expect("decode photo");
        if let Ok(refrgb) = std::fs::read(photo.join("ref_rgb.bin")) {
            assert_eq!(img.rgb, refrgb, "Rust int_synth diverges from the python integer-forward reference");
            println!("trained vector photo: Rust == python integer-forward (max|Δ|=0)");
        }
        let model: serde_json::Value =
            serde_json::from_slice(&std::fs::read(photo.join("int_synth.json")).unwrap()).unwrap();
        let n_layers = model["layers"].as_array().unwrap().len();
        let mut hasher = blake3::Hasher::new();
        hasher.update(b"wai:int-synth-rgb\x01");
        hasher.update(&img.rgb);
        let mut entry = serde_json::json!({
            "trained": true, "width": img.width, "height": img.height,
            "layers": n_layers, "rgb_hash": hex(hasher.finalize().as_bytes()),
        });
        if let Ok(bytes) = std::fs::read(photo.join("meta.json"))
            && let Ok(meta) = serde_json::from_slice::<serde_json::Value>(&bytes)
        {
            for k in ["psnr_db", "bpp", "in_channels", "s_log2"] {
                entry[k] = meta[k].clone();
            }
        }
        manifest.insert("photo".to_string(), entry);
        sha.insert("photo".to_string(), serde_json::json!({ "buffer": RGB8, "sha256": wai::pin::sha256_hex(&img.rgb) }));
    }
    std::fs::write(
        dir.join("sha256.json"),
        serde_json::to_string_pretty(&serde_json::Value::Object(sha)).unwrap() + "\n",
    )
    .expect("sha256 sidecar");
    std::fs::write(
        dir.join("expected_hashes.json"),
        serde_json::to_string_pretty(&serde_json::Value::Object(manifest)).unwrap() + "\n",
    )
    .expect("manifest");
    println!("generated int_synth vectors into {}", dir.display());
}

fn verify(dir: &Path) {
    let m: serde_json::Value =
        serde_json::from_slice(&std::fs::read(dir.join("expected_hashes.json")).expect("manifest")).expect("parse");
    let sha: serde_json::Value =
        serde_json::from_slice(&std::fs::read(dir.join("sha256.json")).expect("sha256.json")).expect("parse sha256.json");
    let v = dir.join("vectors");
    let mut fails = 0;
    for (name, spec) in m.as_object().expect("obj") {
        let model_dir = v.join(name);
        let pay = std::fs::read(model_dir.join("latent.bin")).expect("payload");
        let img = decode(&pay, &model_dir.join("model")).expect("decode");
        let mut hasher = blake3::Hasher::new();
        hasher.update(b"wai:int-synth-rgb\x01");
        hasher.update(&img.rgb);
        let got = hex(hasher.finalize().as_bytes());
        let ok = got == spec["rgb_hash"].as_str().unwrap()
            && Some(wai::pin::sha256_hex(&img.rgb).as_str()) == sha[name]["sha256"].as_str()
            && img.width == spec["width"].as_u64().unwrap() as u32
            && img.height == spec["height"].as_u64().unwrap() as u32;
        println!("{} {name}  ({}×{}, {} layers)", if ok { "PASS" } else { "FAIL" },
            img.width, img.height, spec["layers"].as_u64().unwrap());
        if !ok {
            fails += 1;
        }
    }
    if fails > 0 {
        eprintln!("{fails} failed — integer synthesis is NOT bit-exact here");
        std::process::exit(1);
    }
    println!("all vectors pass: integer-synthesis decode is bit-exact on this machine");
}
