//! Deterministic packaging + attested SegmentMap — `wai.video.package`.
//!
//! The ingest/packaging layer: a source is sliced into segments, each with an
//! index and a content hash, and the ordered `(index -> content_hash)` map is
//! the SegmentMap the manifest layer references. Two things the incumbent
//! packagers leave implicit become explicit and checkable here.
//!
//! - **The SegmentMap is signed.** `segment_map_root` is a Merkle commitment
//!   over the segment leaves, so a `wai.video.manifest` that references segment
//!   `i` can be checked against the package that actually produced it — the two
//!   receipts compose.
//! - **The dedup is a fact, not a billing line.** Identical content — an ad
//!   played in every break, a bumper, a replay, the overlap between simulcast
//!   feeds — hashes identically and is stored and encoded once.
//!   [`PackageReceipt::verify`] recomputes `unique_segments` and `stored_bytes`
//!   from the segments themselves, so the storage/encode saving cannot be
//!   overstated. `uj_per_unique_segment` is the encode-once energy, and it is
//!   honest precisely because the uniqueness is recomputable.
//!
//! Only the measurement — `joules_micro`, its acquisition class and the class's
//! evidence — is measured-attested: signed, and not re-measurable. Reuses the
//! manifest module's `EnergyProvenance` and the world Merkle floor's Ed25519 + `GrantRef`.


use ed25519_dalek::{Signer, SigningKey};

use crate::quantum_ops::GrantRef;
use crate::energy_class::{figure_and_class_cohere, write_optional};
use crate::wai_video_manifest::EnergyProvenance;
use crate::world::receipt::merkle_root;

const DOMAIN_RECEIPT: &[u8] = b"wai:video-package-receipt\x01";
const DOMAIN_RECEIPT_ID: &[u8] = b"wai:video-package-receipt-id\x01";
const DOMAIN_SEG: &[u8] = b"wai:video-package-seg\x01";
const DOMAIN_MEDIA: &[u8] = b"wai:video-package-media\x01";

/// One segment in the SegmentMap.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SegmentEntry {
    pub index: u64,
    /// BLAKE3 over the segment's media bytes — the identity that makes two
    /// byte-identical segments dedup.
    pub content_hash: [u8; 32],
    pub duration_ms: u64,
    pub byte_len: u64,
}

impl SegmentEntry {
    fn leaf(&self) -> [u8; 32] {
        let mut h = blake3::Hasher::new();
        h.update(DOMAIN_SEG);
        h.update(&self.index.to_be_bytes());
        h.update(&self.content_hash);
        h.update(&self.duration_ms.to_be_bytes());
        h.update(&self.byte_len.to_be_bytes());
        *h.finalize().as_bytes()
    }
}

/// A source GOP (group-of-pictures): the smallest independently-decodable unit,
/// which a segment boundary must fall on. `(media bytes, presentation duration)`.
pub type Gop<'a> = (&'a [u8], u64);

/// BLAKE3 over concatenated GOP media bytes — a segment's content identity.
fn media_hash(gops: &[Gop]) -> [u8; 32] {
    let mut h = blake3::Hasher::new();
    h.update(DOMAIN_MEDIA);
    for (bytes, _) in gops {
        h.update(&(bytes.len() as u64).to_be_bytes());
        h.update(bytes);
    }
    *h.finalize().as_bytes()
}

/// Deterministic segmentation: accumulate GOPs until the segment reaches
/// `target_ms`, then close it on the GOP boundary. A pure function of the source
/// and the target, so the SegmentMap is reproducible and content that repeats
/// byte-for-byte lands on identical `content_hash`es.
pub fn segment(gops: &[Gop], target_ms: u64) -> Vec<SegmentEntry> {
    let mut out = Vec::new();
    let mut cur: Vec<Gop> = Vec::new();
    let mut dur = 0u64;
    let mut bytes = 0u64;
    let close = |cur: &mut Vec<Gop>, dur: &mut u64, bytes: &mut u64, out: &mut Vec<SegmentEntry>| {
        if !cur.is_empty() {
            out.push(SegmentEntry {
                index: out.len() as u64,
                content_hash: media_hash(cur),
                duration_ms: *dur,
                byte_len: *bytes,
            });
            cur.clear();
            *dur = 0;
            *bytes = 0;
        }
    };
    for &(b, d) in gops {
        cur.push((b, d));
        dur += d;
        bytes += b.len() as u64;
        if dur >= target_ms {
            close(&mut cur, &mut dur, &mut bytes, &mut out);
        }
    }
    close(&mut cur, &mut dur, &mut bytes, &mut out);
    out
}

