//! Asset provenance — `wai.asset.provenance` (lever 4).
//!
//! The other cargo classes are the "carries whatever it can reconstruct"
//! half of WAI: a world replays, a mix sums, a cloud interpolates. This
//! module is the other half — **"verifies whatever it carries"** — for
//! the content that *cannot* be reduced to instructions: a photograph, a
//! recording, a telemetry stream. WAI does not decode or regenerate such
//! an asset; it carries the opaque bytes (content-addressed in a WAI
//! envelope) and attaches a signed, energy-accounted **lineage**.
//!
//! A [`ProvStep`] is one processing step — an operation, the content hash
//! of its input and output, the exact work, and the measured/carried
//! joules. An [`AssetReceipt`] is a JWP group over those steps: the
//! asset's content hash is the identity, the per-step hashes are the
//! Merkle leaves (JWP's exact domains, reused), and one Ed25519 signature
//! covers asset hash + capability + the chain + total joules. It chains
//! across edits via `parent_receipt_hash`, like every other receipt here.
//!
//! What this buys without reconstruction: **tamper-evidence** (recompute
//! the asset's content hash, compare), **lineage** (the signed step
//! chain), and **energy accounting** (measured joules per step) — for
//! bytes WAI never has to understand.
//!
//! Each step's figure carries its acquisition class as an optional label
//! (`extensions/energy-measurement.md` §2.8), because the steps of one
//! lineage can be metered by different parties in different ways. A step
//! with no label keeps the legacy leaf hash; a labelled step hashes under the
//! `0x02` step domain with its class after the figure. A receipt with no
//! labelled step keeps the legacy signing payload byte for byte; one with any
//! labelled step signs under the `0x02` receipt domain and writes every
//! step's optional class (`0` when unlabelled) after that step's figure.
//! `total_joules_micro` carries no label of its own: a sum over steps of
//! different classes is only as strong as its weakest step. An unmetered step
//! adds its marker `0` to the sum, so a total over steps that include one covers
//! only the measured steps (`Partial`, energy-measurement §2) and a reader says
//! so; it is never the whole cost of the lineage.

use crate::energy_class::{
    check_label, json_with_label, label_coheres, label_from_json, labelled_domain, write_label,
    write_optional, AcquisitionError, EnergyProvenance,
};
use crate::world::receipt::merkle_root;
use ed25519_dalek::{Signer, SigningKey};

const DOMAIN_RECEIPT: &[u8] = b"wai:asset-receipt\x01";
const DOMAIN_STEP: &[u8] = b"wai:prov-step\x01";

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

fn from_hex32(s: &str) -> Option<[u8; 32]> {
    if s.len() != 64 || !s.bytes().all(|b| b.is_ascii_hexdigit()) {
        return None;
    }
    let mut out = [0u8; 32];
    for (i, c) in s.as_bytes().chunks(2).enumerate() {
        out[i] = u8::from_str_radix(std::str::from_utf8(c).ok()?, 16).ok()?;
    }
    Some(out)
}

/// One processing step in an asset's lineage.
#[derive(Clone, Debug, PartialEq)]
pub struct ProvStep {
    /// Operation, e.g. `"capture"`, `"encode.jpeg"`, `"derive.thumbnail"`.
    pub op: String,
    /// Content hash of the input (all-zero for an originating step).
    pub input_hash: [u8; 32],
    /// Content hash of the output.
    pub output_hash: [u8; 32],
    /// Exact, deterministic work (e.g. pixels or samples processed).
    pub work: u64,
    /// Measured/carried microjoules for the step; `0` is the unmetered marker,
    /// never a figure of zero joules (energy-measurement §2).
    pub joules_micro: u64,
    /// How the step's `joules_micro` was acquired — an optional label
    /// (`extensions/energy-measurement.md` §2.8). `None` is an unlabelled
    /// figure, as every step sealed before the label existed carries; it is
    /// never read as any class. Bound into the step's leaf hash when present.
    pub energy_provenance: Option<EnergyProvenance>,
}

impl ProvStep {
    /// An unlabelled step, as every step was before the label existed.
    pub fn new(
        op: impl Into<String>,
        input_hash: [u8; 32],
        output_hash: [u8; 32],
        work: u64,
        joules_micro: u64,
    ) -> ProvStep {
        ProvStep { op: op.into(), input_hash, output_hash, work, joules_micro, energy_provenance: None }
    }

