//! Session receipts — extension §8, a profile of the JWP receipt chain.
//!
//! A live or replayed world session receipts as a JWP **group**:
//!
//! - **Objects** are op batches: the concatenated `CausalOp` bytes for
//!   one checkpoint interval, content-hashed with domain
//!   `wai:world-ops\x01` exactly as jwp-receipts §2.1 hashes an
//!   envelope. Batch *k* covers the ops whose `effect_tick` falls in
//!   `(checkpoint[k-1].tick, checkpoint[k].tick]`, so each batch pairs
//!   one-to-one with the `state_hash` it produces.
//! - The group is **Merkle-bound** over those object hashes using JWP's
//!   exact leaf/node domain separators (this is a profile, not a fork).
//! - One **Ed25519 signature** over the world-profile signing payload,
//!   which carries — beyond the JWP base fields — the world hash, the
//!   checkpoint `state_hash` list, and the session's measured joules.
//!
//! Joules are a **carried** quantity, never fabricated: the builder
//! takes a microjoule figure from the caller (the sink-metering path).
//! An unmetered session passes `0` and the receipt says so honestly.
//!
//! The figure's acquisition class is an optional label
//! (`extensions/energy-measurement.md` §2.8): [`SessionReceipt::build`] seals
//! an unlabelled receipt whose signed bytes are exactly the legacy ones, and
//! [`SessionReceipt::with_energy_class`] labels and re-signs it.

use super::op::CausalOp;
use super::replay::Checkpoint;
use super::DOMAIN_OPS_BATCH;
use crate::energy_class::{
    check_label, json_with_label, label_coheres, labelled_domain, write_label, AcquisitionError,
    EnergyProvenance,
};
use ed25519_dalek::{Signer, SigningKey};

/// JWP Merkle domain separators (jwp-receipts §2.2), reused verbatim —
/// this profile builds the *same* tree shape JWP defines.
const MERKLE_LEAF: &[u8] = b"moq-jwp:merkle-leaf\x01";
const MERKLE_NODE: &[u8] = b"moq-jwp:merkle-node\x01";

/// World-profile group-signing domain.
const DOMAIN_RECEIPT: &[u8] = b"wai:world-receipt\x01";

/// A signed, energy-accounted record of one world session.
#[derive(Clone, Debug, PartialEq)]
pub struct SessionReceipt {
    /// Identity of the world being played: BLAKE3 over the scene
    /// (contract + tick-0 graph), domain `wai:world-scene\x01`. Two
    /// sessions of the same world share this hash.
    pub world_hash: [u8; 32],
    /// Monotonic group counter under `signer_pubkey`; `0` for a
    /// standalone session.
    pub group_id: u64,
    /// `(tick, state_hash)` at each checkpoint — the §7 list. Batch k
    /// produced `checkpoints[k]`.
    pub checkpoints: Vec<Checkpoint>,
    /// Per-batch content hashes (the Merkle leaves, in order).
    pub object_hashes: Vec<[u8; 32]>,
    /// Merkle root over `object_hashes` (JWP §2.2).
    pub merkle_root: [u8; 32],
    /// Measured microjoules for the session; `0` is the unmetered marker, never a
    /// figure of zero joules (energy-measurement §2). Joules
    /// are *carried*, never fabricated — a real meter (RAPL natively, a
    /// declared wall-clock estimate in a browser) supplies the figure;
    /// it is NOT derived inside this builder, and a sink that didn't
    /// measure signs `0` honestly.
    pub joules_micro: u64,
    /// How `joules_micro` was acquired — an optional label
    /// (`extensions/energy-measurement.md` §2.8). `None` is an unlabelled
    /// figure, as every receipt sealed before the label existed carries; it is
    /// never read as any class. Inside the signature when present.
    pub energy_provenance: Option<EnergyProvenance>,
    /// Deterministic node-steps simulated (Σ over ticks of node count) —
    /// exact and reproducible on any sink, unlike `joules_micro`. This
    /// is the honest work measure: joules are derived FROM it with a
    /// declared power coefficient, never asserted independently.
    pub work_steps: u64,
    /// BLAKE3 of the previous group receipt's canonical bytes
    /// ([`Self::receipt_hash`]), or `None` for the first in a chain.
    /// Successive sealed sessions of a world form an append-only,
    /// signed timeline (JWP `parent_receipt_hash`, generalized to state).
    pub parent_receipt_hash: Option<[u8; 32]>,
    pub signer_pubkey: [u8; 32],
    pub signer_id: String,
    pub sig: [u8; 64],
}

