//! `wai_stage_bench` — measure one staged set natively (`extensions/staged-delivery.md`;
//! the measurement method of the staged-delivery design, §9).
//!
//!   wai_stage_bench --source <file> [--set <dir>] [--limit OUT,SYMBOLS,WORKING,MACS] [--runs N] [--out report.json] <envelope>…
//!
//! The envelopes are a stage set: bases (no `refines`), decoded through
//! `staged::decode_base` against the parameter set in `--set` (none for
//! `wai.video.int_motion`), and their refinement layers. Every base is decoded under the
//! cost limit `--limit` gives (by default `staged::TOOL_DEFAULT_LIMIT`), applied before
//! any symbol is decoded; every layer is held to that limit's output, symbols and
//! working bytes, and every chain to `staged::TOOL_DEFAULT_CHAIN_BUDGET`. Stage 0 is every base decoded;
//! stage `k` is every layer whose `refines.layer` is `k` applied, in the order given.
//! `--source` holds what the units approximate, at their final shape: each unit's
//! canonical buffer as the last stage shapes it, base by base in the order given and
//! unit by unit. A unit whose stage holds another number of samples (below an up2 or a
//! rate layer) has no quality figure at that stage, and the clip's figure counts the
//! units that have one (`units_measured`).
//!
//! For each stage it reports:
//! - **bytes**: the envelopes' bytes as delivered, split into `WAI1` framing, manifest
//!   and payload, and the cumulative bytes through that stage. JSON and base64 are
//!   never counted; there is no receipt, index or transport byte in a set on disk.
//! - **quality**: for 8-bit units, PSNR = 10·log10(255²·N / SSE) per unit and over the
//!   whole clip (N samples, SSE the exact integer sum of squared differences); for
//!   16-bit units, SNR = 10·log10(Σx² / Σ(x−y)²) on the samples. A unit equal to its
//!   source has no finite figure and is reported as `null` with `"exact": true`.
//! - **time**: the microseconds the stage took (decoding the bases; applying the
//!   layers), median and interquartile range over `--runs` repetitions (at least 20),
//!   single-threaded, in this process.
//! - **the direct comparator** (§9.3): for stage `k ≥ 1` whose layers share one step
//!   `q` and the identity predictor, one layer per base coded straight from the base's
//!   stage 0 to the source at that `q` — the cost of reaching stage `k` in one layer
//!   instead of `k`. Its bytes and quality are reported beside the stage's.
//!
//! Every unit's stage hash at every stage is reported too, so a run can be checked
//! against `wai stage apply --json`. Energy is reported as `unmetered`: this tool reads
//! no meter. The figures describe this set on this machine and build, nothing wider.

use std::collections::BTreeMap;
use std::process::ExitCode;
use std::time::Instant;

use serde_json::{json, Value};
use wai::container::{Prev, Refines, Wai};
use wai::staged::wir1::{self, Below, Next, Predictor, RefineParams};
use wai::staged::{BaseInfo, ChainConfig, HoldBound, ReceiptEvidence, Samples, StageSink, UnitBuf};

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

fn opt<'a>(args: &'a [String], flag: &str) -> Option<&'a str> {
    args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).map(String::as_str)
}

fn positional(args: &[String]) -> Vec<&str> {
    let mut out = Vec::new();
    let mut skip = false;
    for a in args {
        if skip {
            skip = false;
        } else if a.starts_with("--") {
            skip = true;
        } else {
            out.push(a.as_str());
        }
    }
    out
}

fn read(path: &str) -> Result<Vec<u8>, String> {
    std::fs::read(path).map_err(|e| format!("read {path}: {e}"))
}

/// The cost limit: `--limit OUT,SYMBOLS,WORKING,MACS`, or the tools' stated default.
fn limit(args: &[String]) -> Result<wai::int_payload::DeclaredCost, String> {
    opt(args, "--limit").map_or(Ok(wai::staged::TOOL_DEFAULT_LIMIT), wai::staged::parse_limit)
}

/// The sink's configuration: each layer held to the cost limit's output, symbols and
/// working bytes, each chain to the stated chain budget.
fn cfg(limit: &wai::int_payload::DeclaredCost) -> ChainConfig {
    ChainConfig {
        bound: HoldBound { max_envelopes: 4096, max_age_ms: u64::MAX },
        cost_limit: Some(wai::staged::layer_limit(limit)),
        budget: wai::staged::TOOL_DEFAULT_CHAIN_BUDGET,
        record_budget: u64::MAX,
        live: None,
    }
}

