//! Quantum Bill of Materials — `wai.quantum.qbom`
//! (extensions/quantum-ops § Provenance).
//!
//! An open quantum operations stack runs the job — submits the program, manages the
//! device, returns counts. What it does not do is **sign what the result is made
//! of**: which compiled circuit ran (and that it was equivalent to the source),
//! against which calibration, under which learned noise model, with which error
//! mitigation, on a register prepared how. Each of those stages, in this repo,
//! already emits a signed receipt. A **Quantum Bill of Materials** assembles the
//! stage receipts of one computation into a single signed, content-addressed
//! **provenance manifest** — the audit sidecar the operations layer doesn't itself
//! produce.
//!
//! The QBOM lists each stage by `(kind, receipt_hash, joules)`, binds them to a
//! job identity and program hash under one signature, and exposes a **provenance
//! root** (a content hash over the whole manifest) plus the **total energy** across
//! stages. Verification is layered: the QBOM verifies its own signature and root;
//! presented alongside the actual stage receipts, [`Qbom::binds`] confirms each
//! listed hash matches a receipt that itself verifies — so a swapped or forged
//! stage breaks the bill. It is byte-exact and reproducible: the same computation
//! rebuilds the same root from the git repo.
//!
//! This is a manifest, not a new physics claim — it inherits exactly the honesty
//! of the receipts it bundles, and adds the property that a quantum result now
//! carries a single verifiable statement of everything it was made of.

use crate::quantum_ops::content_hash;
use ed25519_dalek::{Signature, SigningKey, Verifier, VerifyingKey};

const DOMAIN_QBOM: &[u8] = b"wai:quantum-qbom\x01";
const DOMAIN_QBOM_ID: &[u8] = b"wai:quantum-qbom-id\x01";

/// The canonical stages of a quantum computation, in pipeline order. A QBOM need
/// not carry all of them; `Qbom::completeness` reports which are present.
pub const CANONICAL_STAGES: &[&str] = &[
    "compile",
    "rearrange",
    "calibrate",
    "noise_learn",
    "execute",
    "mitigate",
    "decode",
];

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 {
        return None;
    }
    let mut out = [0u8; 32];
    for (i, c) in s.as_bytes().as_chunks::<2>().0.iter().enumerate() {
        out[i] = u8::from_str_radix(std::str::from_utf8(c).ok()?, 16).ok()?;
    }
    Some(out)
}
fn from_hex64(s: &str) -> Option<[u8; 64]> {
    if s.len() != 128 {
        return None;
    }
    let mut out = [0u8; 64];
    for (i, c) in s.as_bytes().as_chunks::<2>().0.iter().enumerate() {
        out[i] = u8::from_str_radix(std::str::from_utf8(c).ok()?, 16).ok()?;
    }
    Some(out)
}

/// One line item: a stage of the computation, referenced by its receipt hash.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct QbomEntry {
    pub kind: String,
    pub receipt_hash: [u8; 32],
    pub joules_micro: u64,
}

/// A signed Quantum Bill of Materials for one computation.
#[derive(Clone, Debug, PartialEq)]
pub struct Qbom {
    pub job_id: String,
    pub program_hash: [u8; 32],
    pub device_id: String,
    pub entries: Vec<QbomEntry>,
    pub signer_pubkey: [u8; 32],
    pub signer_id: String,
    pub sig: [u8; 64],
}

/// Builder for a QBOM before it is sealed.
#[derive(Clone, Debug)]
pub struct QbomBuilder {
    job_id: String,
    program_hash: [u8; 32],
    device_id: String,
    entries: Vec<QbomEntry>,
}

impl QbomBuilder {
    pub fn new(job_id: impl Into<String>, program_hash: [u8; 32], device_id: impl Into<String>) -> QbomBuilder {
        QbomBuilder {
            job_id: job_id.into(),
            program_hash,
            device_id: device_id.into(),
            entries: Vec::new(),
        }
    }
    /// Add a stage by its receipt hash and metered energy. Chainable.
    pub fn stage(mut self, kind: impl Into<String>, receipt_hash: [u8; 32], joules_micro: u64) -> QbomBuilder {
        self.entries.push(QbomEntry { kind: kind.into(), receipt_hash, joules_micro });
        self
    }
    pub fn seal(self, signer: &SigningKey, signer_id: impl Into<String>) -> Qbom {
        let mut q = Qbom {
            job_id: self.job_id,
            program_hash: self.program_hash,
            device_id: self.device_id,
            entries: self.entries,
            signer_pubkey: signer.verifying_key().to_bytes(),
            signer_id: signer_id.into(),
            sig: [0u8; 64],
        };
        q.sig = signer.sign_payload(&q.signing_payload());
        q
    }
}

