//! ISO BMFF boxes as live-provenance §3.2, §4.1, §4.5 and §7.2 write and read them: the
//! top-level layout of a segment, the C2PA Manifest Box, the VSI `emsg` box, the chunks, the
//! track's `moov`, and the BMFF hash of C2PA 2.4 §18.6.2 over whole-box exclusions.

use sha2::{Digest, Sha256};

/// The extended type of C2PA's `uuid` box (C2PA 2.4 Appendix A.5.1).
pub const C2PA_UUID: [u8; 16] =
    [0xD8, 0xFE, 0xC3, 0xD6, 0x1B, 0x0E, 0x48, 0x3C, 0x92, 0x97, 0x58, 0x28, 0x87, 0x7E, 0xC4, 0x81];
/// The scheme of a VSI box (C2PA 2.4 §19.4.2).
pub const VSI_SCHEME: &str = "urn:c2pa:verifiable-segment-info";
/// The value of a VSI box (C2PA 2.4 §19.4.2).
pub const VSI_VALUE: &str = "fseg";
/// The `box_purpose` of a C2PA Manifest Box (Appendix A.5.3).
pub const PURPOSE_MANIFEST: &str = "manifest";
/// The `box_purpose` of the standard store once an update store follows it (Appendix A.5.3).
pub const PURPOSE_ORIGINAL: &str = "original";
/// The `box_purpose` of the update store at the end of a file (Appendix A.5.3).
pub const PURPOSE_UPDATE: &str = "update";

/// `tfhd` flags: default-base-is-moof only (§3.2, packaging §7.6).
pub const TFHD_FLAGS: u32 = 0x0002_0000;
/// `trun` flags: data offset, sample duration, size and flags (§3.2, packaging §7.6).
pub const TRUN_FLAGS: u32 = 0x0000_0701;

/// One box: its type and its span `[start, end)` in the bytes it was read from.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BoxRef {
    pub ty: [u8; 4],
    pub start: usize,
    pub end: usize,
}

impl BoxRef {
    pub fn len(&self) -> usize {
        self.end - self.start
    }

    pub fn is_empty(&self) -> bool {
        self.end == self.start
    }

    /// The box's content after its 8-byte header.
    pub fn body<'a>(&self, b: &'a [u8]) -> &'a [u8] {
        &b[self.start + 8..self.end]
    }
}

fn u32_at(b: &[u8], at: usize) -> Option<u32> {
    b.get(at..at + 4).map(|s| u32::from_be_bytes(s.try_into().expect("four bytes")))
}

fn u64_at(b: &[u8], at: usize) -> Option<u64> {
    b.get(at..at + 8).map(|s| u64::from_be_bytes(s.try_into().expect("eight bytes")))
}

/// The boxes of `b[from..to]`, one after another, each with a 32-bit size of at least 8 that
/// fits. A size of 0 or 1 (to the end, or 64-bit) is not written by this profile and does not
/// read.
pub fn boxes_in(b: &[u8], from: usize, to: usize) -> Option<Vec<BoxRef>> {
    let mut out = Vec::new();
    let mut p = from;
    while p < to {
        let size = u32_at(b, p)? as usize;
        if size < 8 || to - p < size {
            return None;
        }
        let ty: [u8; 4] = b[p + 4..p + 8].try_into().expect("four bytes");
        out.push(BoxRef { ty, start: p, end: p + size });
        p += size;
    }
    Some(out)
}

/// The top-level boxes of `b`.
pub fn top_level(b: &[u8]) -> Option<Vec<BoxRef>> {
    boxes_in(b, 0, b.len())
}

/// The children of a container box `r`, after `skip` bytes of its own fields.
pub fn children(b: &[u8], r: &BoxRef, skip: usize) -> Option<Vec<BoxRef>> {
    if r.len() < 8 + skip {
        return None;
    }
    boxes_in(b, r.start + 8 + skip, r.end)
}

/// A box of type `ty` around `payload`.
pub fn write_box(ty: &[u8; 4], payload: &[u8]) -> Vec<u8> {
    let mut o = Vec::with_capacity(8 + payload.len());
    o.extend_from_slice(&((8 + payload.len()) as u32).to_be_bytes());
    o.extend_from_slice(ty);
    o.extend_from_slice(payload);
    o
}