/// Content hash of one op batch (extension §8 / jwp-receipts §2.1).
pub fn batch_content_hash(batch_bytes: &[u8]) -> [u8; 32] {
    let mut h = blake3::Hasher::new();
    h.update(DOMAIN_OPS_BATCH);
    h.update(batch_bytes);
    *h.finalize().as_bytes()
}

fn merkle_leaf(content_hash: &[u8; 32]) -> [u8; 32] {
    let mut h = blake3::Hasher::new();
    h.update(MERKLE_LEAF);
    h.update(content_hash);
    *h.finalize().as_bytes()
}

fn merkle_node(left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] {
    let mut h = blake3::Hasher::new();
    h.update(MERKLE_NODE);
    h.update(left);
    h.update(right);
    *h.finalize().as_bytes()
}

/// JWP §2.2 Merkle root: duplicate the final leaf on odd levels; the
/// empty-group root is all-zero.
pub fn merkle_root(object_hashes: &[[u8; 32]]) -> [u8; 32] {
    if object_hashes.is_empty() {
        return [0u8; 32];
    }
    let mut level: Vec<[u8; 32]> = object_hashes.iter().map(merkle_leaf).collect();
    while level.len() > 1 {
        if level.len() % 2 == 1 {
            level.push(*level.last().unwrap());
        }
        level = level
            .chunks_exact(2)
            .map(|pair| merkle_node(&pair[0], &pair[1]))
            .collect();
    }
    level[0]
}

/// Split a canonically-ordered op log into per-checkpoint batches.
/// Batch `k` is the ops with `effect_tick` in
/// `(prev_tick, checkpoints[k].tick]` (prev_tick = 0 for k=0).
pub fn batch_ops(ops: &[CausalOp], checkpoints: &[Checkpoint]) -> Vec<Vec<u8>> {
    let mut batches = Vec::with_capacity(checkpoints.len());
    let mut prev = 0u64;
    let mut idx = 0usize;
    for cp in checkpoints {
        let mut bytes = Vec::new();
        while idx < ops.len() && ops[idx].effect_tick <= cp.tick {
            // ops are canonically ordered, so effect_tick is non-decreasing
            debug_assert!(ops[idx].effect_tick > prev || prev == 0);
            bytes.extend_from_slice(&ops[idx].to_bytes());
            idx += 1;
        }
        batches.push(bytes);
        prev = cp.tick;
    }
    batches
}

impl SessionReceipt {
    /// The signed bytes: world-profile signing payload (extends JWP
    /// §2.3 with the world hash, checkpoint list, and joules).
    fn signing_payload(&self) -> Vec<u8> {
        let label = self.energy_provenance.as_ref();
        let mut out = labelled_domain(DOMAIN_RECEIPT, label);
        out.extend_from_slice(&self.world_hash);
        out.extend_from_slice(&self.group_id.to_be_bytes());
        out.extend_from_slice(&self.merkle_root);
        out.extend_from_slice(&(self.object_hashes.len() as u64).to_be_bytes());
        out.extend_from_slice(&self.joules_micro.to_be_bytes());
        write_label(label, &mut out);
        out.extend_from_slice(&self.work_steps.to_be_bytes());
        // parent: 1-byte present-flag then the hash (domain-stable: a
        // first receipt and a chained one can never collide).
        match self.parent_receipt_hash {
            Some(h) => {
                out.push(1);
                out.extend_from_slice(&h);
            }
            None => out.push(0),
        }
        out.extend_from_slice(&(self.checkpoints.len() as u64).to_be_bytes());
        for cp in &self.checkpoints {
            out.extend_from_slice(&cp.tick.to_be_bytes());
            out.extend_from_slice(&cp.hash);
        }
        out.extend_from_slice(&self.signer_pubkey);
        let id = self.signer_id.as_bytes();
        out.extend_from_slice(&(id.len() as u32).to_be_bytes());
        out.extend_from_slice(id);
        out
    }