    /// Label the step's figure with its acquisition class. Refuses whatever
    /// [`AssetReceipt::verify`] would refuse: legacy `HwShunt`, a class beside
    /// a zero figure, a refined class without its evidence.
    pub fn with_energy_class(mut self, class: EnergyProvenance) -> Result<ProvStep, AcquisitionError> {
        check_label(self.joules_micro, Some(&class))?;
        self.energy_provenance = Some(class);
        Ok(self)
    }

    /// Content hash of the step — the Merkle leaf.
    pub fn content_hash(&self) -> [u8; 32] {
        let label = self.energy_provenance.as_ref();
        let mut h = blake3::Hasher::new();
        h.update(&labelled_domain(DOMAIN_STEP, label));
        let op = self.op.as_bytes();
        h.update(&(op.len() as u32).to_be_bytes());
        h.update(op);
        h.update(&self.input_hash);
        h.update(&self.output_hash);
        h.update(&self.work.to_be_bytes());
        h.update(&self.joules_micro.to_be_bytes());
        let mut l = Vec::new();
        write_label(label, &mut l);
        h.update(&l);
        *h.finalize().as_bytes()
    }
}

/// A signed, energy-accounted lineage for a content-addressed asset.
#[derive(Clone, Debug, PartialEq)]
pub struct AssetReceipt {
    /// Content hash of the carried asset (BLAKE3 over its complete WAI
    /// envelope, [`crate::world::asset_content_hash`]). Equals the final
    /// step's `output_hash`.
    pub asset_hash: [u8; 32],
    /// Informational media class, e.g. `"image"`, `"audio"`.
    pub media: String,
    /// The WAI capability the asset declares, e.g. `wai.image.jpeg`.
    pub capability: String,
    pub group_id: u64,
    pub steps: Vec<ProvStep>,
    pub merkle_root: [u8; 32],
    /// Σ of the steps' carried joules. A step carrying the unmetered marker adds
    /// nothing, so with one the total covers only the measured steps.
    pub total_joules_micro: u64,
    pub parent_receipt_hash: Option<[u8; 32]>,
    pub signer_pubkey: [u8; 32],
    pub signer_id: String,
    pub sig: [u8; 64],
}

impl AssetReceipt {
    /// Whether any step carries a label — which puts the receipt's signing
    /// payload under the `0x02` domain.
    fn labelled(&self) -> Option<&EnergyProvenance> {
        self.steps.iter().find_map(|s| s.energy_provenance.as_ref())
    }

    fn signing_payload(&self) -> Vec<u8> {
        let any_label = self.labelled();
        let mut o = labelled_domain(DOMAIN_RECEIPT, any_label);
        o.extend_from_slice(&self.asset_hash);
        for s in [self.media.as_bytes(), self.capability.as_bytes()] {
            o.extend_from_slice(&(s.len() as u32).to_be_bytes());
            o.extend_from_slice(s);
        }
        o.extend_from_slice(&self.group_id.to_be_bytes());
        o.extend_from_slice(&self.merkle_root);
        o.extend_from_slice(&(self.steps.len() as u64).to_be_bytes());
        o.extend_from_slice(&self.total_joules_micro.to_be_bytes());
        match self.parent_receipt_hash {
            Some(h) => {
                o.push(1);
                o.extend_from_slice(&h);
            }
            None => o.push(0),
        }
        for s in &self.steps {
            let op = s.op.as_bytes();
            o.extend_from_slice(&(op.len() as u32).to_be_bytes());
            o.extend_from_slice(op);
            o.extend_from_slice(&s.input_hash);
            o.extend_from_slice(&s.output_hash);
            o.extend_from_slice(&s.work.to_be_bytes());
            o.extend_from_slice(&s.joules_micro.to_be_bytes());
            if any_label.is_some() {
                write_optional(s.energy_provenance.as_ref(), &mut o);
            }
        }
        o.extend_from_slice(&self.signer_pubkey);
        let id = self.signer_id.as_bytes();
        o.extend_from_slice(&(id.len() as u32).to_be_bytes());
        o.extend_from_slice(id);
        o
    }

