//! A title's evaluation as a sink's session (`extensions/staged-measurement.md` §8.7,
//! jwp-receipts §2.4.12).
//!
//! A tool that decodes a title as a sink does — every base, then every layer, all on the
//! primary path, no deadline — records the session a [`SessionRecorder`] would: one
//! decoder entry per (capability, pin, object) of the evaluation, in its first-use order; every
//! presentation with its output's digest; every object of the title received whole, as its
//! bytes, so the recorder derives its figures and checks the claim against the envelopes
//! when it seals it. Its output tree is the title's (staged-measurement Appendix C), so a
//! session claim over a fixture commits to the root its golden states.

use crate::claim_receipt::{SessionHost, SessionPath};
use crate::codecs::DeclaredTier;
use crate::container::Wai;
use crate::session_recorder::{DecoderKey, SessionRecorder};

use super::run::Evaluation;
use super::title::Title;
use super::MeasureError;

/// This build's decoders, as a decoder entry names them: `builtin:wai@<version>`.
pub fn builtin_decoder() -> String {
    format!("builtin:wai@{}", env!("CARGO_PKG_VERSION"))
}

fn unhex32(s: &str) -> Option<[u8; 32]> {
    if s.len() != 64 {
        return None;
    }
    let mut out = [0u8; 32];
    for (i, b) in out.iter_mut().enumerate() {
        *b = u8::from_str_radix(s.get(2 * i..2 * i + 2)?, 16).ok()?;
    }
    Some(out)
}

/// The recorder of the session that decoding `t` as a sink does is: `host` the decoder
/// host, `session_nonce` fresh. Gives it with the session's entry object, the title's
/// first base: the first envelope a sink dispatches. The record carries no energy; the
/// caller adds a figure where it measured one.
pub fn recorder_of(t: &Title, e: &Evaluation, host: SessionHost, session_nonce: [u8; 32]) -> Result<(SessionRecorder, [u8; 32]), MeasureError> {
    let invalid = |s: String| MeasureError::Invalid(s);
    let first = e.bases.first().ok_or_else(|| invalid(format!("{}: no base", t.dir)))?;
    // The entry object's manifest `media`, as written.
    let base = t.bases.first().ok_or_else(|| invalid(format!("{}: no base", t.dir)))?;
    let media = Wai::from_bytes(&base.obj.bytes).map_err(|x| invalid(format!("{}: {x}", base.obj.path)))?.media().to_owned();
    let mut r = SessionRecorder::new(session_nonce, &media, host, None);
    let mut handles = Vec::new();
    for (capability, pin, object) in &e.decoders {
        let pin = pin.as_deref().map(|p| unhex32(p).ok_or_else(|| invalid(format!("{}: pin {p} is not a SHA-256", t.dir)))).transpose()?;
        // That pin's declared tier as SPEC §3.1 reads it, from the base envelope the entry
        // decoded, which pinned it: an unregistered token, or none, reads as `none`.
        let tier = pin.map(|_| {
            t.bases
                .iter()
                .zip(&e.bases)
                .filter(|(_, (d, _, _))| d.info.content_hash == *object)
                .find_map(|(b, _)| Wai::from_bytes(&b.obj.bytes).ok()?.prior()?.tier().and_then(DeclaredTier::from_token))
                .unwrap_or(DeclaredTier::None)
                .token()
                .to_owned()
        });
        // The title's unit that is unit 0 of the object's base: the base's own, for a layer
        // too.
        let first_unit = e
            .objects
            .iter()
            .find(|o| o.content_hash == *object)
            .and_then(|o| e.bases.get(o.base))
            .map(|(d, _, _)| d.info.first_unit)
            .ok_or_else(|| invalid(format!("{}: a decoder names an object the title does not hold", t.dir)))?;
        handles.push(r.decoder(DecoderKey {
            capability: capability.clone(),
            rendition: None,
            object: *object,
            first_unit,
            path: SessionPath::Primary,
            decoder: builtin_decoder(),
            variant: None,
            pin,
            tier,
            module_status: None,
        }));
    }
    for p in &e.presentations {
        let h = *handles.get(usize::from(p.p.decoder)).ok_or_else(|| invalid(format!("{}: a presentation names no decoder", t.dir)))?;
        r.presented(p.p.unit, p.p.layers, h, p.p.out_digest, None).map_err(|x| invalid(format!("{}: {x}", t.dir)))?;
    }
    for o in &e.objects {
        let held = t.bases.iter().map(|b| &b.obj).chain(&t.layers).find(|x| x.path == o.path);
        let bytes = held.ok_or_else(|| invalid(format!("{}: {} is not one of the title's objects", t.dir, o.path)))?;
        r.received_envelope(&bytes.bytes, SessionPath::Primary);
    }
    Ok((r, first.0.info.content_hash))
}