/// Add a base to the sink, refusing one the sink does not take whole (its units past
/// the chain budget).
fn add_base(sink: &mut StageSink, file: &str, info: BaseInfo, units: Vec<UnitBuf>) -> Result<(), String> {
    let h = info.content_hash;
    let reports = sink.add_base(info, units);
    if sink.chain(&h).is_none() {
        let why = reports.iter().find_map(|r| r.envelope).map_or("not taken", |r| r.code());
        return Err(format!("{file}: the base is refused ({why})"));
    }
    Ok(())
}

/// One envelope of the set, read once.
struct Env {
    file: String,
    bytes: Vec<u8>,
    /// 0 for a base, `refines.layer` for a layer.
    stage: u8,
    manifest: usize,
    payload: usize,
    /// For a layer: its step and predictor.
    q: Option<(u32, Predictor)>,
}

fn read_env(file: &str) -> Result<Env, String> {
    let bytes = read(file)?;
    let w = Wai::from_bytes(&bytes).map_err(|e| format!("{file}: {e}"))?;
    let manifest = u32::from_le_bytes(bytes[4..8].try_into().map_err(|_| format!("{file}: no manifest length"))?) as usize;
    let (stage, q) = match w.refines() {
        None => (0, None),
        Some(Err(e)) => return Err(format!("{file}: {e}")),
        Some(Ok(r)) => {
            let kind = wai::staged::unit_kind(w.capability()).ok_or_else(|| format!("{file}: not a residual refinement"))?;
            let p = wir1::parse(&w.payload, kind).map_err(|e| format!("{file}: {e}"))?;
            (r.layer, Some((p.q, p.predictor)))
        }
    };
    Ok(Env { file: file.to_string(), stage, manifest, payload: w.payload.len(), q, bytes })
}

fn quantiles(mut v: Vec<f64>) -> Value {
    v.sort_by(|a, b| a.partial_cmp(b).expect("finite"));
    let at = |p: f64| {
        let x = p * (v.len() - 1) as f64;
        let (lo, hi) = (x.floor() as usize, x.ceil() as usize);
        v[lo] + (v[hi] - v[lo]) * (x - lo as f64)
    };
    json!({"median": round2(at(0.5)), "p25": round2(at(0.25)), "p75": round2(at(0.75))})
}

fn round2(x: f64) -> f64 {
    (x * 100.0).round() / 100.0
}

/// The squared error of a unit against its source, and its sample count and signal
/// energy (for SNR).
fn sse(s: &Samples, src: &[u8]) -> (u64, u64, u64) {
    match s {
        Samples::U8(v) => (v.iter().zip(src).map(|(a, b)| (*a as i64 - *b as i64).pow(2) as u64).sum(), v.len() as u64, 0),
        Samples::I16(v) => {
            let x: Vec<i64> = src.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]]) as i64).collect();
            let e = v.iter().zip(&x).map(|(a, b)| (*a as i64 - b).pow(2) as u64).sum();
            (e, v.len() as u64, x.iter().map(|b| (b * b) as u64).sum())
        }
    }
}

/// PSNR (8-bit) or SNR (16-bit) in dB, `None` for an exact match.
fn quality(err: u64, n: u64, energy: u64, pcm: bool) -> Option<f64> {
    if err == 0 {
        return None;
    }
    let db = if pcm { 10.0 * (energy as f64 / err as f64).log10() } else { 10.0 * (255.0 * 255.0 * n as f64 / err as f64).log10() };
    Some(round2(db))
}

fn quality_json(err: u64, n: u64, energy: u64, pcm: bool) -> Value {
    let mut v = json!({"sse": err, "samples": n});
    v[if pcm { "snr_db" } else { "psnr_db" }] = json!(quality(err, n, energy, pcm));
    v["exact"] = json!(err == 0);
    v
}

/// The refinement capability over a base's units (as `wai stage refine` picks it).
fn refine_cap(info: &BaseInfo, units: &[UnitBuf]) -> &'static str {
    use wai::codecs::{capability_media, Media};
    match units.first().map(|u| &u.samples) {
        Some(Samples::I16(_)) => wai::codecs::CAP_AUDIO_INT_REFINE,
        _ if capability_media(&info.capability) == Some(Media::Video) => wai::codecs::CAP_VIDEO_INT_REFINE,
        _ if capability_media(&info.capability) == Some(Media::Volumetric) => wai::codecs::CAP_SPLAT_INT_REFINE,
        _ => wai::codecs::CAP_IMAGE_INT_REFINE,
    }
}