/// The count of distinct segment contents, and the bytes a dedup store keeps
/// (one copy per distinct content). Both recomputed from the segments so a
/// receipt cannot overstate the saving (`content_dedup::totals_checked`, the one
/// dedup computation of the crate). `None` for a map that gives one content hash two
/// byte lengths: it has no dedup accounting, and a receipt over it does not verify.
fn dedup_stats(segments: &[SegmentEntry]) -> Option<(u64, u64)> {
    let t = crate::content_dedup::totals_checked(segments.iter().map(|s| (s.content_hash, s.byte_len))).ok()?;
    Some((t.unique, t.stored_bytes))
}

/// A signed SegmentMap with an attested dedup accounting.
#[derive(Clone, Debug, PartialEq)]
pub struct PackageReceipt {
    pub asset_id: String,
    /// Merkle commitment over the segment leaves — the SegmentMap identity a
    /// manifest is checked against.
    pub segment_map_root: [u8; 32],
    pub segments: Vec<SegmentEntry>,
    pub total_segments: u64,
    /// Distinct `content_hash`es — recomputed in verify().
    pub unique_segments: u64,
    pub total_bytes: u64,
    /// Bytes a dedup store keeps (one copy per distinct content) — recomputed.
    pub stored_bytes: u64,
    pub work_segment_ops: u64,
    pub joules_micro: u64,
    pub energy_provenance: Option<EnergyProvenance>,
    pub grant: GrantRef,
    pub parent_receipt_hash: Option<[u8; 32]>,
    pub signer_pubkey: [u8; 32],
    pub signer_id: String,
    pub sig: [u8; 64],
}

impl PackageReceipt {
    /// The bytes the signature covers, under the key the receipt names.
    pub fn signing_payload(&self) -> Vec<u8> {
        let mut o = DOMAIN_RECEIPT.to_vec();
        o.extend_from_slice(&(self.asset_id.len() as u32).to_be_bytes());
        o.extend_from_slice(self.asset_id.as_bytes());
        o.extend_from_slice(&self.segment_map_root);
        o.extend_from_slice(&self.total_segments.to_be_bytes());
        o.extend_from_slice(&self.unique_segments.to_be_bytes());
        o.extend_from_slice(&self.total_bytes.to_be_bytes());
        o.extend_from_slice(&self.stored_bytes.to_be_bytes());
        o.extend_from_slice(&self.work_segment_ops.to_be_bytes());
        o.extend_from_slice(&self.joules_micro.to_be_bytes());
        write_optional(self.energy_provenance.as_ref(), &mut o);
        o.extend_from_slice(&self.grant.grant_hash);
        o.extend_from_slice(&(self.grant.capability.len() as u32).to_be_bytes());
        o.extend_from_slice(self.grant.capability.as_bytes());
        o.extend_from_slice(&self.grant.joule_ceiling_micro.to_be_bytes());
        match self.grant.funds_ceiling {
            Some(f) => {
                o.push(1);
                o.extend_from_slice(&f.to_be_bytes());
            }
            None => o.push(0),
        }
        match self.parent_receipt_hash {
            Some(h) => {
                o.push(1);
                o.extend_from_slice(&h);
            }
            None => o.push(0),
        }
        o.extend_from_slice(&(self.segments.len() as u64).to_be_bytes());
        for s in &self.segments {
            o.extend_from_slice(&s.leaf());
        }
        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
    }