/// A decode figure measured under energy-measurement §3, as a session seals it: the figure
/// of one completion with its class, the bracket's values, and the power reader.
#[derive(Clone, Debug)]
pub struct MeasuredFigure {
    pub joules_micro: u64,
    pub class: crate::energy_class::EnergyProvenance,
    pub bracket: Option<crate::energy_report::BracketValues>,
    pub power_reader: String,
}

/// The host a session states where its figure was bracketed: the decode workers the
/// bracket's burst ran on (its `in_flight`), written as §2.4.12 writes them. A session
/// whose figure has no bracket keeps `host`. A bracket [`seal_session`] publishes ran on a
/// power of two of workers, so this is its `in_flight` itself.
pub fn host_for(host: SessionHost, figure: Option<&MeasuredFigure>) -> SessionHost {
    match figure.and_then(|f| f.bracket) {
        Some(b) => SessionHost { threads: SessionHost::rounded_threads(b.in_flight), ..host },
        None => host,
    }
}

/// Seal the session a sink decoding `t` would seal, with `figure` as its `energy` where
/// one was measured and published (staged-measurement §8.7), and, beside a figure, the
/// `measurement` claim on it — its bracket and build, signed with the same key and
/// chained after the session (energy-measurement §6). The session's output tree is
/// checked to be the evaluation's before anything is returned, and, where the figure was
/// bracketed, the bracket to have run on a power of two of workers and the host's
/// `threads` to be that number ([`host_for`]): the session and its paired `measurement`
/// claim state one machine, and the claim's `in_flight` names no more of it than the
/// session's rounded `threads` (jwp-receipts §6.7). A meter bursts on
/// [`crate::energy_burst::power_of_two_workers`] for that.
#[allow(clippy::too_many_arguments)]
pub fn seal_session(
    t: &Title,
    e: &Evaluation,
    host: SessionHost,
    session_nonce: [u8; 32],
    figure: Option<&MeasuredFigure>,
    key: &ed25519_dalek::SigningKey,
    signer_id: &str,
    issued_at: &str,
) -> Result<(crate::claim_receipt::ClaimReceipt, Option<crate::claim_receipt::ClaimReceipt>), String> {
    use crate::claim_receipt::{ClaimBody, SessionEnergy};
    use crate::energy_figure::EnergyFigures;
    if let Some(b) = figure.and_then(|f| f.bracket) {
        if b.in_flight != SessionHost::rounded_threads(b.in_flight) {
            return Err(format!(
                "the figure's bracket ran on {} workers, which its report would publish beside a session that rounds them to {}; burst on a power of two",
                b.in_flight,
                SessionHost::rounded_threads(b.in_flight)
            ));
        }
        if host.threads != b.in_flight {
            return Err(format!("the session's host states {} decode workers, its figure's bracket ran on {}", host.threads, b.in_flight));
        }
    }
    let (mut r, entry) = recorder_of(t, e, host, session_nonce).map_err(|x| x.to_string())?;
    r.energy(figure.map(|f| SessionEnergy::new(f.joules_micro, &f.class)));
    let s = r.finish(entry, key, signer_id, 0, None, issued_at).map_err(|x| format!("the session is refused ({x})"))?;
    let ClaimBody::ReconstructSession(b) = &s.claim.body else { unreachable!("a recorder seals a session") };
    if b.output_root != e.output.map(|c| c.root) || Some(b.leaf_count) != e.output.map(|c| c.leaf_count) {
        return Err("the session's output tree is not the title's".into());
    }
    let report = match figure {
        None => None,
        Some(f) => {
            let view = s.claim.figure_view().ok_or("the session does not verify")?;
            let report = crate::energy_figure::meter_report::emit(
                &view,
                None,
                f.bracket,
                crate::energy_report::BuildRecord::of_this_build(),
                &f.power_reader,
                key,
                signer_id,
                Some(&s.claim),
                issued_at,
            )?;
            check_host_agrees(&s.claim, &report)?;
            Some(report)
        }
    };
    Ok((s.claim, report))
}

