//! Every stage of a staged set, decoded through the reference sink.
//!
//! A base is decoded by `staged::decode_base_pinned`, the one base decode the sink, the
//! `wai stage` tool and the browser binding share: SPEC §4's dispatch through the
//! integer decoders, the pin verified against the parameter set (SPEC §3.1), and the
//! title's cost limit applied to the payload's declared cost before any symbol is
//! decoded or any output allocated. Its layers are applied by `staged::StageSink`
//! (`extensions/staged-delivery.md` §4) under the title's chain budget, with the same
//! limit on each layer's declared cost (Appendix A.5), one layer index at a time; each
//! unit's output at each stage is handed to the caller as it is reached and not kept,
//! so a replay holds one output per unit. This is a composition of library functions
//! for measurement, not a sink's dispatch path: it requires every layer to apply to
//! every unit it covers, as a set delivered whole and in order does, and refuses
//! anything else.

use std::collections::BTreeMap;

use crate::container::Wai;
use crate::staged::{tier, BaseInfo, ChainConfig, HoldBound, ReceiptEvidence, StageRefusal, StageSink, UnitBuf};

use super::title::{Limits, TitleBase};
use super::MeasureError;

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

/// A base decoded: its identity as the sink holds it, its units, and the parameter set its
/// decode read, named as a pin names it: the SHA-256 of the pin it verified, or, unpinned,
/// the digest a pin over the bytes it read would name (`staged::BaseDecode::set`); `None`
/// for a capability that reads no parameter set.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DecodedBase {
    pub info: BaseInfo,
    pub units: Vec<UnitBuf>,
    pub pin: Option<String>,
}

/// Decode a base under `limits`, and require what its case states: its stage-0 output,
/// the capability it was decoded through, whether that was the fallback, its tier and
/// whether its pin resolved. A refusal carries the decode's code.
pub fn decode_base(b: &TitleBase, limits: &Limits) -> Result<DecodedBase, MeasureError> {
    checked(b, decode_pinned(b, limits)?)
}

/// The base's decode alone: `staged::decode_base_pinned` against its parameter set, under
/// the title's cost limit.
fn decode_pinned(b: &TitleBase, limits: &Limits) -> Result<crate::staged::BaseDecode, MeasureError> {
    let set: BTreeMap<String, &[u8]> = b.set.iter().map(|s| (s.member.clone(), s.bytes.as_slice())).collect();
    crate::staged::decode_base_pinned(&b.obj.bytes, None, &set, Some(&limits.cost), None)
        .map_err(|e| MeasureError::Refused { at: b.obj.path.clone(), code: e.code().to_string(), detail: e.to_string() })
}

/// A base's decode checked against what its case states.
fn checked(b: &TitleBase, d: crate::staged::BaseDecode) -> Result<DecodedBase, MeasureError> {
    let at = &b.obj.path;
    let mut info = d.info;
    info.first_unit = b.first_unit;
    if let Some(s0) = &b.stage0 {
        let mine: Vec<u8> = d.units.iter().flat_map(|u| u.samples.canonical_bytes()).collect();
        if &mine != s0 {
            return Err(invalid(format!("{at}: the decode is not the case's stage 0")));
        }
    }
    if let Some(c) = &b.case {
        let got = (info.capability.as_str(), info.via_fallback, tier::token(info.tier), d.pin.is_some());
        let want = (c.decoded_via.as_str(), c.via_fallback, c.tier.as_str(), c.pin_resolved);
        if got != want {
            return Err(invalid(format!("{at}: the decode gives (via, fallback, tier, pin) {got:?}; the case states {want:?}")));
        }
    }
    Ok(DecodedBase { info, units: d.units, pin: d.set })
}

/// The chain configuration a replay runs under: the title's chain budget, its cost
/// limit on each layer, a hold bound no in-order replay reaches, and no live schedule.
pub fn chain_config(limits: &Limits) -> ChainConfig {
    ChainConfig {
        // A replay offers each layer after the layers it needs, so nothing is held; a
        // layer held at all is refused by the replay, whatever the bound.
        bound: HoldBound { max_envelopes: 4096, max_age_ms: u64::MAX },
        cost_limit: Some(crate::staged::layer_limit(&limits.cost)),
        budget: limits.chain_budget,
        // A replay never drops a chain, so it keeps no record of one.
        record_budget: 0,
        live: None,
    }
}