    /// Seal a SegmentMap. The dedup figures are computed here from the segments,
    /// so they cannot be set to a saving the map does not support.
    #[allow(clippy::too_many_arguments)]
    pub fn seal(
        signer: &SigningKey,
        signer_id: impl Into<String>,
        asset_id: impl Into<String>,
        segments: Vec<SegmentEntry>,
        joules_micro: u64,
        energy_provenance: Option<EnergyProvenance>,
        grant: GrantRef,
        parent_receipt_hash: Option<[u8; 32]>,
    ) -> PackageReceipt {
        let leaves: Vec<[u8; 32]> = segments.iter().map(SegmentEntry::leaf).collect();
        // A map that gives one content hash two lengths has no dedup accounting: it is
        // sealed with both figures 0, and `verify` refuses it.
        let (unique_segments, stored_bytes) = dedup_stats(&segments).unwrap_or((0, 0));
        let total_bytes = segments.iter().map(|s| s.byte_len).sum();
        let mut r = PackageReceipt {
            asset_id: asset_id.into(),
            segment_map_root: merkle_root(&leaves),
            total_segments: segments.len() as u64,
            unique_segments,
            total_bytes,
            stored_bytes,
            work_segment_ops: segments.len() as u64,
            segments,
            joules_micro,
            energy_provenance,
            grant,
            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
    }

    /// Signature + internal consistency: the Merkle root, the counts and the
    /// recomputed dedup figures must all follow from the carried segments, and
    /// the energy figure and its label stand or fall together.
    pub fn verify(&self) -> bool {
        let leaves: Vec<[u8; 32]> = self.segments.iter().map(SegmentEntry::leaf).collect();
        if merkle_root(&leaves) != self.segment_map_root {
            return false;
        }
        if self.total_segments != self.segments.len() as u64 {
            return false;
        }
        if self.work_segment_ops != self.segments.len() as u64 {
            return false;
        }
        let Some((uniq, stored)) = dedup_stats(&self.segments) else { return false };
        if uniq != self.unique_segments || stored != self.stored_bytes {
            return false;
        }
        if self.total_bytes != self.segments.iter().map(|s| s.byte_len).sum::<u64>() {
            return false;
        }
        for (i, s) in self.segments.iter().enumerate() {
            if s.index != i as u64 {
                return false;
            }
        }
        // The figure and its class stand or fall together, and a hardware
        // class must carry its declared uncertainty (energy_class).
        if !figure_and_class_cohere(self.joules_micro, self.energy_provenance.as_ref()) {
            return false;
        }
        crate::sig::verify_strict(&self.signer_pubkey, &self.signing_payload(), &self.sig).is_ok()
    }

    /// Whether the segment at `index` in the attested map has this content —
    /// how a `wai.video.manifest` reference is checked against the package.
    pub fn segment_content(&self, index: u64) -> Option<[u8; 32]> {
        self.segments.get(index as usize).map(|s| s.content_hash)
    }

    /// Bytes a dedup store saves versus storing every segment. Verifiable, since
    /// both terms are recomputed in [`PackageReceipt::verify`].
    pub fn bytes_saved(&self) -> u64 {
        self.total_bytes.saturating_sub(self.stored_bytes)
    }

    /// Encode-once energy: microjoules per *distinct* segment, with its
    /// acquisition class. `None` when unmetered.
    pub fn uj_per_unique_segment(&self) -> Option<(f64, EnergyProvenance)> {
        match (&self.energy_provenance, self.unique_segments) {
            (Some(p), u) if self.joules_micro > 0 && u > 0 => {
                Some((self.joules_micro as f64 / u as f64, p.clone()))
            }
            _ => None,
        }
    }

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

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

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

    // A source where the ad GOP (`AD`) recurs: content, ad, content, ad.
    fn source() -> Vec<Gop<'static>> {
        vec![
            (b"content-gop-0000".as_ref(), 2000),
            (b"AD-creative-0001".as_ref(), 2000),
            (b"content-gop-0002".as_ref(), 2000),
            (b"AD-creative-0001".as_ref(), 2000),
        ]
    }

    #[test]
    fn segmentation_is_deterministic_and_dedups_repeats() {
        // target 2000ms -> one GOP per segment; the two ad segments are identical.
        let a = segment(&source(), 2000);
        let b = segment(&source(), 2000);
        assert_eq!(a, b, "segmentation must be a pure function");
        assert_eq!(a.len(), 4);
        assert_eq!(a[1].content_hash, a[3].content_hash, "the repeated ad dedups");
        let (uniq, _) = dedup_stats(&a).unwrap();
        assert_eq!(uniq, 3, "4 segments, 3 distinct");
    }