/// A full box: version, 24 bits of flags, then `payload`.
pub fn full_box(ty: &[u8; 4], version: u8, flags: u32, payload: &[u8]) -> Vec<u8> {
    let mut p = Vec::with_capacity(4 + payload.len());
    p.push(version);
    p.extend_from_slice(&flags.to_be_bytes()[1..]);
    p.extend_from_slice(payload);
    write_box(ty, &p)
}

/// A full box's version and flags.
fn version_flags(b: &[u8], r: &BoxRef) -> Option<(u8, u32)> {
    if r.len() < 12 {
        return None;
    }
    let v = u32_at(b, r.start + 8)?;
    Some(((v >> 24) as u8, v & 0x00FF_FFFF))
}

// ───────────────────────────── the C2PA boxes ─────────────────────────────

/// The C2PA Manifest Box of §4.1 around a manifest store.
pub fn manifest_box(store: &[u8]) -> Vec<u8> {
    c2pa_box(PURPOSE_MANIFEST, store)
}

/// A C2PA `uuid` box of `purpose` around a manifest store (§4.1, §4.7).
pub fn c2pa_box(purpose: &str, store: &[u8]) -> Vec<u8> {
    let mut p = Vec::with_capacity(16 + 4 + purpose.len() + 1 + 8 + store.len());
    p.extend_from_slice(&C2PA_UUID);
    p.extend_from_slice(&[0, 0, 0, 0]);
    p.extend_from_slice(purpose.as_bytes());
    p.push(0);
    p.extend_from_slice(&[0u8; 8]);
    p.extend_from_slice(store);
    write_box(b"uuid", &p)
}

/// Whether `r` is a C2PA `uuid` box.
pub fn is_c2pa_uuid(b: &[u8], r: &BoxRef) -> bool {
    &r.ty == b"uuid" && b.get(r.start + 8..r.start + 24) == Some(&C2PA_UUID[..])
}

/// The manifest store a C2PA Manifest Box holds (§4.1): version 0, flags 0, purpose
/// `manifest`, a zero merkle offset, then the store, which is the rest of the box.
pub fn read_manifest_box<'a>(b: &'a [u8], r: &BoxRef) -> Option<&'a [u8]> {
    read_c2pa_box(b, r, PURPOSE_MANIFEST)
}

/// The store a C2PA `uuid` box of `purpose` holds (§4.1, §4.7).
pub fn read_c2pa_box<'a>(b: &'a [u8], r: &BoxRef, purpose: &str) -> Option<&'a [u8]> {
    let body = r.body(b);
    let rest = body.strip_prefix(&C2PA_UUID[..])?;
    let rest = rest.strip_prefix(&[0u8, 0, 0, 0][..])?;
    let rest = rest.strip_prefix(purpose.as_bytes())?;
    let rest = rest.strip_prefix(&[0u8][..])?;
    let rest = rest.strip_prefix(&[0u8; 8][..])?;
    Some(rest)
}

/// An `emsg` box of version 0 (ISO/IEC 23009-1 §5.10.3.3.3).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Emsg {
    pub scheme_id_uri: String,
    pub value: String,
    pub timescale: u32,
    pub presentation_time_delta: u32,
    pub event_duration: u32,
    pub id: u32,
    pub message_data: Vec<u8>,
}

impl Emsg {
    pub fn to_box(&self) -> Vec<u8> {
        let mut p = Vec::new();
        p.extend_from_slice(self.scheme_id_uri.as_bytes());
        p.push(0);
        p.extend_from_slice(self.value.as_bytes());
        p.push(0);
        for v in [self.timescale, self.presentation_time_delta, self.event_duration, self.id] {
            p.extend_from_slice(&v.to_be_bytes());
        }
        p.extend_from_slice(&self.message_data);
        full_box(b"emsg", 0, 0, &p)
    }