    /// Build and sign a provenance receipt over an ordered step chain.
    /// `asset_hash` MUST equal the last step's `output_hash` (the asset
    /// the receipt vouches for).
    #[allow(clippy::too_many_arguments)]
    pub fn build(
        signer: &SigningKey,
        signer_id: impl Into<String>,
        asset_hash: [u8; 32],
        media: impl Into<String>,
        capability: impl Into<String>,
        group_id: u64,
        steps: Vec<ProvStep>,
        parent_receipt_hash: Option<[u8; 32]>,
    ) -> AssetReceipt {
        let leaves: Vec<[u8; 32]> = steps.iter().map(|s| s.content_hash()).collect();
        let total_joules_micro = steps.iter().map(|s| s.joules_micro).sum();
        let mut r = AssetReceipt {
            asset_hash,
            media: media.into(),
            capability: capability.into(),
            group_id,
            steps,
            merkle_root: merkle_root(&leaves),
            total_joules_micro,
            parent_receipt_hash,
            signer_pubkey: signer.verifying_key().to_bytes(),
            signer_id: signer_id.into(),
            sig: [0u8; 64],
        };
        r.sig = signer.sign(&r.signing_payload()).to_bytes();
        r
    }

    /// Verify the Merkle root recomputes, the step chain is linked
    /// (each step's input is the previous step's output), the final
    /// output equals `asset_hash`, and the signature is valid.
    pub fn verify(&self) -> bool {
        let leaves: Vec<[u8; 32]> = self.steps.iter().map(|s| s.content_hash()).collect();
        if merkle_root(&leaves) != self.merkle_root {
            return false;
        }
        // chain linkage: step k input == step k-1 output
        for w in self.steps.windows(2) {
            if w[1].input_hash != w[0].output_hash {
                return false;
            }
        }
        if let Some(last) = self.steps.last()
            && last.output_hash != self.asset_hash
        {
            return false;
        }
        if self.steps.iter().map(|s| s.joules_micro).sum::<u64>() != self.total_joules_micro {
            return false;
        }
        if !self.steps.iter().all(|s| label_coheres(s.joules_micro, s.energy_provenance.as_ref())) {
            return false;
        }
        crate::sig::verify_strict(&self.signer_pubkey, &self.signing_payload(), &self.sig).is_ok()
    }

    /// Confirm the receipt vouches for exactly this asset envelope:
    /// recompute the content hash and compare. The tamper test — a
    /// 1-bit change to the envelope breaks it.
    /// The total as a reader presents it (energy-measurement §2): `UNMETERED`
    /// when every step carries the unmetered marker, never `0 µJ`; the figure
    /// labelled `Whole` when every step is measured; otherwise the figure
    /// labelled `Partial`, naming the unmetered steps the total does not cover.
    ///
    /// The figure carries the class the total reads at (energy-measurement §2.8): the
    /// weakest class among the steps with a figure, or `unlabelled` when one of them
    /// carries no label, so a total is never shown without what it is.
    pub fn total_label(&self) -> String {
        let unmetered: Vec<&str> = self.steps.iter().filter(|s| s.joules_micro == 0).map(|s| s.op.as_str()).collect();
        let metered = self.steps.iter().filter(|s| s.joules_micro > 0);
        let class = if metered.clone().any(|s| s.energy_provenance.is_none()) {
            "unlabelled"
        } else {
            metered.filter_map(|s| s.energy_provenance.as_ref()).min_by_key(|p| p.rank()).map_or("unlabelled", EnergyProvenance::label)
        };
        if unmetered.len() == self.steps.len() {
            "UNMETERED (no step measured)".to_owned()
        } else if unmetered.is_empty() {
            format!("{} µJ ({class}; Whole)", self.total_joules_micro)
        } else {
            format!(
                "{} µJ ({class}) over the measured steps only (Partial; unmetered: {})",
                self.total_joules_micro,
                unmetered.join(", ")
            )
        }
    }

    pub fn matches_asset(&self, envelope_bytes: &[u8]) -> bool {
        crate::world::asset_content_hash(envelope_bytes) == self.asset_hash
    }

