//! A measured title (`extensions/staged-measurement.md` Appendix D): a staged set and
//! what a measurement needs beside it.
//!
//! The stage structure is the staged set's own: each base is a staged-conformance base
//! (its envelope, the parameter set it decodes against, and its stage-0 output), and
//! each layer's `refines` names its base, its producers, its layer and its span
//! (`extensions/staged-delivery.md` §3). `title.json` adds only what that model does
//! not carry: the units' duration and capture times (for the delivery simulator), the
//! source the units approximate, the single-layer references of the layering overhead,
//! and the reconstruction digests the fixture tool wrote from its own encoder.
//!
//! A base may also be named by its envelope alone (with its parameter set): such a base
//! has no stage-0 output or case facts to check its decode against, and the corpus uses
//! the form only for titles a tool must refuse. A title may lower the limits it is
//! measured under (`limits`), never raise them: each figure is the lower of the title's
//! and the tool's (`staged::TOOL_DEFAULT_LIMIT` and `staged::TOOL_DEFAULT_CHAIN_BUDGET`,
//! or the limit the tool is given). A title also may state the code its evaluation must
//! be refused with (`refused`).
//!
//! Every path is relative to the `wai/` directory, so a title reads the staged corpus's
//! bases where they are. [`Title::from_files`] reads from a map of path → bytes and
//! needs no file system (a browser fetches the files); it reports the first file it
//! needs and does not hold ([`MeasureError::Missing`]), so a caller can fetch it and
//! call again. [`Title::load`] does that from a directory.

use std::collections::BTreeMap;

use serde_json::Value;

use super::MeasureError;

/// The `format` of a `title.json`.
pub const FORMAT: &str = "wai-measure-title/1";

/// An object of the title: its path below `wai/` and its bytes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Obj {
    pub path: String,
    pub bytes: Vec<u8>,
}

/// A file of a base's parameter set: the member name its decode reads it as, its path,
/// its bytes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SetFile {
    pub member: String,
    pub path: String,
    pub bytes: Vec<u8>,
}

/// What a staged-conformance case states of how its base decodes: the capability it
/// was decoded through, whether that was the fallback, its tier, and whether its pin
/// resolved. A measurement's decode must state the same.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CaseFacts {
    pub decoded_via: String,
    pub via_fallback: bool,
    pub tier: String,
    pub pin_resolved: bool,
}

/// A base of the title, as its staged-conformance case states it, or by its envelope
/// alone.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TitleBase {
    pub obj: Obj,
    /// The global index of the base's unit 0.
    pub first_unit: u64,
    /// When the base's content is captured, for a live delivery plan; `None` for a
    /// title delivered on demand.
    pub capture_us: Option<u64>,
    pub set: Vec<SetFile>,
    /// The case's stage-0 output: the units' canonical buffers in unit order; `None`
    /// for a base named by its envelope alone.
    pub stage0: Option<Vec<u8>>,
    /// The case's statement of the decode; `None` for a base named by its envelope.
    pub case: Option<CaseFacts>,
}

/// The limits a title is measured under: the cost limit its bases are decoded under
/// (and whose output, symbols and working bytes its layers are held to,
/// staged-delivery Appendix A.5), and its chains' budget (§4.5).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Limits {
    pub cost: crate::int_payload::DeclaredCost,
    pub chain_budget: u64,
}

impl Default for Limits {
    /// The tools' stated limits.
    fn default() -> Self {
        Limits { cost: crate::staged::TOOL_DEFAULT_LIMIT, chain_budget: crate::staged::TOOL_DEFAULT_CHAIN_BUDGET }
    }
}

/// What a title states of its limits: each one it lowers, `None` where it states none.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct StatedLimits {
    pub cost: Option<crate::int_payload::DeclaredCost>,
    pub chain_budget: Option<u64>,
}