    /// Build and sign a session receipt. `joules_micro` is carried from
    /// the caller's meter (0 if unmetered); `work_steps` is the exact
    /// deterministic node-step count from
    /// [`crate::world::collect_checkpoints_and_work`]; `parent` links
    /// this receipt to the previous one in the world's timeline.
    #[allow(clippy::too_many_arguments)]
    pub fn build(
        signer: &SigningKey,
        signer_id: impl Into<String>,
        world_hash: [u8; 32],
        group_id: u64,
        ops: &[CausalOp],
        checkpoints: &[Checkpoint],
        joules_micro: u64,
        work_steps: u64,
        parent_receipt_hash: Option<[u8; 32]>,
    ) -> SessionReceipt {
        let object_hashes: Vec<[u8; 32]> = batch_ops(ops, checkpoints)
            .iter()
            .map(|b| batch_content_hash(b))
            .collect();
        let merkle_root = merkle_root(&object_hashes);
        let mut receipt = SessionReceipt {
            world_hash,
            group_id,
            checkpoints: checkpoints.to_vec(),
            object_hashes,
            merkle_root,
            joules_micro,
            energy_provenance: None,
            work_steps,
            parent_receipt_hash,
            signer_pubkey: signer.verifying_key().to_bytes(),
            signer_id: signer_id.into(),
            sig: [0u8; 64],
        };
        receipt.sig = signer.sign(&receipt.signing_payload()).to_bytes();
        receipt
    }

    /// Label the figure with its acquisition class and re-sign
    /// (`extensions/energy-measurement.md` §2.8). `signer` must be the key
    /// that sealed the receipt. Refuses whatever [`Self::verify`] would
    /// refuse: legacy `HwShunt`, a class beside a zero figure, a refined class
    /// without its evidence.
    pub fn with_energy_class(
        mut self,
        signer: &SigningKey,
        class: EnergyProvenance,
    ) -> Result<SessionReceipt, AcquisitionError> {
        if signer.verifying_key().to_bytes() != self.signer_pubkey {
            return Err(AcquisitionError::SignerMismatch);
        }
        check_label(self.joules_micro, Some(&class))?;
        self.energy_provenance = Some(class);
        self.sig = signer.sign(&self.signing_payload()).to_bytes();
        Ok(self)
    }

    /// Re-sign whatever the fields now say, bypassing every check — so tests
    /// can present `verify()` with a correctly signed receipt that breaks a
    /// label rule, and see the rule itself refuse it.
    #[cfg(test)]
    pub(crate) fn resign_for_test(&mut self, signer: &SigningKey) {
        self.sig = signer.sign(&self.signing_payload()).to_bytes();
    }

    /// Stable identity of this receipt: `BLAKE3("wai:world-receipt-id" ||
    /// signing_payload || sig)`. The next receipt in the chain names
    /// this value as its `parent_receipt_hash`.
    pub fn receipt_hash(&self) -> [u8; 32] {
        let mut h = blake3::Hasher::new();
        h.update(b"wai:world-receipt-id\x01");
        h.update(&self.signing_payload());
        h.update(&self.sig);
        *h.finalize().as_bytes()
    }

    /// Verify internal consistency (Merkle root recomputes from the
    /// leaves) and the Ed25519 signature. Pairing the receipt to an
    /// actual op log is the verifier's job — see [`Self::matches_ops`].
    pub fn verify(&self) -> bool {
        if merkle_root(&self.object_hashes) != self.merkle_root {
            return false;
        }
        if !label_coheres(self.joules_micro, self.energy_provenance.as_ref()) {
            return false;
        }
        crate::sig::verify_strict(&self.signer_pubkey, &self.signing_payload(), &self.sig).is_ok()
    }

    /// Confirm this receipt commits to exactly the given ops: rebatch
    /// them against the receipt's checkpoints and compare object
    /// hashes. Together with [`Self::verify`] and a replay that
    /// reproduces `checkpoints`, this binds receipt ⇄ session.
    pub fn matches_ops(&self, ops: &[CausalOp]) -> bool {
        let recomputed: Vec<[u8; 32]> = batch_ops(ops, &self.checkpoints)
            .iter()
            .map(|b| batch_content_hash(b))
            .collect();
        recomputed == self.object_hashes
    }