    /// 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();
    }

    pub fn receipt_hash(&self) -> [u8; 32] {
        let mut h = blake3::Hasher::new();
        h.update(b"wai:asset-receipt-id\x01");
        h.update(&self.signing_payload());
        h.update(&self.sig);
        *h.finalize().as_bytes()
    }

    pub fn to_json(&self) -> String {
        let steps: Vec<String> = self
            .steps
            .iter()
            .map(|s| {
                let step = format!(
                    "{{\"input_hash\":\"{}\",\"joules_micro\":{},\"op\":{},\"output_hash\":\"{}\",\"work\":{}}}",
                    hx(&s.input_hash),
                    s.joules_micro,
                    serde_json::to_string(&s.op).unwrap(),
                    hx(&s.output_hash),
                    s.work
                );
                json_with_label(step, s.energy_provenance.as_ref())
            })
            .collect();
        let parent = match self.parent_receipt_hash {
            Some(h) => format!("\"{}\"", hx(&h)),
            None => "null".into(),
        };
        format!(
            "{{\"kind\":\"asset-provenance\",\"asset_hash\":\"{}\",\"capability\":{},\"group_id\":{},\"media\":{},\"parent_receipt_hash\":{},\"receipt_hash\":\"{}\",\"root_hash\":\"{}\",\"sig\":\"{}\",\"signer_id\":{},\"signer_pubkey\":\"{}\",\"steps\":[{}],\"total_joules_micro\":{}}}",
            hx(&self.asset_hash),
            serde_json::to_string(&self.capability).unwrap(),
            self.group_id,
            serde_json::to_string(&self.media).unwrap(),
            parent,
            hx(&self.receipt_hash()),
            hx(&self.merkle_root),
            hx(&self.sig),
            serde_json::to_string(&self.signer_id).unwrap(),
            hx(&self.signer_pubkey),
            steps.join(","),
            self.total_joules_micro,
        )
    }

    /// Parse a receipt from its canonical JSON (for a verifying sink). A
    /// malformed step label is refused rather than dropped, so a parse can
    /// never turn a labelled step into an unlabelled one.
    pub fn from_json(s: &str) -> Option<AssetReceipt> {
        let v: serde_json::Value = serde_json::from_str(s).ok()?;
        let o = v.as_object()?;
        let h = |k: &str| from_hex32(o.get(k)?.as_str()?);
        let str_f = |k: &str| o.get(k).and_then(|x| x.as_str()).map(str::to_owned);
        let u = |k: &str| o.get(k).and_then(|x| x.as_u64());
        let mut steps = Vec::new();
        for sv in o.get("steps")?.as_array()? {
            let so = sv.as_object()?;
            steps.push(ProvStep {
                op: so.get("op")?.as_str()?.to_owned(),
                input_hash: from_hex32(so.get("input_hash")?.as_str()?)?,
                output_hash: from_hex32(so.get("output_hash")?.as_str()?)?,
                work: so.get("work")?.as_u64()?,
                joules_micro: so.get("joules_micro")?.as_u64()?,
                energy_provenance: label_from_json(so).ok()?,
            });
        }
        let parent = match o.get("parent_receipt_hash") {
            Some(serde_json::Value::String(s)) => Some(from_hex32(s)?),
            _ => None,
        };
        Some(AssetReceipt {
            asset_hash: h("asset_hash")?,
            media: str_f("media")?,
            capability: str_f("capability")?,
            group_id: u("group_id")?,
            steps,
            merkle_root: h("root_hash")?,
            total_joules_micro: u("total_joules_micro")?,
            parent_receipt_hash: parent,
            signer_pubkey: from_hex32(o.get("signer_pubkey")?.as_str()?)?,
            signer_id: str_f("signer_id")?,
            sig: {
                let s = o.get("sig")?.as_str()?;
                if s.len() != 128 {
                    return None;
                }
                let mut out = [0u8; 64];
                for (i, c) in s.as_bytes().chunks(2).enumerate() {
                    out[i] = u8::from_str_radix(std::str::from_utf8(c).ok()?, 16).ok()?;
                }
                out
            },
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

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

    // capture (origin) → encode → derive thumbnail. asset_hash = last output.
    fn chain() -> (Vec<ProvStep>, [u8; 32]) {
        let raw = [0x11u8; 32];
        let jpeg = [0x22u8; 32];
        let thumb = [0x33u8; 32];
        let steps = vec![
            ProvStep { op: "capture".into(), input_hash: [0u8; 32], output_hash: raw, work: 12_000_000, joules_micro: 0, energy_provenance: None },
            ProvStep { op: "encode.jpeg".into(), input_hash: raw, output_hash: jpeg, work: 12_000_000, joules_micro: 4200, energy_provenance: None },
            ProvStep { op: "derive.thumbnail".into(), input_hash: jpeg, output_hash: thumb, work: 65_536, joules_micro: 90, energy_provenance: None },
        ];
        (steps, thumb)
    }

    #[test]
    fn receipt_links_verifies_and_tamper_fails() {
        let (steps, asset) = chain();
        let r = AssetReceipt::build(&key(9), "did:web:wai#prov", asset, "image", "wai.image.jpeg", 0, steps, None);
        assert!(r.verify());
        assert_eq!(r.total_joules_micro, 4290);

        // mutate a signed field → sig fails
        let mut bad = r.clone();
        bad.total_joules_micro += 1;
        assert!(!bad.verify());
        // break the chain linkage → fails before the sig even matters
        let mut broke = r.clone();
        broke.steps[1].input_hash[0] ^= 1;
        assert!(!broke.verify());
        // asset_hash must equal the final output
        let mut wrong = r.clone();
        wrong.asset_hash[0] ^= 1;
        assert!(!wrong.verify());
    }

    /// An unmetered step is never presented as 0 J, and a total over steps that
    /// include one says it covers only the measured steps.
    #[test]
    fn the_total_is_labelled_as_a_reader_presents_it() {
        let (steps, asset) = chain();
        let build = |steps: Vec<ProvStep>| AssetReceipt::build(&key(9), "did:web:wai#prov", asset, "image", "wai.image.jpeg", 0, steps, None);
        let partial = build(steps.clone());
        assert_eq!(partial.total_label(), "4290 µJ (unlabelled) over the measured steps only (Partial; unmetered: capture)");
        let none = build(steps.iter().cloned().map(|s| ProvStep { joules_micro: 0, ..s }).collect());
        assert_eq!(none.total_joules_micro, 0);
        assert_eq!(none.total_label(), "UNMETERED (no step measured)", "never 0 µJ");
        let whole = build(steps.into_iter().map(|s| ProvStep { joules_micro: s.joules_micro.max(7), ..s }).collect());
        assert_eq!(whole.total_label(), "4297 µJ (unlabelled; Whole)");
        // Labelled steps: the total reads at the weakest of their classes.
        let chip = crate::energy_class::test_on_chip();
        let labelled: Vec<ProvStep> = whole.steps.iter().cloned()
            .map(|s| if s.joules_micro < 100 { s.with_energy_class(EnergyProvenance::ModelBased) } else { s.with_energy_class(chip.clone()) }.unwrap())
            .collect();
        assert_eq!(build(labelled).total_label(), "4297 µJ (ModelBased; Whole)");
    }

    #[test]
    fn json_round_trips_and_verifies() {
        let (steps, asset) = chain();
        let r = AssetReceipt::build(&key(7), "did:web:wai#prov", asset, "image", "wai.image.jpeg", 2, steps, Some([0xab; 32]));
        let j = r.to_json();
        let _: serde_json::Value = serde_json::from_str(&j).unwrap();
        let back = AssetReceipt::from_json(&j).expect("parse");
        assert_eq!(back, r);
        assert!(back.verify());
    }

    #[test]
    fn receipts_chain() {
        let (steps, asset) = chain();
        let r0 = AssetReceipt::build(&key(9), "id", asset, "image", "wai.image.jpeg", 0, steps.clone(), None);
        let r1 = AssetReceipt::build(&key(9), "id", asset, "image", "wai.image.jpeg", 1, steps, Some(r0.receipt_hash()));
        assert!(r0.verify() && r1.verify());
        assert_eq!(r1.parent_receipt_hash, Some(r0.receipt_hash()));
        assert_ne!(r0.receipt_hash(), r1.receipt_hash());
    }

    /// 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 (steps, asset) = chain();
        let r = AssetReceipt::build(&key(9), "did:web:wai#prov", asset, "image", "wai.image.jpeg", 0, steps, None);
        crate::sig::assert_refuses_forgeries(&r, |r, k, s| (r.signer_pubkey, r.sig) = (k, s), |r| r.signing_payload(), |r| r.verify());
    }
}