impl Limits {
    /// The limits a title is measured under by a tool whose own limits are `tool`: each
    /// figure the lower of the title's and the tool's. A title lowers a limit and never
    /// raises one; a limit past the tool's takes the tool's own (`--limit`), because
    /// whoever writes a title is not whoever runs the tool.
    pub fn within(stated: &StatedLimits, tool: &Limits) -> Limits {
        let t = &tool.cost;
        let cost = stated.cost.map_or(*t, |s| {
            crate::int_payload::DeclaredCost::new(
                s.output_bytes.min(t.output_bytes),
                s.symbols.min(t.symbols),
                s.working_bytes.min(t.working_bytes),
                s.macs.min(t.macs),
            )
        });
        Limits { cost, chain_budget: stated.chain_budget.map_or(tool.chain_budget, |b| b.min(tool.chain_budget)) }
    }
}

/// A single-layer reference (§6.2): the title's bases and the layers given here.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Reference {
    pub label: String,
    pub layers: Vec<Obj>,
}

/// A measured title.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Title {
    pub name: String,
    /// The fixture directory, below `wai/`.
    pub dir: String,
    pub about: String,
    /// The duration of one unit.
    pub unit_us: u64,
    pub bases: Vec<TitleBase>,
    /// The layers, in the order the title lists them.
    pub layers: Vec<Obj>,
    /// The source the units approximate: each unit's canonical buffer at its final
    /// shape, in global unit order.
    pub source: Option<Obj>,
    pub references: Vec<Reference>,
    /// The fixture tool's own reconstruction digests (`expect.json`).
    pub expect: Option<Value>,
    /// What the title states of its limits.
    pub stated: StatedLimits,
    /// The limits it is measured under: [`Limits::within`] the tools' stated limits,
    /// until [`Title::under`] gives a tool's own.
    pub limits: Limits,
    /// The code the title's evaluation must be refused with, for a title that exists
    /// to be refused.
    pub refused: Option<String>,
}

fn invalid(s: impl Into<String>) -> MeasureError {
    MeasureError::Invalid(s.into())
}

fn get<'a>(files: &'a BTreeMap<String, Vec<u8>>, path: &str) -> Result<&'a [u8], MeasureError> {
    files.get(path).map(Vec::as_slice).ok_or_else(|| MeasureError::Missing(path.to_string()))
}

fn json(files: &BTreeMap<String, Vec<u8>>, path: &str) -> Result<Value, MeasureError> {
    serde_json::from_slice(get(files, path)?).map_err(|e| invalid(format!("{path}: {e}")))
}

fn str_of<'a>(v: &'a Value, k: &str, at: &str) -> Result<&'a str, MeasureError> {
    v[k].as_str().ok_or_else(|| invalid(format!("{at}: `{k}` is not a string")))
}

/// A file name inside a directory: one path segment, no `..`.
fn name<'a>(v: &'a Value, at: &str) -> Result<&'a str, MeasureError> {
    let s = v.as_str().ok_or_else(|| invalid(format!("{at}: a file name is not a string")))?;
    if s.is_empty() || s.contains('/') || s.contains('\\') || s == "." || s == ".." {
        return Err(invalid(format!("{at}: `{s}` is not a file name")));
    }
    Ok(s)
}

/// A path below `wai/`: relative, `/`-separated, no `.` or `..` segment.
fn wai_path<'a>(v: &'a Value, at: &str) -> Result<&'a str, MeasureError> {
    let s = v.as_str().ok_or_else(|| invalid(format!("{at}: a path is not a string")))?;
    if s.is_empty() || s.starts_with('/') || s.contains('\\') || s.split('/').any(|p| p.is_empty() || p == "." || p == "..") {
        return Err(invalid(format!("{at}: `{s}` is not a path below wai/")));
    }
    Ok(s)
}

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

impl Title {
    /// Measure the title under a tool's own limits `tool`: each figure the lower of the
    /// title's and the tool's ([`Limits::within`]).
    pub fn under(&mut self, tool: &Limits) {
        self.limits = Limits::within(&self.stated, tool);
    }