    /// Read an `emsg` of version 0 and flags 0: two strings each ended by a zero byte, four
    /// u32s, and the message.
    pub fn read(b: &[u8], r: &BoxRef) -> Option<Emsg> {
        if &r.ty != b"emsg" || version_flags(b, r)? != (0, 0) {
            return None;
        }
        let body = &b[r.start + 12..r.end];
        let z1 = body.iter().position(|&c| c == 0)?;
        let scheme = std::str::from_utf8(&body[..z1]).ok()?.to_string();
        let rest = &body[z1 + 1..];
        let z2 = rest.iter().position(|&c| c == 0)?;
        let value = std::str::from_utf8(&rest[..z2]).ok()?.to_string();
        let rest = &rest[z2 + 1..];
        if rest.len() < 16 {
            return None;
        }
        Some(Emsg {
            scheme_id_uri: scheme,
            value,
            timescale: u32_at(rest, 0)?,
            presentation_time_delta: u32_at(rest, 4)?,
            event_duration: u32_at(rest, 8)?,
            id: u32_at(rest, 12)?,
            message_data: rest[16..].to_vec(),
        })
    }
}

/// The scheme a VSI box holds at offset 12, with its terminating zero byte: the value of
/// §4.5's fourth exclusion.
pub fn vsi_marker() -> Vec<u8> {
    let mut v = VSI_SCHEME.as_bytes().to_vec();
    v.push(0);
    v
}

/// Whether `r` is a VSI box: an `emsg` whose bytes at offset 12 are the scheme and its zero.
pub fn is_vsi(b: &[u8], r: &BoxRef) -> bool {
    let m = vsi_marker();
    &r.ty == b"emsg" && b.get(r.start + 12..r.start + 12 + m.len()) == Some(&m[..])
}

// ───────────────────────────── the hash ─────────────────────────────

/// One exclusion of §4.5: a top-level type and the bytes it must hold at offsets from its
/// first byte.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Exclusion {
    pub ty: [u8; 4],
    pub data: Vec<(u64, Vec<u8>)>,
}

impl Exclusion {
    fn matches(&self, b: &[u8], r: &BoxRef) -> bool {
        r.ty == self.ty
            && self.data.iter().all(|(off, v)| {
                let s = r.start + *off as usize;
                s.checked_add(v.len()).is_some_and(|e| e <= r.end) && &b[s..s + v.len()] == v.as_slice()
            })
    }
}

/// §4.5's four exclusions, in order.
pub fn profile_exclusions() -> Vec<Exclusion> {
    vec![
        Exclusion { ty: *b"uuid", data: vec![(8, C2PA_UUID.to_vec())] },
        Exclusion { ty: *b"ftyp", data: vec![] },
        Exclusion { ty: *b"mfra", data: vec![] },
        Exclusion { ty: *b"emsg", data: vec![(12, vsi_marker())] },
    ]
}

/// The BMFF hash of §4.5 (C2PA 2.4 §18.6.2): SHA-256 over each top-level box no exclusion
/// matches, as its u64 offset then its bytes.
pub fn segment_hash(b: &[u8], top: &[BoxRef], exclusions: &[Exclusion]) -> [u8; 32] {
    let mut h = Sha256::new();
    for r in top {
        if exclusions.iter().any(|e| e.matches(b, r)) {
            continue;
        }
        h.update((r.start as u64).to_be_bytes());
        h.update(&b[r.start..r.end]);
    }
    h.finalize().into()
}

// ───────────────────────────── chunks and the moov ─────────────────────────────

/// One sample of a chunk's `trun`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Sample {
    pub duration: u32,
    pub size: u32,
    pub flags: u32,
}