    #[test]
    fn seal_verify_and_dedup_is_a_fact() {
        let segs = segment(&source(), 2000);
        let g = GrantRef::unbounded("video.package");
        let r = PackageReceipt::seal(&key(1), "did:key:pkg", "asset-abc", segs, 12_000, Some(crate::energy_class::test_on_chip()), g, None);
        assert!(r.verify());
        assert_eq!(r.total_segments, 4);
        assert_eq!(r.unique_segments, 3);
        assert_eq!(r.bytes_saved(), 16, "one 16-byte ad segment is stored once, not twice");
        // A manifest referencing index 3 can be checked against the package.
        assert_eq!(r.segment_content(3), Some(r.segments[1].content_hash));
    }

    #[test]
    fn overstated_dedup_will_not_verify() {
        let segs = segment(&source(), 2000);
        let g = GrantRef::unbounded("video.package");
        let mut r = PackageReceipt::seal(&key(2), "o", "asset", segs, 0, None, g, None);
        assert!(r.verify());
        // Claim more saving than the map supports.
        r.unique_segments = 2;
        r.stored_bytes = 32;
        r.sig = key(2).sign(&r.signing_payload()).to_bytes();
        assert!(!r.verify(), "a dedup figure the segments do not support must fail");
    }

    #[test]
    fn one_content_with_two_lengths_has_no_dedup_accounting() {
        // Segment 3 repeats segment 1's content hash with another length: no store can
        // keep "one copy" of it, so the map has no dedup figure and its receipt fails.
        let mut segs = segment(&source(), 2000);
        segs[3].byte_len += 1;
        let g = GrantRef::unbounded("video.package");
        let r = PackageReceipt::seal(&key(5), "o", "asset", segs, 0, None, g, None);
        assert_eq!((r.unique_segments, r.stored_bytes), (0, 0));
        assert!(!r.verify(), "a map with one hash at two lengths must not verify");
    }

    #[test]
    fn tampered_segment_breaks_the_map() {
        let segs = segment(&source(), 2000);
        let g = GrantRef::unbounded("video.package");
        let mut r = PackageReceipt::seal(&key(3), "o", "asset", segs, 0, None, g, None);
        assert!(r.verify());
        r.segments[0].content_hash = [0xFF; 32]; // swap a segment's content
        assert!(!r.verify(), "a swapped segment must break the SegmentMap root");
    }

    #[test]
    fn encode_once_energy_amortises_over_unique_segments() {
        let segs = segment(&source(), 2000);
        let g = GrantRef::unbounded("video.package");
        // 9000 uJ of packaging over 3 distinct segments.
        let r = PackageReceipt::seal(&key(4), "o", "asset", segs, 9000, Some(EnergyProvenance::ModelBased), g, None);
        let (uj, prov) = r.uj_per_unique_segment().unwrap();
        assert_eq!(prov, EnergyProvenance::ModelBased);
        assert!((uj - 3000.0).abs() < 1e-9, "9000 uJ / 3 unique = 3000 uJ each");
    }


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

    /// A receipt sealed with the legacy `HwShunt` label keeps its signed bytes.
    #[test]
    #[allow(deprecated)] // sealing a legacy HwShunt receipt is the point
    fn legacy_hwshunt_receipt_keeps_its_signed_bytes() {
        let segs = segment(&source(), 2000);
        let g = GrantRef::unbounded("video.package");
        let r = PackageReceipt::seal(&key(1), "did:key:pkg", "asset-abc", segs, 12_000, Some(EnergyProvenance::HwShunt), g, None);
        assert!(r.verify());
        assert_eq!(hex(&r.receipt_hash()), "71a9da59a3202abb2c3fd22ee692b1485a9c674bb8c311f4c82455e6be29ae2b");
    }

    /// 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 segs = segment(&source(), 2000);
        let g = GrantRef::unbounded("video.package");
        let r = PackageReceipt::seal(&key(1), "did:key:pkg", "asset-abc", segs, 12_000, Some(crate::energy_class::test_on_chip()), g, None);
        crate::sig::assert_refuses_forgeries(&r, |r, k, s| (r.signer_pubkey, r.sig) = (k, s), |r| r.signing_payload(), |r| r.verify());
    }
}