/// One stage a unit reached, as a replay reports it.
#[derive(Debug, Clone, Copy)]
pub struct StageEvent<'a> {
    /// The global unit index: the base's `first_unit` plus the unit's index in it.
    pub unit: u64,
    /// The base's index in the title.
    pub base: usize,
    /// The stage's layer count: 1 for the base alone.
    pub layers: u16,
    /// The capability that produced it: the base's, or the layer's.
    pub capability: &'a str,
    /// The parameter set the base's decode read, as a pin names it ([`DecodedBase`]), for
    /// stage 0.
    pub pin: Option<&'a str>,
    /// The content hash of the envelope that produced it: the base's, or the layer's.
    pub object: [u8; 32],
    pub buf: &'a UnitBuf,
}

/// A unit's output at the last stage it reached.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FinalUnit {
    pub unit: u64,
    pub base: usize,
    /// Its stage's layer count.
    pub layers: u16,
    pub buf: UnitBuf,
}

/// What one layer envelope is: its base, its layer index, its units, its capability.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LayerInfo {
    pub base: usize,
    pub layer: u8,
    /// The global units it covers: `[from, from + count)`.
    pub from: u64,
    pub count: u64,
    pub capability: String,
}

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

/// Read a layer envelope's `refines` against the bases.
pub fn layer_info(bytes: &[u8], bases: &[DecodedBase], at: &str) -> Result<LayerInfo, MeasureError> {
    let w = Wai::from_bytes(bytes).map_err(|e| invalid(format!("{at}: {e}")))?;
    let r = w.refines().ok_or_else(|| invalid(format!("{at}: not a layer (no refines)")))?.map_err(|e| invalid(format!("{at}: {e}")))?;
    let b = bases
        .iter()
        .position(|d| hex(&d.info.content_hash) == r.base)
        .ok_or_else(|| invalid(format!("{at}: its base is not one of the title's")))?;
    let n = bases[b].units.len() as u32;
    let (from, count) = r.units(n);
    Ok(LayerInfo { base: b, layer: r.layer, from: bases[b].info.first_unit + from as u64, count: count as u64, capability: w.capability().to_string() })
}