// small helper so we don't need the Signer trait import name-clash
trait SignPayload {
    fn sign_payload(&self, msg: &[u8]) -> [u8; 64];
}
impl SignPayload for SigningKey {
    fn sign_payload(&self, msg: &[u8]) -> [u8; 64] {
        use ed25519_dalek::Signer;
        self.sign(msg).to_bytes()
    }
}

impl Qbom {
    /// Canonical bytes of the manifest body (everything but the signature).
    fn body(&self) -> Vec<u8> {
        let mut o = DOMAIN_QBOM.to_vec();
        o.extend_from_slice(&(self.job_id.len() as u32).to_le_bytes());
        o.extend_from_slice(self.job_id.as_bytes());
        o.extend_from_slice(&self.program_hash);
        o.extend_from_slice(&(self.device_id.len() as u32).to_le_bytes());
        o.extend_from_slice(self.device_id.as_bytes());
        o.extend_from_slice(&(self.entries.len() as u32).to_le_bytes());
        for e in &self.entries {
            o.extend_from_slice(&(e.kind.len() as u32).to_le_bytes());
            o.extend_from_slice(e.kind.as_bytes());
            o.extend_from_slice(&e.receipt_hash);
            o.extend_from_slice(&e.joules_micro.to_le_bytes());
        }
        o
    }

    /// The provenance root: a content hash over the whole manifest body. Stable
    /// and reproducible — the same computation rebuilds the same root.
    pub fn provenance_root(&self) -> [u8; 32] {
        content_hash(&self.body())
    }

    fn signing_payload(&self) -> Vec<u8> {
        let mut o = self.body();
        o.extend_from_slice(&self.signer_pubkey);
        o.extend_from_slice(self.signer_id.as_bytes());
        o
    }

    /// Total metered energy across all stages (µJ).
    pub fn total_joules_micro(&self) -> u64 {
        self.entries.iter().map(|e| e.joules_micro).sum()
    }

    /// Which canonical stages this bill carries, in pipeline order.
    pub fn completeness(&self) -> Vec<String> {
        CANONICAL_STAGES
            .iter()
            .filter(|s| self.entries.iter().any(|e| e.kind == **s))
            .map(|s| s.to_string())
            .collect()
    }

    /// The manifest's own signature is valid.
    pub fn verify(&self) -> bool {
        let Ok(k) = VerifyingKey::from_bytes(&self.signer_pubkey) else {
            return false;
        };
        k.verify(&self.signing_payload(), &Signature::from_bytes(&self.sig))
            .is_ok()
    }

    /// Given the actual stage receipts as `(receipt_hash, receipt_verifies)`, does
    /// the bill bind exactly them? Every entry must match a presented receipt that
    /// itself verifies, and every presented receipt must be listed. A swapped,
    /// forged, or failed-to-verify stage breaks it.
    pub fn binds(&self, receipts: &[([u8; 32], bool)]) -> bool {
        if !self.verify() || receipts.len() != self.entries.len() {
            return false;
        }
        // each entry hash appears among the presented receipts, and that receipt verifies
        self.entries.iter().all(|e| {
            receipts.iter().any(|(h, ok)| *h == e.receipt_hash && *ok)
        })
    }

    pub fn receipt_hash(&self) -> [u8; 32] {
        let mut h = blake3::Hasher::new();
        h.update(DOMAIN_QBOM_ID);
        h.update(&self.signing_payload());
        h.update(&self.sig);
        *h.finalize().as_bytes()
    }

    pub fn to_json(&self) -> String {
        let entries: Vec<String> = self
            .entries
            .iter()
            .map(|e| {
                format!(
                    "{{\"joules_micro\":{},\"kind\":{},\"receipt_hash\":\"{}\"}}",
                    e.joules_micro,
                    serde_json::to_string(&e.kind).unwrap(),
                    hx(&e.receipt_hash)
                )
            })
            .collect();
        format!(
            "{{\"kind\":\"quantum-qbom\",\"device_id\":{},\"entries\":[{}],\"job_id\":{},\
             \"program_hash\":\"{}\",\"provenance_root\":\"{}\",\"receipt_hash\":\"{}\",\"sig\":\"{}\",\
             \"signer_id\":{},\"signer_pubkey\":\"{}\",\"total_joules_micro\":{}}}",
            serde_json::to_string(&self.device_id).unwrap(),
            entries.join(","),
            serde_json::to_string(&self.job_id).unwrap(),
            hx(&self.program_hash),
            hx(&self.provenance_root()),
            hx(&self.receipt_hash()),
            hx(&self.sig),
            serde_json::to_string(&self.signer_id).unwrap(),
            hx(&self.signer_pubkey),
            self.total_joules_micro(),
        )
    }