/// A chunk (§3.2, packaging §7.6): `moof` holding `mfhd` and `traf` (`tfhd`
/// default-base-is-moof, `tfdt` version 1, `trun` version 0 pointing at the `mdat` payload),
/// then `mdat`.
pub fn chunk(sequence_number: u32, track_id: u32, base_media_decode_time: u64, samples: &[Sample], data: &[u8]) -> Vec<u8> {
    let mfhd = full_box(b"mfhd", 0, 0, &sequence_number.to_be_bytes());
    let tfhd = full_box(b"tfhd", 0, TFHD_FLAGS, &track_id.to_be_bytes());
    let tfdt = full_box(b"tfdt", 1, 0, &base_media_decode_time.to_be_bytes());
    let trun_len = 8 + 4 + 4 + 4 + 12 * samples.len();
    let traf_len = 8 + tfhd.len() + tfdt.len() + trun_len;
    let moof_len = 8 + mfhd.len() + traf_len;
    let mut tp = Vec::new();
    tp.extend_from_slice(&(samples.len() as u32).to_be_bytes());
    tp.extend_from_slice(&((moof_len + 8) as i32).to_be_bytes());
    for s in samples {
        tp.extend_from_slice(&s.duration.to_be_bytes());
        tp.extend_from_slice(&s.size.to_be_bytes());
        tp.extend_from_slice(&s.flags.to_be_bytes());
    }
    let trun = full_box(b"trun", 0, TRUN_FLAGS, &tp);
    let traf = write_box(b"traf", &[tfhd, tfdt, trun].concat());
    let moof = write_box(b"moof", &[mfhd, traf].concat());
    debug_assert_eq!(moof.len(), moof_len);
    [moof, write_box(b"mdat", data)].concat()
}

/// What a chunk's `moof` says.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ChunkInfo {
    pub sequence_number: u32,
    pub track_id: u32,
    pub base_media_decode_time: u64,
    /// The sum of the sample durations.
    pub duration: u64,
}

fn only(b: &[u8], r: &BoxRef, skip: usize, types: &[&[u8; 4]]) -> Option<Vec<BoxRef>> {
    let c = children(b, r, skip)?;
    if c.len() != types.len() || c.iter().zip(types).any(|(x, t)| &x.ty != *t) {
        return None;
    }
    Some(c)
}

/// Read a chunk as §3.2 lays one out: `moof` with exactly `mfhd` and `traf`, the `traf` with
/// exactly `tfhd`, `tfdt` and `trun` in this profile's versions and flags, the `trun`'s data
/// offset at the `mdat`'s payload and its sizes summing to it.
pub fn read_chunk(b: &[u8], moof: &BoxRef, mdat: &BoxRef) -> Option<ChunkInfo> {
    if &moof.ty != b"moof" || &mdat.ty != b"mdat" || mdat.start != moof.end {
        return None;
    }
    let mc = only(b, moof, 0, &[b"mfhd", b"traf"])?;
    let (mfhd, traf) = (mc[0], mc[1]);
    if version_flags(b, &mfhd)? != (0, 0) || mfhd.len() != 16 {
        return None;
    }
    let sequence_number = u32_at(b, mfhd.start + 12)?;
    let tc = only(b, &traf, 0, &[b"tfhd", b"tfdt", b"trun"])?;
    let (tfhd, tfdt, trun) = (tc[0], tc[1], tc[2]);
    if version_flags(b, &tfhd)? != (0, TFHD_FLAGS) || tfhd.len() != 16 {
        return None;
    }
    let track_id = u32_at(b, tfhd.start + 12)?;
    if version_flags(b, &tfdt)? != (1, 0) || tfdt.len() != 20 {
        return None;
    }
    let base_media_decode_time = u64_at(b, tfdt.start + 12)?;
    if version_flags(b, &trun)? != (0, TRUN_FLAGS) {
        return None;
    }
    let count = u32_at(b, trun.start + 12)? as usize;
    if trun.len() != 20 + 12 * count {
        return None;
    }
    let offset = i32::from_be_bytes(b[trun.start + 16..trun.start + 20].try_into().expect("four bytes"));
    if offset as i64 != moof.len() as i64 + 8 {
        return None;
    }
    let (mut duration, mut size) = (0u64, 0u64);
    for i in 0..count {
        let at = trun.start + 20 + 12 * i;
        duration += u32_at(b, at)? as u64;
        size += u32_at(b, at + 4)? as u64;
    }
    if size != (mdat.len() - 8) as u64 {
        return None;
    }
    Some(ChunkInfo { sequence_number, track_id, base_media_decode_time, duration })
}

/// What the initialization segment's `moov` says (§11.3 step 1).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Moov {
    pub track_id: u32,
    pub timescale: u32,
}