fn below(s: &Samples) -> Below<'_> {
    match s {
        Samples::U8(v) => Below::U8(v),
        Samples::I16(v) => Below::I16(v),
    }
}

fn run(args: &[String]) -> Result<(), String> {
    let source = read(opt(args, "--source").ok_or("--source <file>")?)?;
    let runs: usize = opt(args, "--runs").unwrap_or("20").parse().map_err(|e| format!("--runs: {e}"))?;
    if runs < 20 {
        return Err("--runs: at least 20, so the quartiles mean something".into());
    }
    let set: BTreeMap<String, Vec<u8>> = match opt(args, "--set") {
        Some(dir) => wai::pin::read_dir_set(std::path::Path::new(dir)).map_err(|e| format!("--set {dir}: {e}"))?,
        None => BTreeMap::new(),
    };
    let files: BTreeMap<String, &[u8]> = set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
    let limit = limit(args)?;
    let envs: Vec<Env> = positional(args).iter().map(|f| read_env(f)).collect::<Result<_, _>>()?;
    let n_stages = envs.iter().map(|e| e.stage).max().ok_or("no envelopes")?;
    if !envs.iter().any(|e| e.stage == 0) {
        return Err("no base among the envelopes".into());
    }

    // A replay, untimed: every stage's units, each base's in order. Every layer must
    // apply whole: the bench measures a set delivered whole and in order.
    let mut bases: Vec<(BaseInfo, Vec<UnitBuf>)> = Vec::new();
    for e in envs.iter().filter(|e| e.stage == 0) {
        bases.push(wai::staged::decode_base(&e.bytes, &files, Some(&limit), None).map_err(|x| format!("{}: {x}", e.file))?);
    }
    let hashes: Vec<[u8; 32]> = bases.iter().map(|(i, _)| i.content_hash).collect();
    // stages[k][b][u]: the unit's samples, stage and stage hash at stage k.
    let mut stages_seen: Vec<Vec<Vec<(Samples, u8, Option<[u8; 32]>)>>> = Vec::new();
    {
        let mut sink = StageSink::new(cfg(&limit));
        for ((info, units), e) in bases.iter().zip(envs.iter().filter(|e| e.stage == 0)) {
            add_base(&mut sink, &e.file, info.clone(), units.clone())?;
        }
        let snap = |sink: &StageSink| -> Vec<Vec<(Samples, u8, Option<[u8; 32]>)>> {
            hashes.iter().map(|h| {
                let c = sink.chain(h).expect("added");
                c.units().iter().enumerate().map(|(u, s)| (s.buf.samples.clone(), s.stage, c.unit_sha256(u as u32))).collect()
            }).collect()
        };
        stages_seen.push(snap(&sink));
        for k in 1..=n_stages {
            for (i, e) in envs.iter().enumerate().filter(|(_, e)| e.stage == k) {
                let reports = sink.offer(&e.bytes, i as u64, ReceiptEvidence::None).map_err(|x| format!("{}: {x}", e.file))?;
                if let Some(bad) = reports.iter().find(|x| x.envelope.is_some() || !x.refused.is_empty() || !x.held.is_empty() || x.held_for_base) {
                    let why = bad.envelope.map(|r| r.code()).or_else(|| bad.refused.first().map(|&(_, r)| r.code())).unwrap_or("held");
                    return Err(format!("{}: not applied whole ({why}); the bench measures a set delivered whole and in order", e.file));
                }
            }
            stages_seen.push(snap(&sink));
        }
    }

    // The source, at the units' final shape: each unit's canonical buffer at the last
    // stage, base by base and unit by unit. A stage whose unit holds another number of
    // samples (below an up2 or rate layer) has no quality figure for that unit.
    let last = stages_seen.last().expect("stage 0");
    let mut src_of: Vec<Vec<&[u8]>> = Vec::new();
    let mut at = 0usize;
    for units in last {
        let mut v = Vec::new();
        for (samples, ..) in units {
            let n = samples.byte_len() as usize;
            v.push(source.get(at..at + n).ok_or("--source is shorter than the units at their last stage")?);
            at += n;
        }
        src_of.push(v);
    }
    if at != source.len() {
        return Err(format!("--source holds {} bytes; the units at their last stage hold {at}", source.len()));
    }
    let pcm = matches!(bases[0].1[0].samples, Samples::I16(_));
    let mut snapshots: Vec<Vec<Value>> = vec![Vec::new(); n_stages as usize + 1];
    let mut clip: Vec<(u64, u64, u64, usize)> = vec![(0, 0, 0, 0); n_stages as usize + 1];
    for (k, per_base) in stages_seen.iter().enumerate() {
        for (b, units) in per_base.iter().enumerate() {
            for (u, (samples, stage, sha)) in units.iter().enumerate() {
                let src = src_of[b][u];
                let mut q = if samples.byte_len() as usize == src.len() {
                    let (err, n, energy) = sse(samples, src);
                    clip[k].0 += err;
                    clip[k].1 += n;
                    clip[k].2 += energy;
                    clip[k].3 += 1;
                    quality_json(err, n, energy, pcm)
                } else {
                    json!({"sse": null, "samples": null, "exact": null, "why": "the unit's shape at this stage is not the source's"})
                };
                q["base"] = json!(hex(&hashes[b]));
                q["unit"] = json!(u);
                q["stage"] = json!(stage);
                q["stage_sha256"] = json!(sha.map(|x| hex(&x)));
                snapshots[k].push(q);
            }
        }
    }

    // Timed replays.
    let mut times: Vec<Vec<f64>> = vec![Vec::new(); n_stages as usize + 1];
    for _ in 0..runs {
        let mut sink = StageSink::new(cfg(&limit));
        let t = Instant::now();
        for e in envs.iter().filter(|e| e.stage == 0) {
            let (info, units) = wai::staged::decode_base(&e.bytes, &files, Some(&limit), None).map_err(|x| format!("{}: {x}", e.file))?;
            add_base(&mut sink, &e.file, info, units)?;
        }
        times[0].push(t.elapsed().as_secs_f64() * 1e6);
        for k in 1..=n_stages {
            let t = Instant::now();
            for (i, e) in envs.iter().enumerate().filter(|(_, e)| e.stage == k) {
                sink.offer(&e.bytes, i as u64, ReceiptEvidence::None).map_err(|x| format!("{}: {x}", e.file))?;
            }
            times[k as usize].push(t.elapsed().as_secs_f64() * 1e6);
        }
    }

    // The direct comparator for each stage whose layers share one step and identity.
    let direct = |k: u8| -> Result<Value, String> {
        let qs: Vec<(u32, Predictor)> = envs.iter().filter(|e| e.stage == k).filter_map(|e| e.q).collect();
        let Some(&(q, p)) = qs.first() else { return Ok(Value::Null) };
        if p != Predictor::Identity || qs.iter().any(|x| *x != (q, p)) {
            return Ok(json!({"q": null, "why": "the stage's layers do not share one step and the identity predictor"}));
        }
        let (mut bytes, mut err, mut n, mut energy) = (0u64, 0u64, 0u64, 0u64);
        for (b, (info, units)) in bases.iter().enumerate() {
            let shape = units[0].shape;
            if units.iter().any(|u| u.shape != shape) {
                return Ok(json!({"q": q, "why": "a base's units differ in shape"}));
            }
            // The target is the source at the units' final shape; an identity layer
            // from stage 0 reaches it only where stage 0 has that shape.
            if units.iter().zip(&src_of[b]).any(|(u, s)| u.samples.byte_len() as usize != s.len()) {
                return Ok(json!({"q": q, "why": "the units change shape between stage 0 and the source; an identity layer from stage 0 cannot reach it"}));
            }
            let belows: Vec<Below> = units.iter().map(|u| below(&u.samples)).collect();
            let targets: Vec<Samples> = src_of[b]
                .iter()
                .map(|s| if pcm { Samples::I16(s.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]])).collect()) } else { Samples::U8(s.to_vec()) })
                .collect();
            let tb: Vec<Below> = targets.iter().map(below).collect();
            let params = RefineParams { q, predictor: Predictor::Identity, rate: (0, 0), table_per_unit: false, max_symbols: 4095 };
            let cap = refine_cap(info, units);
            let kind = wai::staged::unit_kind(cap).expect("a refinement");
            let (payload, recon) = wir1::encode(kind, &belows, shape, &tb, shape, &params).map_err(|e| e.to_string())?;
            let refines = Refines { base: hex(&info.content_hash), prev: Prev::One(hex(&info.content_hash)), layer: 1, span: None };
            let media = wai::codecs::capability_media(cap).map(|m| m.label()).unwrap_or("binary");
            let m = wai::Manifest::new(media, "replicate", wai::ModelRequirement::new(cap), wai::Conditioning::new("wir1")).with_refines(&refines);
            bytes += Wai::new(m, payload).to_emit_bytes().map_err(|e| e.to_string())?.len() as u64;
            for (r, t) in recon.iter().zip(&targets) {
                let s = match r {
                    Next::U8(v) => Samples::U8(v.clone()),
                    Next::I16(v) => Samples::I16(v.clone()),
                };
                let (e, m, x) = sse(&s, &t.canonical_bytes());
                err += e;
                n += m;
                energy += x;
            }
        }
        let mut v = quality_json(err, n, energy, pcm);
        v["q"] = json!(q);
        v["bytes"] = json!(bytes);
        Ok(v)
    };

    let mut stages = Vec::new();
    let mut cumulative = 0u64;
    for k in 0..=n_stages {
        let es: Vec<&Env> = envs.iter().filter(|e| e.stage == k).collect();
        let bytes: u64 = es.iter().map(|e| e.bytes.len() as u64).sum();
        cumulative += bytes;
        let (err, n, energy, measured) = clip[k as usize];
        let mut quality = if measured == 0 { json!({"sse": null, "samples": null, "exact": null, "why": "no unit has the source's shape at this stage"}) } else { quality_json(err, n, energy, pcm) };
        quality["units_measured"] = json!(measured);
        stages.push(json!({
            "stage": k,
            "envelopes": es.len(),
            "bytes": bytes,
            "framing_bytes": es.iter().map(|e| (e.bytes.len() - e.manifest - e.payload) as u64).sum::<u64>(),
            "manifest_bytes": es.iter().map(|e| e.manifest as u64).sum::<u64>(),
            "payload_bytes": es.iter().map(|e| e.payload as u64).sum::<u64>(),
            "cumulative_bytes": cumulative,
            "clip": quality,
            "time_us": quantiles(times[k as usize].clone()),
            "direct": if k == 0 { Value::Null } else { direct(k)? },
            "units": snapshots[k as usize],
        }));
    }
    let report = json!({
        "format": "wai-stage-bench/1",
        "conditions": {
            "runs": runs,
            "profile": if cfg!(debug_assertions) { "debug" } else { "release" },
            "arch": std::env::consts::ARCH,
            "os": std::env::consts::OS,
            "threads": 1,
            "energy": "unmetered",
            "receipts": "none: a set on disk carries no delivery receipts",
            "quality": if pcm { "SNR = 10·log10(Σx² / Σ(x−y)²) on the 16-bit samples" } else { "PSNR = 10·log10(255²·N / SSE), SSE the exact integer sum over N samples" },
        },
        "envelopes": envs.iter().map(|e| json!({
            "file": e.file, "object": hex(&wai::jwp_hash::content_hash(&e.bytes)), "stage": e.stage, "bytes": e.bytes.len(),
            "manifest_bytes": e.manifest, "payload_bytes": e.payload, "q": e.q.map(|x| x.0), "predictor": e.q.map(|x| x.1.token()),
        })).collect::<Vec<_>>(),
        "stages": stages,
    });

    println!("stage  envelopes      bytes  cumulative  clip quality       time µs (median [p25, p75])   direct: bytes  quality");
    for s in report["stages"].as_array().expect("stages") {
        let qk = if pcm { "snr_db" } else { "psnr_db" };
        let q = |v: &Value| match v[qk].as_f64() {
            Some(x) => format!("{x:.2} dB"),
            None if v["exact"] == true => "exact".to_string(),
            None => "-".to_string(),
        };
        let d = &s["direct"];
        let direct = if d.is_null() || d["bytes"].is_null() { "-".to_string() } else { format!("{:>9}  {}", d["bytes"], q(d)) };
        println!("{:>5}  {:>9}  {:>9}  {:>10}  {:>16}   {:>9} [{}, {}]   {}",
                 s["stage"], s["envelopes"], s["bytes"], s["cumulative_bytes"], q(&s["clip"]),
                 s["time_us"]["median"], s["time_us"]["p25"], s["time_us"]["p75"], direct);
    }
    if let Some(out) = opt(args, "--out") {
        let mut text = serde_json::to_string_pretty(&report).map_err(|e| e.to_string())?;
        text.push('\n');
        std::fs::write(out, text).map_err(|e| format!("write {out}: {e}"))?;
        println!("report → {out}");
    }
    Ok(())
}

fn main() -> ExitCode {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.is_empty() {
        eprintln!("usage: wai_stage_bench --source <file> [--set <dir>] [--limit OUT,SYMBOLS,WORKING,MACS] [--runs N≥20] [--out report.json] <envelope>…");
        return ExitCode::from(2);
    }
    match run(&args) {
        Ok(()) => ExitCode::SUCCESS,
        Err(e) => {
            eprintln!("error: {e}");
            ExitCode::from(1)
        }
    }
}