/// Apply `layers` (envelope bytes, with a name for messages) to the decoded bases under
/// `limits`, layer index by layer index and, within one, in the order given. `on_stage`
/// sees every stage every unit reaches, stage 0 first, as it is reached; the result is
/// each unit's output at its last stage, in unit order. A base past the chain budget is
/// refused (`chain-over-limit`), and so is a layer that does not apply to every unit it
/// covers (its reason's code).
pub fn replay(
    bases: &[DecodedBase],
    layers: &[(&str, &[u8])],
    limits: &Limits,
    on_stage: &mut dyn FnMut(StageEvent<'_>) -> Result<(), MeasureError>,
) -> Result<Vec<FinalUnit>, MeasureError> {
    let mut sink = StageSink::new(chain_config(limits));
    let mut seen = std::collections::BTreeSet::new();
    for (bi, b) in bases.iter().enumerate() {
        for (u, buf) in b.units.iter().enumerate() {
            let unit = b.info.first_unit + u as u64;
            if !seen.insert(unit) {
                return Err(invalid(format!("two bases give the global unit {unit}")));
            }
            on_stage(StageEvent { unit, base: bi, layers: 1, capability: &b.info.capability, pin: b.pin.as_deref(), object: b.info.content_hash, buf })?;
        }
        let reports = sink.add_base(b.info.clone(), b.units.clone());
        if let Some(r) = reports.iter().find(|r| r.envelope == Some(StageRefusal::ChainOverLimit)) {
            return Err(MeasureError::Refused {
                at: format!("base {bi} ({})", hex(&r.content_hash)),
                code: StageRefusal::ChainOverLimit.code().into(),
                detail: "the base's units pass the chain budget".into(),
            });
        }
        if !reports.is_empty() {
            return Err(invalid("a base released something on decode: two bases with one content hash"));
        }
    }
    let mut infos = Vec::new();
    for (name, bytes) in layers {
        infos.push(layer_info(bytes, bases, name)?);
    }
    let top = infos.iter().map(|l| l.layer).max().unwrap_or(0);
    for k in 1..=top {
        for (i, ((name, bytes), info)) in layers.iter().zip(&infos).enumerate() {
            if info.layer != k {
                continue;
            }
            let reports = sink.offer(bytes, i as u64, ReceiptEvidence::None).map_err(|e| invalid(format!("{name}: {e}")))?;
            let r = reports.first().filter(|_| reports.len() == 1).ok_or_else(|| invalid(format!("{name}: {} reports for one offer", reports.len())))?;
            if let Some(why) = r.envelope {
                return Err(MeasureError::Refused { at: name.to_string(), code: why.code().into(), detail: r.detail.unwrap_or("the layer is refused").into() });
            }
            if let Some(&(u, why)) = r.refused.first() {
                return Err(MeasureError::Refused { at: name.to_string(), code: why.code().into(), detail: format!("unit {u}") });
            }
            if !r.held.is_empty() || r.held_for_base || r.applied.len() as u64 != info.count {
                return Err(invalid(format!("{name}: does not apply to every unit it covers")));
            }
            let base = &bases[info.base];
            let chain = sink.chain(&base.info.content_hash).expect("a base the sink holds");
            for &u in &r.applied {
                let st = chain.unit(u).expect("an applied unit");
                let unit = base.info.first_unit + u as u64;
                on_stage(StageEvent {
                    unit,
                    base: info.base,
                    layers: st.stage as u16 + 1,
                    capability: &info.capability,
                    pin: None,
                    object: r.content_hash,
                    buf: &st.buf,
                })?;
            }
        }
    }
    let mut finals = Vec::new();
    for (bi, b) in bases.iter().enumerate() {
        let chain = sink.chain(&b.info.content_hash).expect("a base the sink holds");
        for (u, st) in chain.units().iter().enumerate() {
            finals.push(FinalUnit { unit: b.info.first_unit + u as u64, base: bi, layers: st.stage as u16 + 1, buf: st.buf.clone() });
        }
    }
    finals.sort_by_key(|f| f.unit);
    Ok(finals)
}

/// One decode step of a title, the case staged-measurement §7.1 times: a base decoded,
/// giving its units at one layer, or one layer applied, giving the units it covers at its
/// stage.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Step {
    /// The title's base at this index.
    Base(usize),
    /// The title's layer at this index.
    Layer(usize),
}

/// The steps of a title in the order [`replay`] takes them: every base, then each layer
/// index in turn and, within one, the layers in title order.
pub fn steps(t: &super::title::Title, bases: &[DecodedBase]) -> Result<Vec<Step>, MeasureError> {
    let infos: Vec<LayerInfo> = t.layers.iter().map(|o| layer_info(&o.bytes, bases, &o.path)).collect::<Result<_, _>>()?;
    let mut out: Vec<Step> = (0..t.bases.len()).map(Step::Base).collect();
    let top = infos.iter().map(|l| l.layer).max().unwrap_or(0);
    for k in 1..=top {
        out.extend(infos.iter().enumerate().filter(|(_, l)| l.layer == k).map(|(i, _)| Step::Layer(i)));
    }
    Ok(out)
}

/// A step made ready to run: every step before it in [`steps`] order done, so running it
/// does the step's own work and nothing else.
pub enum Prepared<'a> {
    Base { base: &'a TitleBase, limits: &'a Limits },
    Layer { sink: StageSink, name: &'a str, bytes: &'a [u8], index: u64, first_unit: u64, base: [u8; 32] },
}

/// Make `step` of `t` ready: for a layer, the bases (already decoded, `bases`) taken by a
/// sink under the title's limits and every earlier layer applied, each to every unit it
/// covers, as [`replay`] applies them.
pub fn prepare<'a>(t: &'a super::title::Title, bases: &[DecodedBase], step: Step) -> Result<Prepared<'a>, MeasureError> {
    let order = steps(t, bases)?;
    match step {
        Step::Base(b) => {
            let base = t.bases.get(b).ok_or_else(|| invalid(format!("no base {b}")))?;
            Ok(Prepared::Base { base, limits: &t.limits })
        }
        Step::Layer(i) => {
            let obj = t.layers.get(i).ok_or_else(|| invalid(format!("no layer {i}")))?;
            let mut sink = StageSink::new(chain_config(&t.limits));
            for b in bases {
                sink.add_base(b.info.clone(), b.units.clone());
            }
            for s in order.iter().take_while(|s| **s != step) {
                if let Step::Layer(j) = s {
                    let o = &t.layers[*j];
                    apply(&mut sink, &o.path, &o.bytes, *j as u64)?;
                }
            }
            let info = layer_info(&obj.bytes, bases, &obj.path)?;
            let b = &bases[info.base];
            Ok(Prepared::Layer { sink, name: &obj.path, bytes: &obj.bytes, index: i as u64, first_unit: b.info.first_unit, base: b.info.content_hash })
        }
    }
}