/// Whether a session and the `measurement` claim on its figure state one machine, each
/// rounded as §2.4.12 writes a worker count: where the report carries a bracket, its
/// `in_flight` is the session's host `threads` — the same power of two.
pub fn check_host_agrees(session: &crate::claim_receipt::ClaimReceipt, report: &crate::claim_receipt::ClaimReceipt) -> Result<(), String> {
    use crate::claim_receipt::ClaimBody;
    let (ClaimBody::ReconstructSession(s), ClaimBody::Measurement(m)) = (&session.body, &report.body) else {
        return Err("not a session and a measurement claim".into());
    };
    match m.bracket {
        Some(b) if s.host.threads != b.in_flight => Err(format!(
            "the session states {} decode workers and its report's bracket {}",
            s.host.threads, b.in_flight
        )),
        _ => Ok(()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::claim_receipt::ClaimBody;
    use crate::energy_report::BuildRecord;
    use ed25519_dalek::SigningKey;

    /// Every fixture of the corpus, decoded as a sink and sealed as a session, commits to
    /// the output tree its evaluation gives — the root and leaf count its golden states —
    /// and checks against what it recorded; one completion decodes as many stages as the
    /// session has leaves.
    #[test]
    fn a_fixtures_session_commits_to_its_evaluation_tree() {
        let wai = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("..");
        let corpus = wai.join("staged-measure-conformance/fixtures");
        let mut names: Vec<String> = std::fs::read_dir(&corpus)
            .unwrap()
            .filter_map(|d| d.ok())
            .filter(|d| d.path().join("title.json").is_file())
            .map(|d| d.file_name().to_string_lossy().into_owned())
            .collect();
        names.sort();
        assert!(!names.is_empty());
        let key = SigningKey::from_bytes(&[0x44; 32]);
        let host = SessionHost { profile: "native-ffi".into(), target: None, build: BuildRecord::of_this_build(), threads: 1 };
        for name in names {
            let t = Title::load(&wai, &format!("staged-measure-conformance/fixtures/{name}")).unwrap();
            let e = super::super::run::evaluate(&t).unwrap();
            let (r, entry) = recorder_of(&t, &e, host.clone(), [3; 32]).unwrap();
            assert_eq!(entry, e.objects[0].content_hash, "{name}: the entry object is the first base");
            let s = r.finish(entry, &key, "did:example:sink", 0, None, "2026-10-06T12:00:00Z").unwrap();
            let ClaimBody::ReconstructSession(b) = &s.claim.body else { panic!("a session") };
            let c = e.output.unwrap();
            assert_eq!((b.output_root, b.leaf_count), (Some(c.root), c.leaf_count), "{name}: the session's tree is the title's");
            assert_eq!(b.decoders.len(), e.decoders.len(), "{name}");
            assert_eq!(b.counted_primary, e.objects.iter().map(|o| o.counted.total).sum::<u64>(), "{name}");
            for (d, (cap, pin, object)) in b.decoders.iter().zip(&e.decoders) {
                assert_eq!((&d.capability, &d.object), (cap, object));
                assert_eq!(d.pin.is_some(), pin.is_some());
                assert_eq!(d.tier.is_some(), pin.is_some());
            }
            s.claim.verify().unwrap();
            s.claim.check_session_outputs(&s.presentations, &|_| None).unwrap();
            s.claim.check_session_objects(&s.objects).unwrap();
            assert_eq!(s.envelopes.len(), e.objects.len(), "{name}: the recorder holds every object it received");
            let held: Vec<&[u8]> = s.envelopes.iter().map(Vec::as_slice).collect();
            s.claim.check_session_envelopes(&held, Some(&s.objects), Some(&s.presentations)).unwrap();
            assert!(s.objects.iter().all(|o| o.presented), "{name}: every object of a completion is presented");
            assert_eq!(super::super::run::completion(&t).unwrap(), c.leaf_count, "{name}: one completion is every stage");
        }
    }

    /// SPEC §3.1, *Pin by default*: a base its encoder left unpinned is still decoded against
    /// the set the sink holds, and the session's decoder entry names that set by the digest a
    /// pin over the bytes read names (jwp-receipts §2.4.12), with the tier `none`. The
    /// fixture's own pin, a listing of the same files, is that digest; the pinned base
    /// names it as before, with the tier its pin declares.
    #[test]
    fn an_unpinned_decode_names_the_set_it_read() {
        let wai = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("..");
        let key = SigningKey::from_bytes(&[0x44; 32]);
        let host = SessionHost { profile: "native-ffi".into(), target: None, build: BuildRecord::of_this_build(), threads: 1 };
        let hex = |b: &[u8; 32]| b.iter().map(|x| format!("{x:02x}")).collect::<String>();
        let entry_of = |t: &Title| {
            let e = super::super::run::evaluate(t).unwrap();
            let (r, entry) = recorder_of(t, &e, host.clone(), [3; 32]).unwrap();
            let s = r.finish(entry, &key, "did:example:sink", 0, None, "2026-10-06T12:00:00Z").unwrap();
            s.claim.verify().unwrap();
            let ClaimBody::ReconstructSession(b) = s.claim.body else { panic!("a session") };
            let d = b.decoders[0].clone();
            (d.pin.as_ref().map(hex), d.tier)
        };
        let mut t = Title::load(&wai, "staged-measure-conformance/fixtures/still").unwrap();
        // The base alone: its layers name the pinned base by its content hash, and its source
        // is the size of their final stage.
        t.layers.clear();
        t.references.clear();
        t.source = None;
        t.expect = None;
        let pinned = Wai::from_bytes(&t.bases[0].obj.bytes).unwrap();
        let pin = pinned.prior().expect("the fixture's base is pinned").sha256.clone();
        assert_eq!(entry_of(&t), (Some(pin.clone()), Some("decode-equivalence".to_string())));

        let mut unpinned = Wai::from_bytes(&t.bases[0].obj.bytes).unwrap();
        unpinned.manifest.conditioning.prior = None;
        unpinned.manifest.intent = "improve".into();
        assert!(unpinned.check_unpinned().is_some(), "the planted base is reported prior_unpinned");
        t.bases[0].obj.bytes = unpinned.to_bytes().unwrap();
        t.bases[0].case = None;
        assert_eq!(entry_of(&t), (Some(pin), Some("none".to_string())), "the set it read, by the digest its pin would have named");
    }

    /// A metered session, as `wai_stage_measure bench --sign` seals one: the figure a live
    /// from-rest bracket gives under `OnChipCounter`, in the session's `energy`, and the
    /// `measurement` claim on it by the same key, chained after the session and checked
    /// against its figure. An unmetered session gets no report.
    #[cfg(not(target_arch = "wasm32"))]
    #[test]
    fn a_metered_session_is_sealed_with_its_report() {
        use crate::claim_receipt::{verify_claim_json, ClaimBody};
        use crate::energy_burst::Bracket;
        use crate::energy_class::{CompletionCount, CounterFiltering, WindowOpening};
        let wai = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("..");
        let t = Title::load(&wai, "staged-measure-conformance/fixtures/still").unwrap();
        let e = super::super::run::evaluate(&t).unwrap();
        let key = SigningKey::from_bytes(&[0x44; 32]);
        let host = SessionHost { profile: "native-ffi".into(), target: None, build: BuildRecord::of_this_build(), threads: 1 };
        // A live, quiet idle bracket and a 10 W marginal over 3 s, 60 completions from
        // rest, on `n` workers.
        let rep = |a: f64, b: f64, n: usize| (0..n).map(|i| if i % 2 == 0 { a } else { b }).collect::<Vec<f64>>();
        let figure_on = |n: u64| {
            let b = Bracket {
                idle_w: rep(19.9, 20.1, 10),
                busy_w: rep(29.9, 30.1, 14),
                window_s: 3.0,
                idle_span_s: 2.0,
                count: CompletionCount { in_flight: n, completions: 60, opening: WindowOpening::FromRest },
                source: crate::energy_burst::PowerSource::Counter,
            };
            let (uj, class) = b.on_chip_figure(CounterFiltering::Unknown).expect("a figure");
            MeasuredFigure { joules_micro: uj, class, bracket: b.values(), power_reader: "a counter reader".into() }
        };
        // A bracket on ten workers would publish ten beside a session that writes eight:
        // refused, whatever the host states.
        let ten = figure_on(10);
        for h in [host.clone(), host_for(host.clone(), Some(&ten))] {
            let refused = seal_session(&t, &e, h, [5; 32], Some(&ten), &key, "did:example:sink", "2026-10-06T12:00:00Z");
            let why = refused.expect_err("a bracket on a worker count the session rounds is refused");
            assert!(why.contains("burst on a power of two"), "refused before anything is sealed: {why}");
        }
        // On eight, a host stating one is refused, and the host the figure gives states
        // eight.
        let f = figure_on(8);
        let uj = f.joules_micro;
        let refused = seal_session(&t, &e, host.clone(), [5; 32], Some(&f), &key, "did:example:sink", "2026-10-06T12:00:00Z");
        assert!(refused.is_err(), "a host that misstates the bracket's workers is refused");
        let metered_host = host_for(host.clone(), Some(&f));
        assert_eq!(metered_host.threads, 8);
        let (session, report) = seal_session(&t, &e, metered_host.clone(), [5; 32], Some(&f), &key, "did:example:sink", "2026-10-06T12:00:00Z").unwrap();
        let ClaimBody::ReconstructSession(body) = &session.body else { panic!("a session") };
        assert_eq!(body.energy.as_ref().map(|x| (x.joules_micro, x.energy_provenance.as_str())), Some((uj, "OnChipCounter")));
        let report = report.expect("a report beside a figure");
        let report = verify_claim_json(&report.to_json()).unwrap();
        assert_eq!((report.subject, report.seq, report.parent_receipt_hash), (session.link_hash(), 1, Some(session.link_hash())));
        report.check_figure_against(crate::energy_figure::EnergyFigures::figure_view(&session).as_ref()).unwrap();
        report.check_parent(Some(&session)).unwrap();
        check_host_agrees(&session, &report).unwrap();
        let ClaimBody::ReconstructSession(sb) = &session.body else { panic!("a session") };
        let ClaimBody::Measurement(mb) = &report.body else { panic!("a report") };
        assert_eq!(sb.host.threads, mb.bracket.unwrap().in_flight, "the report names the session's rounded workers");
        assert!(crate::energy_burst::power_of_two_workers().is_power_of_two());
        // A session that states eight workers beside a report that publishes ten — sealed
        // by hand, as seal_session never would — does not agree: the report names more of
        // the machine than the session.
        let (mut r, entry) = recorder_of(&t, &e, metered_host.clone(), [6; 32]).unwrap();
        r.energy(Some(crate::claim_receipt::SessionEnergy::new(ten.joules_micro, &ten.class)));
        let eight = r.finish(entry, &key, "did:example:sink", 0, None, "2026-10-06T12:00:00Z").unwrap().claim;
        let view = crate::energy_figure::EnergyFigures::figure_view(&eight).unwrap();
        let ten_report = crate::energy_figure::meter_report::emit(
            &view,
            None,
            ten.bracket,
            crate::energy_report::BuildRecord::of_this_build(),
            &ten.power_reader,
            &key,
            "did:example:sink",
            Some(&eight),
            "2026-10-06T12:00:00Z",
        )
        .unwrap();
        assert!(check_host_agrees(&eight, &ten_report).is_err(), "eight workers beside a bracket of ten");
        let (unmetered, none) = seal_session(&t, &e, host, [5; 32], None, &key, "did:example:sink", "2026-10-06T12:00:00Z").unwrap();
        assert!(none.is_none());
        assert!(check_host_agrees(&unmetered, &report).is_err(), "a session of one worker beside a bracket of eight");
        let session = unmetered;
        let ClaimBody::ReconstructSession(body) = &session.body else { panic!("a session") };
        assert_eq!(body.energy, None, "an unmetered session writes null, never 0");
    }
}