    /// An operator-envelope sidecar: the provenance a job/result exchange can carry
    /// alongside the operations stack's own job record. Names the job + device it
    /// belongs to and the provenance root that anchors the bill.
    pub fn to_operator_envelope(&self) -> String {
        format!(
            "{{\"job_id\":{},\"device_id\":{},\"program_hash\":\"{}\",\"provenance\":{{\
             \"kind\":\"wai.quantum.qbom\",\"root\":\"{}\",\"stages\":{},\"total_joules_micro\":{},\
             \"signer_id\":{}}}}}",
            serde_json::to_string(&self.job_id).unwrap(),
            serde_json::to_string(&self.device_id).unwrap(),
            hx(&self.program_hash),
            hx(&self.provenance_root()),
            self.entries.len(),
            self.total_joules_micro(),
            serde_json::to_string(&self.signer_id).unwrap(),
        )
    }

    pub fn from_json(s: &str) -> Option<Qbom> {
        let v: serde_json::Value = serde_json::from_str(s).ok()?;
        let o = v.as_object()?;
        let mut entries = Vec::new();
        for e in o.get("entries")?.as_array()? {
            let eo = e.as_object()?;
            entries.push(QbomEntry {
                kind: eo.get("kind")?.as_str()?.to_owned(),
                receipt_hash: from_hex32(eo.get("receipt_hash")?.as_str()?)?,
                joules_micro: eo.get("joules_micro")?.as_u64()?,
            });
        }
        Some(Qbom {
            job_id: o.get("job_id")?.as_str()?.to_owned(),
            program_hash: from_hex32(o.get("program_hash")?.as_str()?)?,
            device_id: o.get("device_id")?.as_str()?.to_owned(),
            entries,
            signer_pubkey: from_hex32(o.get("signer_pubkey")?.as_str()?)?,
            signer_id: o.get("signer_id")?.as_str()?.to_owned(),
            sig: from_hex64(o.get("sig")?.as_str()?)?,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::quantum_ops::{CompileReceipt, GrantRef, RearrangeReceipt};

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

    fn sample_qbom() -> (Qbom, Vec<([u8; 32], bool)>) {
        // two real signed stage receipts + one directly-sealed stage
        let cr = CompileReceipt::seal(
            &key(9), "did:key:op", "clifford-tableau",
            content_hash(b"src"), content_hash(b"tgt"), content_hash(b"cpl"),
            true, 4, 2048, 120_000, GrantRef::unbounded("quantum.compile"), None,
        );
        let rr = RearrangeReceipt::seal(
            &key(9), "did:key:op", "hungarian-lsap",
            content_hash(b"init"), content_hash(b"tg"), content_hash(b"plan"),
            true, 40, 300, 90_000, GrantRef::unbounded("quantum.rearrange"), None,
        );
        let q = QbomBuilder::new("job-0xABCD", content_hash(b"program"), "sim:transmon:cluster-a")
            .stage("compile", cr.receipt_hash(), cr.joules_micro)
            .stage("rearrange", rr.receipt_hash(), rr.joules_micro)
            .seal(&key(9), "did:key:op");
        let receipts = vec![
            (cr.receipt_hash(), cr.verify()),
            (rr.receipt_hash(), rr.verify()),
        ];
        (q, receipts)
    }

    #[test]
    fn qbom_seals_and_verifies() {
        let (q, _) = sample_qbom();
        assert!(q.verify());
        assert_eq!(q.total_joules_micro(), 120_000 + 90_000);
        assert_eq!(q.completeness(), vec!["compile", "rearrange"]);
    }

    #[test]
    fn qbom_binds_its_real_receipts() {
        let (q, receipts) = sample_qbom();
        assert!(q.binds(&receipts), "the bill must bind exactly its stage receipts");
    }

    #[test]
    fn swapped_stage_breaks_the_bill() {
        let (q, mut receipts) = sample_qbom();
        // swap one receipt hash for an unrelated one → no longer binds
        receipts[0].0 = content_hash(b"a forged stage");
        assert!(!q.binds(&receipts));
    }

    #[test]
    fn failed_stage_breaks_the_bill() {
        let (q, mut receipts) = sample_qbom();
        // a listed stage whose receipt does not verify breaks the bill
        receipts[1].1 = false;
        assert!(!q.binds(&receipts));
    }

    #[test]
    fn tampering_with_the_manifest_breaks_signature() {
        let (mut q, _) = sample_qbom();
        q.entries[0].joules_micro += 1; // alter the bill after sealing
        assert!(!q.verify());
    }

    #[test]
    fn provenance_root_is_reproducible() {
        let (a, _) = sample_qbom();
        let (b, _) = sample_qbom();
        assert_eq!(a.provenance_root(), b.provenance_root());
    }

    #[test]
    fn json_round_trips() {
        let (q, _) = sample_qbom();
        let back = Qbom::from_json(&q.to_json()).unwrap();
        assert_eq!(back, q);
        assert!(back.verify());
    }
}