    /// Canonical JSON (sorted keys, lowercase hex, no insignificant
    /// whitespace) — the wire form for the receipt sidecar.
    pub fn to_json(&self) -> String {
        let hx = |b: &[u8]| b.iter().map(|x| format!("{x:02x}")).collect::<String>();
        let cps: Vec<String> = self
            .checkpoints
            .iter()
            .map(|c| format!("{{\"hash\":\"{}\",\"tick\":{}}}", hx(&c.hash), c.tick))
            .collect();
        let objs: Vec<String> = self
            .object_hashes
            .iter()
            .map(|h| format!("\"{}\"", hx(h)))
            .collect();
        let parent = match self.parent_receipt_hash {
            Some(h) => format!("\"{}\"", hx(&h)),
            None => "null".into(),
        };
        let json = format!(
            "{{\"kind\":\"world-session\",\"checkpoints\":[{}],\"group_id\":{},\"joules_micro\":{},\"leaf_count\":{},\"object_hashes\":[{}],\"parent_receipt_hash\":{},\"receipt_hash\":\"{}\",\"root_hash\":\"{}\",\"sig\":\"{}\",\"signer_id\":{},\"signer_pubkey\":\"{}\",\"work_steps\":{},\"world_hash\":\"{}\"}}",
            cps.join(","),
            self.group_id,
            self.joules_micro,
            self.object_hashes.len(),
            objs.join(","),
            parent,
            hx(&self.receipt_hash()),
            hx(&self.merkle_root),
            hx(&self.sig),
            serde_json::to_string(&self.signer_id).unwrap(),
            hx(&self.signer_pubkey),
            self.work_steps,
            hx(&self.world_hash),
        );
        json_with_label(json, self.energy_provenance.as_ref())
    }
}