    /// The title in the fixture directory `dir` (below `wai/`), from `files` (path
    /// below `wai/` → bytes).
    pub fn from_files(dir: &str, files: &BTreeMap<String, Vec<u8>>) -> Result<Title, MeasureError> {
        let at = format!("{dir}/title.json");
        let t = json(files, &at)?;
        if t["format"] != FORMAT {
            return Err(invalid(format!("{at}: format is not {FORMAT}")));
        }
        let unit_us = t["unit_us"].as_u64().filter(|&u| u >= 1).ok_or_else(|| invalid(format!("{at}: unit_us is not an integer ≥ 1")))?;
        let local = |v: &Value, what: &str| -> Result<Obj, MeasureError> {
            let path = format!("{dir}/{}", name(v, &format!("{at}: {what}"))?);
            Ok(Obj { bytes: get(files, &path)?.to_vec(), path })
        };
        let mut bases = Vec::new();
        for (i, b) in t["bases"].as_array().ok_or_else(|| invalid(format!("{at}: no bases")))?.iter().enumerate() {
            let first_unit = b["first_unit"].as_u64().ok_or_else(|| invalid(format!("{at}: base {i}: first_unit")))?;
            let capture_us = match &b["capture_us"] {
                Value::Null => None,
                v => Some(v.as_u64().ok_or_else(|| invalid(format!("{at}: base {i}: capture_us")))?),
            };
            let read_set = |set: &Value, at: &str| -> Result<Vec<SetFile>, MeasureError> {
                let mut out = Vec::new();
                for (member, p) in set.as_object().ok_or_else(|| invalid(format!("{at}: no decode set")))? {
                    let sp = wai_path(p, at)?;
                    out.push(SetFile { member: member.clone(), path: sp.to_string(), bytes: get(files, sp)?.to_vec() });
                }
                Ok(out)
            };
            if !b["envelope"].is_null() {
                // A base by its envelope alone.
                let path = wai_path(&b["envelope"], &format!("{at}: base {i}"))?.to_string();
                let set = read_set(&b["set"], &format!("{at}: base {i}"))?;
                bases.push(TitleBase { obj: Obj { bytes: get(files, &path)?.to_vec(), path }, first_unit, capture_us, set, stage0: None, case: None });
                continue;
            }
            let case = wai_path(&b["case"], &format!("{at}: base {i}"))?;
            let cpath = format!("{case}/case.json");
            let c = json(files, &cpath)?;
            let k = b["base"].as_u64().ok_or_else(|| invalid(format!("{at}: base {i}: `base` is not an index")))? as usize;
            let cb = c["bases"].get(k).ok_or_else(|| invalid(format!("{cpath}: no base {k}")))?;
            let path = format!("{case}/{}", name(&cb["file"], &cpath)?);
            let bytes = get(files, &path)?.to_vec();
            if hex(&crate::jwp_hash::content_hash(&bytes)) != str_of(cb, "object", &cpath)? {
                return Err(invalid(format!("{path}: not the content hash {cpath} states")));
            }
            if cb["decode"]["kind"] != "int" {
                return Err(invalid(format!("{cpath}: base {k} is not decoded by an integer decoder")));
            }
            let set = read_set(&cb["decode"]["set"], &cpath)?;
            let s0 = format!("{case}/{}", name(&cb["stage0"], &cpath)?);
            let stage0 = get(files, &s0)?.to_vec();
            if crate::pin::sha256_hex(&stage0) != str_of(cb, "stage0_sha256", &cpath)? {
                return Err(invalid(format!("{s0}: not the SHA-256 {cpath} states")));
            }
            let facts = CaseFacts {
                decoded_via: str_of(cb, "decoded_via", &cpath)?.to_string(),
                via_fallback: cb["via_fallback"].as_bool().ok_or_else(|| invalid(format!("{cpath}: via_fallback")))?,
                tier: str_of(cb, "tier", &cpath)?.to_string(),
                pin_resolved: cb["pin_resolved"].as_bool().ok_or_else(|| invalid(format!("{cpath}: pin_resolved")))?,
            };
            bases.push(TitleBase { obj: Obj { path, bytes }, first_unit, capture_us, set, stage0: Some(stage0), case: Some(facts) });
        }
        if bases.is_empty() {
            return Err(invalid(format!("{at}: no bases")));
        }
        let layers = t["layers"]
            .as_array()
            .ok_or_else(|| invalid(format!("{at}: `layers` is not an array")))?
            .iter()
            .map(|v| local(v, "layers"))
            .collect::<Result<Vec<_>, _>>()?;
        let source = match &t["source"] {
            Value::Null => None,
            s => {
                let o = local(&s["file"], "source")?;
                if crate::pin::sha256_hex(&o.bytes) != str_of(s, "sha256", &at)? {
                    return Err(invalid(format!("{}: not the SHA-256 {at} states", o.path)));
                }
                Some(o)
            }
        };
        let mut references = Vec::new();
        for r in t["references"].as_array().map(Vec::as_slice).unwrap_or_default() {
            let label = str_of(r, "label", &at)?.to_string();
            let layers = r["layers"]
                .as_array()
                .ok_or_else(|| invalid(format!("{at}: reference `{label}` has no layers")))?
                .iter()
                .map(|v| local(v, "references"))
                .collect::<Result<Vec<_>, _>>()?;
            references.push(Reference { label, layers });
        }
        let expect = match &t["expect"] {
            Value::Null => None,
            v => Some(json(files, &format!("{dir}/{}", name(v, &at)?))?),
        };
        let stated = match &t["limits"] {
            Value::Null => StatedLimits::default(),
            l => {
                let cost = match &l["cost"] {
                    Value::Null => None,
                    c => {
                        let f: Vec<u64> = c.as_array().filter(|a| a.len() == 4).and_then(|a| a.iter().map(Value::as_u64).collect()).ok_or_else(|| invalid(format!("{at}: limits.cost is not four integers")))?;
                        Some(crate::int_payload::DeclaredCost::new(f[0], f[1], f[2], f[3] as u128))
                    }
                };
                let chain_budget = match &l["chain_budget"] {
                    Value::Null => None,
                    v => Some(v.as_u64().ok_or_else(|| invalid(format!("{at}: limits.chain_budget")))?),
                };
                StatedLimits { cost, chain_budget }
            }
        };
        let limits = Limits::within(&stated, &Limits::default());
        let refused = match &t["refused"] {
            Value::Null => None,
            v => Some(v.as_str().ok_or_else(|| invalid(format!("{at}: refused is not a string")))?.to_string()),
        };
        Ok(Title {
            stated,
            limits,
            refused,
            name: str_of(&t, "title", &at)?.to_string(),
            dir: dir.to_string(),
            about: t["about"].as_str().unwrap_or_default().to_string(),
            unit_us,
            bases,
            layers,
            source,
            references,
            expect,
        })
    }