/// Offer one layer and require it applied to every unit it covers, as [`replay`] does;
/// gives the units it applied to.
fn apply(sink: &mut StageSink, name: &str, bytes: &[u8], index: u64) -> Result<Vec<u32>, MeasureError> {
    let reports = sink.offer(bytes, index, ReceiptEvidence::None).map_err(|e| invalid(format!("{name}: {e}")))?;
    let r = reports.first().filter(|_| reports.len() == 1).ok_or_else(|| invalid(format!("{name}: {} reports for one offer", reports.len())))?;
    if let Some(why) = r.envelope {
        return Err(MeasureError::Refused { at: name.to_string(), code: why.code().into(), detail: r.detail.unwrap_or("the layer is refused").into() });
    }
    if let Some(&(u, why)) = r.refused.first() {
        return Err(MeasureError::Refused { at: name.to_string(), code: why.code().into(), detail: format!("unit {u}") });
    }
    if !r.held.is_empty() || r.held_for_base {
        return Err(invalid(format!("{name}: does not apply to every unit it covers")));
    }
    Ok(r.applied.clone())
}

/// A step that ran, its outputs not yet read: what a bench times is [`Prepared::run`]
/// alone, and [`Ran::outputs`] checks and hashes them after the clock stops.
pub enum Ran<'a> {
    Base { base: &'a TitleBase, decoded: crate::staged::BaseDecode },
    Layer { sink: StageSink, name: &'a str, applied: Vec<u32>, first_unit: u64, base: [u8; 32] },
}

impl<'a> Prepared<'a> {
    /// Run the step's own work and nothing else: `staged::decode_base_pinned` for a base,
    /// the sink's offer of the layer for a layer.
    pub fn run(self) -> Result<Ran<'a>, MeasureError> {
        match self {
            Prepared::Base { base, limits } => Ok(Ran::Base { base, decoded: decode_pinned(base, limits)? }),
            Prepared::Layer { mut sink, name, bytes, index, first_unit, base } => {
                let applied = apply(&mut sink, name, bytes, index)?;
                Ok(Ran::Layer { sink, name, applied, first_unit, base })
            }
        }
    }
}

impl Ran<'_> {
    /// The outputs the step produced, as (unit, layers, SHA-256 of its canonical buffer),
    /// the unit's global index; a base checked first against its case's stage 0 and
    /// decode, as [`decode_base`] checks one.
    pub fn outputs(self) -> Result<Vec<(u64, u16, [u8; 32])>, MeasureError> {
        use crate::output_tree::out_digest;
        match self {
            Ran::Base { base, decoded } => {
                let d = checked(base, decoded)?;
                Ok(d.units.iter().enumerate().map(|(u, buf)| (d.info.first_unit + u as u64, 1, out_digest(&buf.samples.canonical_bytes()))).collect())
            }
            Ran::Layer { sink, name, applied, first_unit, base } => {
                let chain = sink.chain(&base).ok_or_else(|| invalid(format!("{name}: its base is gone")))?;
                applied
                    .iter()
                    .map(|&u| {
                        let st = chain.unit(u).ok_or_else(|| invalid(format!("{name}: unit {u} applied and not held")))?;
                        Ok((first_unit + u as u64, st.stage as u16 + 1, out_digest(&st.buf.samples.canonical_bytes())))
                    })
                    .collect()
            }
        }
    }
}

/// The first output of a run that is not the reference's: a (unit, layers) the reference
/// gives another digest, or none (staged-measurement §7.3: such a run is not a timing).
pub fn differs_from(reference: &BTreeMap<(u64, u16), [u8; 32]>, out: &[(u64, u16, [u8; 32])]) -> Option<(u64, u16)> {
    out.iter().find(|(u, l, d)| reference.get(&(*u, *l)) != Some(d)).map(|(u, l, _)| (*u, *l))
}