/// Read a `moov` of one track: its `tkhd`'s `track_ID`, its `mdhd`'s `timescale` (versions 0
/// and 1), and an `mvex` holding a `trex`.
pub fn read_moov(b: &[u8], moov: &BoxRef) -> Option<Moov> {
    let mc = children(b, moov, 0)?;
    let traks: Vec<&BoxRef> = mc.iter().filter(|r| &r.ty == b"trak").collect();
    let mvex = mc.iter().find(|r| &r.ty == b"mvex")?;
    if traks.len() != 1 || !children(b, mvex, 0)?.iter().any(|r| &r.ty == b"trex") {
        return None;
    }
    let tc = children(b, traks[0], 0)?;
    let tkhd = tc.iter().find(|r| &r.ty == b"tkhd")?;
    let mdia = tc.iter().find(|r| &r.ty == b"mdia")?;
    let mdc = children(b, mdia, 0)?;
    let mdhd = mdc.iter().find(|r| &r.ty == b"mdhd")?;
    // The field after the creation and modification times: 32-bit times in version 0,
    // 64-bit in version 1.
    let field = |r: &BoxRef| -> Option<u32> {
        let at = match version_flags(b, r)?.0 {
            0 => 20,
            1 => 28,
            _ => return None,
        };
        if r.len() < at + 4 {
            return None;
        }
        u32_at(b, r.start + at)
    };
    let track_id = field(tkhd)?;
    let timescale = field(mdhd)?;
    if timescale == 0 {
        return None;
    }
    Some(Moov { track_id, timescale })
}

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

    #[test]
    fn a_chunk_reads_back() {
        let s = [Sample { duration: 1, size: 3, flags: 0x0200_0000 }, Sample { duration: 1, size: 2, flags: 0x0101_0000 }];
        let c = chunk(7, 1, 40, &s, b"abcde");
        let top = top_level(&c).unwrap();
        assert_eq!(top.len(), 2);
        let info = read_chunk(&c, &top[0], &top[1]).unwrap();
        assert_eq!(info, ChunkInfo { sequence_number: 7, track_id: 1, base_media_decode_time: 40, duration: 2 });
        // A data offset that misses the mdat payload, or sizes that do not sum to it, do not read.
        let mut bad = c.clone();
        let last = bad.len() - 1;
        bad.truncate(last);
        let n = bad.len();
        let mdat_start = top[1].start;
        bad[mdat_start..mdat_start + 4].copy_from_slice(&((n - mdat_start) as u32).to_be_bytes());
        let t2 = top_level(&bad).unwrap();
        assert!(read_chunk(&bad, &t2[0], &t2[1]).is_none());
    }

    #[test]
    fn an_emsg_reads_back_and_is_recognised() {
        let e = Emsg {
            scheme_id_uri: VSI_SCHEME.into(),
            value: VSI_VALUE.into(),
            timescale: 25,
            presentation_time_delta: 0,
            event_duration: 60,
            id: 1,
            message_data: vec![1, 2, 3],
        };
        let b = e.to_box();
        let r = top_level(&b).unwrap()[0];
        assert!(is_vsi(&b, &r));
        assert_eq!(Emsg::read(&b, &r).unwrap(), e);
    }

    #[test]
    fn the_hash_skips_excluded_boxes_and_binds_offsets() {
        let a = write_box(b"free", b"x");
        let m = manifest_box(b"store");
        let c = write_box(b"moov", b"y");
        let one = [m.clone(), a.clone(), c.clone()].concat();
        let two = [m.clone(), c.clone(), a.clone()].concat();
        let ex = profile_exclusions();
        let h1 = segment_hash(&one, &top_level(&one).unwrap(), &ex);
        let h2 = segment_hash(&two, &top_level(&two).unwrap(), &ex);
        assert_ne!(h1, h2, "moving a box changes the hash");
        let other = [manifest_box(b"other store"), a, c].concat();
        let h3 = segment_hash(&other, &top_level(&other).unwrap(), &ex);
        assert_ne!(h1, h3, "the manifest box is excluded, but its length moves the offsets after it");
        let same_len = [manifest_box(b"STORE"), write_box(b"free", b"x"), write_box(b"moov", b"y")].concat();
        assert_eq!(h1, segment_hash(&same_len, &top_level(&same_len).unwrap(), &ex));
    }
}