    /// The title in `dir` (below `wai_root`), reading each file it needs from disk.
    #[cfg(not(target_arch = "wasm32"))]
    pub fn load(wai_root: &std::path::Path, dir: &str) -> Result<Title, MeasureError> {
        let mut files: BTreeMap<String, Vec<u8>> = BTreeMap::new();
        loop {
            match Title::from_files(dir, &files) {
                Err(MeasureError::Missing(p)) => {
                    let b = std::fs::read(wai_root.join(&p)).map_err(|e| invalid(format!("{p}: {e}")))?;
                    files.insert(p, b);
                }
                other => return other,
            }
        }
    }

    /// The title's content identity: BLAKE3 over the measured content, the paths and
    /// bytes of its objects (bases, layers and the references' layers), its bases'
    /// parameter-set files and its source, in path order, each path and each file
    /// length-prefixed. Labels, verification aids (`stage0`, `expect.json`) and
    /// `title.json` itself are not part of it.
    pub fn content_id(&self) -> [u8; 32] {
        let mut all: BTreeMap<&str, &[u8]> = BTreeMap::new();
        for b in &self.bases {
            all.insert(&b.obj.path, &b.obj.bytes);
            for s in &b.set {
                all.insert(&s.path, &s.bytes);
            }
        }
        for o in self.layers.iter().chain(self.references.iter().flat_map(|r| &r.layers)) {
            all.insert(&o.path, &o.bytes);
        }
        if let Some(s) = &self.source {
            all.insert(&s.path, &s.bytes);
        }
        let mut h = blake3::Hasher::new();
        h.update(b"wai:measure-title\x01");
        for (p, b) in all {
            h.update(&(p.len() as u64).to_be_bytes());
            h.update(p.as_bytes());
            h.update(&(b.len() as u64).to_be_bytes());
            h.update(b);
        }
        *h.finalize().as_bytes()
    }
}