/// Verify an ordered chain of session receipts forms a valid
/// append-only timeline: every receipt's signature verifies, the first
/// has no parent, each subsequent `parent_receipt_hash` equals the
/// previous receipt's [`SessionReceipt::receipt_hash`], `group_id`
/// increments by one, and (optionally) all share one `world_hash` and
/// signer. Returns the index of the first broken link, or `Ok(())`.
pub fn verify_chain(chain: &[SessionReceipt]) -> Result<(), usize> {
    for (i, r) in chain.iter().enumerate() {
        if !r.verify() {
            return Err(i);
        }
        if i == 0 {
            if r.parent_receipt_hash.is_some() {
                return Err(0);
            }
        } else {
            let prev = &chain[i - 1];
            if r.parent_receipt_hash != Some(prev.receipt_hash())
                || r.group_id != prev.group_id + 1
                || r.world_hash != prev.world_hash
            {
                return Err(i);
            }
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::world::sim::{kin_mutation, pack_fx};
    use crate::world::{mint_node_id, Fx};
    use ed25519_dalek::SigningKey;

    fn k(s: u8) -> SigningKey {
        SigningKey::from_bytes(&[s; 32])
    }

    fn sample_ops() -> Vec<CausalOp> {
        let a = k(1);
        let n = mint_node_id(&a.verifying_key().to_bytes(), 1);
        vec![
            CausalOp::sign(
                &a, 1, 1, vec![], n, kin_mutation::CREATE,
                pack_fx(&[Fx::ZERO; 7]),
            ),
            CausalOp::sign(&a, 2, 50, vec![], n, kin_mutation::SET_SCALAR, pack_fx(&[Fx::from_int(3)])),
            CausalOp::sign(&a, 3, 150, vec![], n, kin_mutation::SET_SCALAR, pack_fx(&[Fx::from_int(9)])),
        ]
    }

    fn cps() -> Vec<Checkpoint> {
        vec![
            Checkpoint { tick: 100, hash: [0x11; 32] },
            Checkpoint { tick: 200, hash: [0x22; 32] },
        ]
    }

    #[test]
    fn build_verifies_and_tamper_fails() {
        let r = SessionReceipt::build(&k(9), "did:key:test", [0xab; 32], 0, &sample_ops(), &cps(), 12_345, 999, None);
        assert!(r.verify());
        assert!(r.matches_ops(&sample_ops()));

        let mut bad = r.clone();
        bad.joules_micro += 1; // signed field
        assert!(!bad.verify(), "mutated joules must break the signature");

        let mut bad2 = r.clone();
        bad2.merkle_root[0] ^= 1;
        assert!(!bad2.verify(), "mutated root must fail");
    }

    #[test]
    fn batching_pairs_ops_with_checkpoint_intervals() {
        let batches = batch_ops(&sample_ops(), &cps());
        assert_eq!(batches.len(), 2);
        // ticks 1 and 50 land in (0,100]; tick 150 in (100,200].
        assert!(!batches[0].is_empty());
        assert!(!batches[1].is_empty());
        // Different op sets → different object hashes → different root.
        let r1 = SessionReceipt::build(&k(9), "id", [0; 32], 0, &sample_ops(), &cps(), 0, 0, None);
        let mut fewer = sample_ops();
        fewer.pop();
        let r2 = SessionReceipt::build(&k(9), "id", [0; 32], 0, &fewer, &cps(), 0, 0, None);
        assert_ne!(r1.merkle_root, r2.merkle_root);
        assert!(!r1.matches_ops(&fewer));
    }

    #[test]
    fn empty_session_has_zero_root() {
        let r = SessionReceipt::build(&k(9), "id", [0; 32], 0, &[], &[], 0, 0, None);
        assert_eq!(r.merkle_root, [0u8; 32]);
        assert!(r.verify());
    }

    #[test]
    fn receipts_chain_into_an_append_only_timeline() {
        let wh = [0x55; 32];
        let r0 = SessionReceipt::build(&k(9), "id", wh, 0, &sample_ops(), &cps(), 0, 100, None);
        let r1 = SessionReceipt::build(
            &k(9), "id", wh, 1, &sample_ops(), &cps(), 0, 200, Some(r0.receipt_hash()),
        );
        let r2 = SessionReceipt::build(
            &k(9), "id", wh, 2, &sample_ops(), &cps(), 0, 300, Some(r1.receipt_hash()),
        );
        assert_eq!(super::verify_chain(&[r0.clone(), r1.clone(), r2.clone()]), Ok(()));

        // Break a link: wrong parent.
        let bad = SessionReceipt::build(&k(9), "id", wh, 2, &sample_ops(), &cps(), 0, 300, Some([0xff; 32]));
        assert_eq!(super::verify_chain(&[r0.clone(), r1.clone(), bad]), Err(2));

        // Break a link: non-monotonic group_id.
        let skip = SessionReceipt::build(
            &k(9), "id", wh, 5, &sample_ops(), &cps(), 0, 300, Some(r1.receipt_hash()),
        );
        assert_eq!(super::verify_chain(&[r0.clone(), r1.clone(), skip]), Err(2));

        // First receipt must have no parent.
        let r0_bad = SessionReceipt::build(&k(9), "id", wh, 0, &sample_ops(), &cps(), 0, 100, Some([1; 32]));
        assert_eq!(super::verify_chain(&[r0_bad]), Err(0));

        // receipt_hash is deterministic + sensitive to content.
        assert_eq!(r0.receipt_hash(), r0.receipt_hash());
        assert_ne!(r0.receipt_hash(), r1.receipt_hash());
    }

    #[test]
    fn work_steps_is_signed_and_carried_joules_is_too() {
        let a = SessionReceipt::build(&k(9), "id", [0; 32], 0, &sample_ops(), &cps(), 1000, 555, None);
        assert!(a.verify());
        // Mutating either signed field breaks the signature.
        let mut w = a.clone();
        w.work_steps += 1;
        assert!(!w.verify());
        let mut j = a.clone();
        j.joules_micro += 1;
        assert!(!j.verify());
    }

    #[test]
    fn json_is_canonical_and_stable() {
        let r = SessionReceipt::build(&k(9), "did:key:z6Mk", [0xcd; 32], 1, &sample_ops(), &cps(), 7, 42, None);
        let j = r.to_json();
        // sorted keys: checkpoints < group_id < joules_micro < leaf_count
        assert!(j.find("\"checkpoints\"").unwrap() < j.find("\"group_id\"").unwrap());
        assert!(j.find("\"joules_micro\"").unwrap() < j.find("\"leaf_count\"").unwrap());
        // parses as valid JSON
        let _: serde_json::Value = serde_json::from_str(&j).unwrap();
    }

    /// The signature is verified strictly (SPEC §7 *Signatures*): a key of small order, a
    /// non-canonical spelling of one, and an R of small order are each refused, though a
    /// lenient check accepts each.
    #[test]
    fn the_signature_is_verified_strictly() {
        let r = SessionReceipt::build(&k(9), "did:key:test", [0xab; 32], 0, &sample_ops(), &cps(), 12_345, 999, None);
        crate::sig::assert_refuses_forgeries(&r, |r, k, s| (r.signer_pubkey, r.sig) = (k, s), |r| r.signing_payload(), |r| r.verify());
    }
}
