arcly_stream/packager/

fmp4.rs

//! A native, dependency-free **fragmented-MP4 / CMAF** muxer (`feature = "fmp4"`).
//!
//! Implements the [`Muxer`](super::Muxer) trait, producing the two artifacts
//! LL-HLS fMP4 egress needs:
//!
//! * an **initialization segment** (`ftyp` + `moov`, carrying the decoder
//!   configuration, e.g. `avcC`/`hvcC`) via
//!   [`init_segment`](super::Muxer::init_segment), referenced once from the
//!   playlist with `#EXT-X-MAP`; and
//! * **media fragments** (`moof` + `mdat`) per segment, with a `trun` sample
//!   table, signed composition-time offsets, and a `tfdt` base-media-decode-time
//!   for seamless concatenation.
//!
//! The **fragment path here is codec-generic** — samples are length-prefixed and
//! described structurally — so the same `moof`/`mdat` machinery serves any
//! codec; only the init segment's sample entry is codec-specific. This release
//! builds init segments for **H.264 (`avc1`)** and, behind their codec features,
//! **H.265 (`hvc1`, `codec-h265`)**, **AV1 (`av01`, `codec-av1`)** and **VVC
//! (`vvc1`, `codec-vvc`)** — all via one shared box writer. For H.265 and VVC the
//! profile-tier-level is read from the SPS, so the `hvcC`/`vvcC` and the HLS
//! codec string carry real values.
//!
//! ```no_run
//! # #[cfg(feature = "storage-fs")]
//! # async fn demo(handle: arcly_stream::StreamHandle) -> arcly_stream::Result<()> {
//! use arcly_stream::packager::{Fmp4Muxer, HlsSegmenter, Packager};
//! use arcly_stream::storage::FsStorage;
//!
//! // 1s parts, 6 segments of LL-HLS, written as init.m4s + segN.m4s.
//! let mut seg = HlsSegmenter::new(Fmp4Muxer::new(), FsStorage::new("/var/hls"), "live/cam", 6, 5)
//!     .low_latency(1.0);
//! let mut sub = handle.subscribe_resilient();
//! while let Some(frame) = sub.recv().await { seg.push(&frame).await?; }
//! seg.finish().await?;
//! # Ok(())
//! # }
//! ```

use super::Muxer;
use crate::{CodecId, MediaFrame, Result, StreamError};
use bytes::Bytes;

const TIMESCALE: u32 = 1000; // milliseconds — matches MediaFrame pts/dts units.
const TRACK_ID: u32 = 1;
const FALLBACK_DURATION: u64 = TIMESCALE as u64 / 30; // ~33 ms when unknown.

/// A buffered sample within the current fragment (stored length-prefixed / AVCC).
struct Sample {
    data: Bytes,
    dts: i64,
    cts: i32, // pts − dts (composition offset)
    is_key: bool,
    duration: Option<u64>,
}

/// Native fragmented-MP4 muxer for a single video track.
///
/// Construct with [`new`](Self::new), hand to an
/// [`HlsSegmenter`](super::HlsSegmenter) (or drive the [`Muxer`] trait directly).
/// The init segment is built for **H.264** (`avc1`) and, behind their codec
/// features, **H.265** (`hvc1`), **AV1** (`av01`) and **VVC** (`vvc1`); the
/// `moof`/`mdat` fragment path is codec-generic.
/// `Send` — owned by a single segmenter task. Audio frames are skipped (use a
/// dedicated audio rendition).
pub struct Fmp4Muxer {
    codec: Option<CodecId>,
    width: u16,
    height: u16,
    codec_str: Option<String>,
    seq: u32,
    base_decode_time: u64,
    samples: Vec<Sample>,
}

impl Default for Fmp4Muxer {
    fn default() -> Self {
        Self::new()
    }
}

impl Fmp4Muxer {
    /// A new muxer, awaiting the first CONFIG access unit to build its init
    /// segment.
    pub fn new() -> Self {
        Self {
            codec: None,
            width: 0,
            height: 0,
            codec_str: None,
            seq: 1,
            base_decode_time: 0,
            samples: Vec::new(),
        }
    }

    /// Build the `avc1` sample-entry box (with `avcC`) and codec string from an
    /// Annex-B SPS/PPS config.
    fn ingest_h264_config(&mut self, config_record: &[u8]) -> Result<Vec<u8>> {
        let mut sps: Option<&[u8]> = None;
        let mut pps: Option<&[u8]> = None;
        for nal in crate::codec::h264::iter_nals_annexb(config_record) {
            match nal.first().map(|b| b & 0x1f) {
                Some(t) if t == crate::codec::h264::NAL_SPS => sps = Some(nal),
                Some(t) if t == crate::codec::h264::NAL_PPS => pps = Some(nal),
                _ => {}
            }
        }
        let sps = sps.ok_or_else(|| StreamError::codec("fmp4: no SPS in H.264 config"))?;
        let pps = pps.ok_or_else(|| StreamError::codec("fmp4: no PPS in H.264 config"))?;
        if sps.len() < 4 {
            return Err(StreamError::codec("fmp4: truncated SPS"));
        }
        let info = crate::codec::h264::parse_sps(sps)
            .ok_or_else(|| StreamError::codec("fmp4: unparseable SPS"))?;
        self.width = info.width as u16;
        self.height = info.height as u16;
        // avc1.PPCCLL — profile / constraint flags / level, from the SPS bytes.
        self.codec_str = Some(format!("avc1.{:02X}{:02X}{:02X}", sps[1], sps[2], sps[3]));

        // avcC (AVCDecoderConfigurationRecord).
        let mut c = Vec::with_capacity(16 + sps.len() + pps.len());
        c.push(1); // configurationVersion
        c.push(sps[1]); // AVCProfileIndication
        c.push(sps[2]); // profile_compatibility
        c.push(sps[3]); // AVCLevelIndication
        c.push(0xFF); // 6 reserved bits | lengthSizeMinusOne = 3
        c.push(0xE1); // 3 reserved bits | numOfSPS = 1
        c.extend_from_slice(&(sps.len() as u16).to_be_bytes());
        c.extend_from_slice(sps);
        c.push(1); // numOfPPS
        c.extend_from_slice(&(pps.len() as u16).to_be_bytes());
        c.extend_from_slice(pps);
        Ok(build_avc1(self.width, self.height, &c))
    }

    /// Build the `hvc1` sample-entry box (with `hvcC`) and codec string from an
    /// Annex-B VPS/SPS/PPS config access unit.
    #[cfg(feature = "codec-h265")]
    fn ingest_h265_config(&mut self, config_record: &[u8]) -> Result<Vec<u8>> {
        use crate::codec::{h265::H265, CodecParser};
        let (params, hvcc) = crate::codec::h265::hvcc_config_record(config_record)
            .ok_or_else(|| StreamError::codec("fmp4: no parsable SPS in H.265 config"))?;
        self.width = params.width as u16;
        self.height = params.height as u16;
        self.codec_str = Some(H265::hls_codec_string(&params));
        Ok(build_hvc1(self.width, self.height, &hvcc))
    }

    /// Build the `av01` sample-entry box (with `av1C`) and codec string from an
    /// AV1 sequence-header config access unit.
    #[cfg(feature = "codec-av1")]
    fn ingest_av1_config(&mut self, config_record: &[u8]) -> Result<Vec<u8>> {
        use crate::codec::{av1::Av1, CodecParser};
        let (params, av1c) = crate::codec::av1::av1c_config_record(config_record)
            .ok_or_else(|| StreamError::codec("fmp4: no AV1 sequence header in config"))?;
        self.width = params.width as u16;
        self.height = params.height as u16;
        self.codec_str = Some(Av1::hls_codec_string(&params));
        Ok(build_av01(self.width, self.height, &av1c))
    }

    /// Build the `vvc1` sample-entry box (with `vvcC`) and codec string from a
    /// VVC parameter-set config access unit.
    #[cfg(feature = "codec-vvc")]
    fn ingest_vvc_config(&mut self, config_record: &[u8]) -> Result<Vec<u8>> {
        use crate::codec::{vvc::Vvc, CodecParser};
        let (params, vvcc) =
            crate::codec::vvc::vvcc_config_record(config_record).ok_or_else(|| {
                StreamError::codec(
                    "fmp4: VVC config needs an SPS without a PTL block (parser limit)",
                )
            })?;
        self.width = params.width as u16;
        self.height = params.height as u16;
        self.codec_str = Some(Vvc::hls_codec_string(&params));
        Ok(build_vvc1(self.width, self.height, &vvcc))
    }

    /// Per-sample durations: inter-DTS deltas, with the last sample falling back
    /// to its declared frame duration (or the previous delta).
    fn durations(&self) -> Vec<u64> {
        let n = self.samples.len();
        let mut out = Vec::with_capacity(n);
        for i in 0..n {
            let dur = if i + 1 < n {
                (self.samples[i + 1].dts - self.samples[i].dts).max(0) as u64
            } else {
                self.samples[i]
                    .duration
                    .or_else(|| out.last().copied())
                    .unwrap_or(FALLBACK_DURATION)
            };
            out.push(dur);
        }
        out
    }
}

impl Muxer for Fmp4Muxer {
    fn extension(&self) -> &'static str {
        "m4s"
    }

    fn start_segment(&mut self) -> Result<()> {
        self.samples.clear();
        Ok(())
    }

    fn write(&mut self, frame: &MediaFrame) -> Result<()> {
        // Single video track; skip audio (use a separate audio rendition).
        if !frame.is_video() {
            return Ok(());
        }
        // NAL codecs (H.264, VVC) are length-prefixed (4-byte) in mp4, not
        // Annex-B; AV1 temporal units are already in the size-delimited form
        // mp4 carries.
        let is_nal = matches!(
            self.codec,
            Some(CodecId::H264) | Some(CodecId::H265) | Some(CodecId::VVC)
        );
        let data = if is_nal {
            crate::codec::h264::annexb_to_avcc(&frame.data) // codec-agnostic NAL → 4-byte length
        } else {
            frame.data.clone()
        };
        if data.is_empty() {
            return Ok(());
        }
        self.samples.push(Sample {
            data,
            dts: frame.dts,
            cts: (frame.pts - frame.dts) as i32,
            is_key: frame.is_keyframe(),
            duration: frame.duration,
        });
        Ok(())
    }

    fn finish_segment(&mut self) -> Result<Bytes> {
        if self.samples.is_empty() {
            return Ok(Bytes::new());
        }
        let durations = self.durations();
        let total: u64 = durations.iter().sum();

        let moof = build_moof(self.seq, self.base_decode_time, &self.samples, &durations);
        let mut out = moof;
        // mdat box wrapping the concatenated sample data.
        let mdat_len: usize = self.samples.iter().map(|s| s.data.len()).sum();
        out.extend_from_slice(&((mdat_len + 8) as u32).to_be_bytes());
        out.extend_from_slice(b"mdat");
        for s in &self.samples {
            out.extend_from_slice(&s.data);
        }

        self.seq += 1;
        self.base_decode_time += total;
        self.samples.clear();
        Ok(Bytes::from(out))
    }

    fn init_segment(&mut self, codec: CodecId, config_record: &[u8]) -> Result<Option<Bytes>> {
        let sample_entry = match codec {
            CodecId::H264 => self.ingest_h264_config(config_record)?,
            #[cfg(feature = "codec-h265")]
            CodecId::H265 => self.ingest_h265_config(config_record)?,
            #[cfg(feature = "codec-av1")]
            CodecId::AV1 => self.ingest_av1_config(config_record)?,
            #[cfg(feature = "codec-vvc")]
            CodecId::VVC => self.ingest_vvc_config(config_record)?,
            _ => {
                return Err(StreamError::UnsupportedCodec(format!(
                    "fmp4 init segment: {codec:?} not supported in this build"
                )))
            }
        };
        self.codec = Some(codec);
        let mut seg = build_ftyp();
        seg.extend_from_slice(&build_moov(self.width, self.height, &sample_entry));
        Ok(Some(Bytes::from(seg)))
    }

    fn codec_string(&self) -> Option<String> {
        self.codec_str.clone()
    }
}

// ── ISO-BMFF box writers ─────────────────────────────────────────────────────

/// `size(4) + type(4) + body`.
fn bx(typ: &[u8; 4], body: &[u8]) -> Vec<u8> {
    let mut v = Vec::with_capacity(8 + body.len());
    v.extend_from_slice(&((body.len() + 8) as u32).to_be_bytes());
    v.extend_from_slice(typ);
    v.extend_from_slice(body);
    v
}

/// A FullBox: `size + type + (version<<24 | flags) + body`.
fn full(typ: &[u8; 4], version: u8, flags: u32, body: &[u8]) -> Vec<u8> {
    let mut b = Vec::with_capacity(4 + body.len());
    b.extend_from_slice(&(((version as u32) << 24) | (flags & 0x00FF_FFFF)).to_be_bytes());
    b.extend_from_slice(body);
    bx(typ, &b)
}

const MATRIX_IDENTITY: [u32; 9] = [0x0001_0000, 0, 0, 0, 0x0001_0000, 0, 0, 0, 0x4000_0000];

fn put_matrix(v: &mut Vec<u8>) {
    for m in MATRIX_IDENTITY {
        v.extend_from_slice(&m.to_be_bytes());
    }
}

fn build_ftyp() -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(b"iso5"); // major_brand
    b.extend_from_slice(&0u32.to_be_bytes()); // minor_version
    for brand in [b"iso5", b"iso6", b"mp41"] {
        b.extend_from_slice(brand);
    }
    bx(b"ftyp", &b)
}

/// Build `moov` around a prebuilt visual sample entry (`avc1` / `av01` box).
fn build_moov(width: u16, height: u16, entry: &[u8]) -> Vec<u8> {
    let mut body = Vec::new();
    body.extend_from_slice(&build_mvhd());
    body.extend_from_slice(&build_trak(width, height, entry));
    body.extend_from_slice(&build_mvex());
    bx(b"moov", &body)
}

fn build_mvhd() -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(&0u32.to_be_bytes()); // creation_time
    b.extend_from_slice(&0u32.to_be_bytes()); // modification_time
    b.extend_from_slice(&TIMESCALE.to_be_bytes());
    b.extend_from_slice(&0u32.to_be_bytes()); // duration (fragmented → 0)
    b.extend_from_slice(&0x0001_0000u32.to_be_bytes()); // rate 1.0
    b.extend_from_slice(&0x0100u16.to_be_bytes()); // volume 1.0
    b.extend_from_slice(&0u16.to_be_bytes()); // reserved
    b.extend_from_slice(&[0u8; 8]); // reserved
    put_matrix(&mut b);
    b.extend_from_slice(&[0u8; 24]); // pre_defined
    b.extend_from_slice(&2u32.to_be_bytes()); // next_track_ID
    full(b"mvhd", 0, 0, &b)
}

fn build_trak(width: u16, height: u16, entry: &[u8]) -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(&build_tkhd(width, height));
    b.extend_from_slice(&build_mdia(entry));
    bx(b"trak", &b)
}

fn build_tkhd(width: u16, height: u16) -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(&0u32.to_be_bytes()); // creation_time
    b.extend_from_slice(&0u32.to_be_bytes()); // modification_time
    b.extend_from_slice(&TRACK_ID.to_be_bytes());
    b.extend_from_slice(&0u32.to_be_bytes()); // reserved
    b.extend_from_slice(&0u32.to_be_bytes()); // duration
    b.extend_from_slice(&[0u8; 8]); // reserved
    b.extend_from_slice(&0u16.to_be_bytes()); // layer
    b.extend_from_slice(&0u16.to_be_bytes()); // alternate_group
    b.extend_from_slice(&0u16.to_be_bytes()); // volume (0 for video)
    b.extend_from_slice(&0u16.to_be_bytes()); // reserved
    put_matrix(&mut b);
    b.extend_from_slice(&((width as u32) << 16).to_be_bytes()); // 16.16 width
    b.extend_from_slice(&((height as u32) << 16).to_be_bytes()); // 16.16 height
    full(b"tkhd", 0, 0x0000_0007, &b) // enabled | in_movie | in_preview
}

fn build_mdia(entry: &[u8]) -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(&build_mdhd());
    b.extend_from_slice(&build_hdlr());
    b.extend_from_slice(&build_minf(entry));
    bx(b"mdia", &b)
}

fn build_mdhd() -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(&0u32.to_be_bytes()); // creation_time
    b.extend_from_slice(&0u32.to_be_bytes()); // modification_time
    b.extend_from_slice(&TIMESCALE.to_be_bytes());
    b.extend_from_slice(&0u32.to_be_bytes()); // duration
    b.extend_from_slice(&0x55C4u16.to_be_bytes()); // language 'und'
    b.extend_from_slice(&0u16.to_be_bytes()); // pre_defined
    full(b"mdhd", 0, 0, &b)
}

fn build_hdlr() -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(&0u32.to_be_bytes()); // pre_defined
    b.extend_from_slice(b"vide"); // handler_type
    b.extend_from_slice(&[0u8; 12]); // reserved
    b.extend_from_slice(b"VideoHandler\0");
    full(b"hdlr", 0, 0, &b)
}

fn build_minf(entry: &[u8]) -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(&full(b"vmhd", 0, 1, &[0u8; 8])); // graphicsmode + opcolor
    b.extend_from_slice(&build_dinf());
    b.extend_from_slice(&build_stbl(entry));
    bx(b"minf", &b)
}

fn build_dinf() -> Vec<u8> {
    // dref with a single self-contained 'url ' entry (flags = 1).
    let url = full(b"url ", 0, 1, &[]);
    let mut dref_body = Vec::new();
    dref_body.extend_from_slice(&1u32.to_be_bytes()); // entry_count
    dref_body.extend_from_slice(&url);
    let dref = full(b"dref", 0, 0, &dref_body);
    bx(b"dinf", &dref)
}

fn build_stbl(entry: &[u8]) -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(&build_stsd(entry));
    b.extend_from_slice(&full(b"stts", 0, 0, &0u32.to_be_bytes())); // entry_count 0
    b.extend_from_slice(&full(b"stsc", 0, 0, &0u32.to_be_bytes()));
    let mut stsz = Vec::new();
    stsz.extend_from_slice(&0u32.to_be_bytes()); // sample_size
    stsz.extend_from_slice(&0u32.to_be_bytes()); // sample_count
    b.extend_from_slice(&full(b"stsz", 0, 0, &stsz));
    b.extend_from_slice(&full(b"stco", 0, 0, &0u32.to_be_bytes()));
    bx(b"stbl", &b)
}

fn build_stsd(entry: &[u8]) -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(&1u32.to_be_bytes()); // entry_count
    b.extend_from_slice(entry);
    full(b"stsd", 0, 0, &b)
}

/// The common `VisualSampleEntry` header shared by `avc1` / `av01`, ending with
/// the codec-specific configuration box (`avcC` / `av1C`).
fn visual_sample_entry(typ: &[u8; 4], width: u16, height: u16, config_box: &[u8]) -> Vec<u8> {
    let mut b = Vec::new();
    b.extend_from_slice(&[0u8; 6]); // reserved
    b.extend_from_slice(&1u16.to_be_bytes()); // data_reference_index
    b.extend_from_slice(&0u16.to_be_bytes()); // pre_defined
    b.extend_from_slice(&0u16.to_be_bytes()); // reserved
    b.extend_from_slice(&[0u8; 12]); // pre_defined[3]
    b.extend_from_slice(&width.to_be_bytes());
    b.extend_from_slice(&height.to_be_bytes());
    b.extend_from_slice(&0x0048_0000u32.to_be_bytes()); // horizresolution 72dpi
    b.extend_from_slice(&0x0048_0000u32.to_be_bytes()); // vertresolution 72dpi
    b.extend_from_slice(&0u32.to_be_bytes()); // reserved
    b.extend_from_slice(&1u16.to_be_bytes()); // frame_count
    b.extend_from_slice(&[0u8; 32]); // compressorname
    b.extend_from_slice(&0x0018u16.to_be_bytes()); // depth
    b.extend_from_slice(&0xFFFFu16.to_be_bytes()); // pre_defined = -1
    b.extend_from_slice(config_box);
    bx(typ, &b)
}

/// `avc1` sample entry wrapping an `avcC` (AVCDecoderConfigurationRecord) body.
fn build_avc1(width: u16, height: u16, avcc: &[u8]) -> Vec<u8> {
    visual_sample_entry(b"avc1", width, height, &bx(b"avcC", avcc))
}

/// `av01` sample entry wrapping an `av1C` (AV1CodecConfigurationRecord) body.
#[cfg(feature = "codec-av1")]
fn build_av01(width: u16, height: u16, av1c: &[u8]) -> Vec<u8> {
    visual_sample_entry(b"av01", width, height, &bx(b"av1C", av1c))
}

/// `vvc1` sample entry wrapping a `vvcC` (VvcDecoderConfigurationRecord) body.
#[cfg(feature = "codec-vvc")]
fn build_vvc1(width: u16, height: u16, vvcc: &[u8]) -> Vec<u8> {
    visual_sample_entry(b"vvc1", width, height, &bx(b"vvcC", vvcc))
}

/// `hvc1` sample entry wrapping an `hvcC` (HEVCDecoderConfigurationRecord) body.
#[cfg(feature = "codec-h265")]
fn build_hvc1(width: u16, height: u16, hvcc: &[u8]) -> Vec<u8> {
    visual_sample_entry(b"hvc1", width, height, &bx(b"hvcC", hvcc))
}

fn build_mvex() -> Vec<u8> {
    let mut trex = Vec::new();
    trex.extend_from_slice(&TRACK_ID.to_be_bytes());
    trex.extend_from_slice(&1u32.to_be_bytes()); // default_sample_description_index
    trex.extend_from_slice(&0u32.to_be_bytes()); // default_sample_duration
    trex.extend_from_slice(&0u32.to_be_bytes()); // default_sample_size
    trex.extend_from_slice(&0u32.to_be_bytes()); // default_sample_flags
    bx(b"mvex", &full(b"trex", 0, 0, &trex))
}

/// Sync sample (keyframe) flags vs. a non-sync delta sample.
fn sample_flags(is_key: bool) -> u32 {
    if is_key {
        0x0200_0000 // sample_depends_on = 2 (I-frame), is_non_sync = 0
    } else {
        0x0101_0000 // sample_depends_on = 1, sample_is_non_sync_sample = 1
    }
}

fn build_moof(seq: u32, base_decode_time: u64, samples: &[Sample], durations: &[u64]) -> Vec<u8> {
    // mfhd
    let mfhd = full(b"mfhd", 0, 0, &seq.to_be_bytes());

    // tfhd: default-base-is-moof so data_offset is relative to the moof start.
    let tfhd = full(b"tfhd", 0, 0x0002_0000, &TRACK_ID.to_be_bytes());

    // tfdt v1: base media decode time.
    let tfdt = full(b"tfdt", 1, 0, &base_decode_time.to_be_bytes());

    // trun v1: per-sample duration, size, flags, signed composition offset.
    let trun_flags: u32 = 0x0001 | 0x0100 | 0x0200 | 0x0400 | 0x0800;
    let mut trun_body = Vec::new();
    trun_body.extend_from_slice(&(samples.len() as u32).to_be_bytes()); // sample_count
    let data_offset_pos_in_body = trun_body.len();
    trun_body.extend_from_slice(&0i32.to_be_bytes()); // data_offset (patched below)
    for (s, &dur) in samples.iter().zip(durations) {
        trun_body.extend_from_slice(&(dur as u32).to_be_bytes());
        trun_body.extend_from_slice(&(s.data.len() as u32).to_be_bytes());
        trun_body.extend_from_slice(&sample_flags(s.is_key).to_be_bytes());
        trun_body.extend_from_slice(&s.cts.to_be_bytes());
    }
    let trun = full(b"trun", 1, trun_flags, &trun_body);

    // traf = tfhd + tfdt + trun
    let mut traf_body = Vec::new();
    traf_body.extend_from_slice(&tfhd);
    traf_body.extend_from_slice(&tfdt);
    let trun_pos_in_traf_body = traf_body.len();
    traf_body.extend_from_slice(&trun);
    let traf = bx(b"traf", &traf_body);

    // moof = mfhd + traf
    let mut moof_body = Vec::new();
    moof_body.extend_from_slice(&mfhd);
    let traf_pos_in_moof_body = moof_body.len();
    moof_body.extend_from_slice(&traf);
    let mut moof = bx(b"moof", &moof_body);

    // Patch trun.data_offset = moof size + mdat header (8) so it points at the
    // first sample byte in the following mdat.
    // Absolute index of the data_offset field within `moof`:
    //   moof header(8) + traf_pos_in_moof_body + traf header(8)
    //     + trun_pos_in_traf_body + trun header+fullbox(12) + data_offset_pos_in_body
    let off = 8 + traf_pos_in_moof_body + 8 + trun_pos_in_traf_body + 12 + data_offset_pos_in_body;
    let data_offset = (moof.len() + 8) as i32;
    moof[off..off + 4].copy_from_slice(&data_offset.to_be_bytes());
    moof
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{CodecId, FrameFlags, MediaFrame};
    use bytes::Bytes;

    // A real baseline SPS (profile 66, level 31) + a minimal PPS, Annex-B framed.
    fn h264_config() -> Vec<u8> {
        let sps = [0x67u8, 0x42, 0x00, 0x1F, 0xF4, 0x02, 0x80, 0x2D, 0x80];
        let pps = [0x68u8, 0xCE, 0x3C, 0x80];
        let mut v = vec![0, 0, 0, 1];
        v.extend_from_slice(&sps);
        v.extend_from_slice(&[0, 0, 0, 1]);
        v.extend_from_slice(&pps);
        v
    }

    fn annexb_frame(pts: i64, is_key: bool) -> MediaFrame {
        let nal_type = if is_key { 0x65 } else { 0x41 }; // IDR vs non-IDR slice
        let data = Bytes::from(vec![0, 0, 0, 1, nal_type, 0xAA, 0xBB]);
        MediaFrame::new_video(pts, pts, data, CodecId::H264, is_key)
    }

    /// Walk the top-level boxes and assert their sizes tile the buffer exactly.
    fn box_types(buf: &[u8]) -> Vec<String> {
        let mut types = Vec::new();
        let mut i = 0;
        while i + 8 <= buf.len() {
            let size = u32::from_be_bytes(buf[i..i + 4].try_into().unwrap()) as usize;
            let typ = String::from_utf8_lossy(&buf[i + 4..i + 8]).to_string();
            types.push(typ);
            assert!(size >= 8 && i + size <= buf.len(), "box size out of range");
            i += size;
        }
        assert_eq!(i, buf.len(), "boxes must tile the buffer exactly");
        types
    }

    #[test]
    fn init_segment_has_ftyp_moov_and_codec_string() {
        let mut m = Fmp4Muxer::new();
        let init = m
            .init_segment(CodecId::H264, &h264_config())
            .unwrap()
            .expect("init segment");
        assert_eq!(box_types(&init), vec!["ftyp", "moov"]);
        // Sample entry + decoder config present.
        assert!(find(&init, b"avc1"));
        assert!(find(&init, b"avcC"));
        assert!(find(&init, b"mvex"));
        assert_eq!(m.codec_string().as_deref(), Some("avc1.42001F"));
    }

    #[test]
    fn unsupported_codec_init_is_rejected() {
        let mut m = Fmp4Muxer::new();
        // VP9 has no fMP4 init-segment support in this release.
        assert!(m.init_segment(CodecId::VP9, &[0, 0]).is_err());
    }

    #[cfg(feature = "codec-h265")]
    #[test]
    fn h265_init_segment_has_hvc1_and_hvcc() {
        use crate::codec::testutil::BitWriter;

        // Minimal HEVC SPS for 1920x1088 (Main, level 4.0), mirroring the codec
        // module's fixture, plus a VPS and PPS so the hvcC carries all arrays.
        let mut w = BitWriter::default();
        w.bits(0, 4); // sps_video_parameter_set_id
        w.bits(0, 3); // sps_max_sub_layers_minus1
        w.bit(0); // sps_temporal_id_nesting_flag
        w.bits(0, 2); // general_profile_space
        w.bit(0); // general_tier_flag
        w.bits(1, 5); // general_profile_idc = Main
        w.bits(0, 32); // compatibility flags
        w.bits(0, 32); // constraint (hi)
        w.bits(0, 16); // constraint (lo)
        w.bits(120, 8); // general_level_idc
        w.ue(0); // sps_seq_parameter_set_id
        w.ue(1); // chroma_format_idc
        w.ue(1920); // pic_width
        w.ue(1080); // pic_height
        w.bit(0); // conformance_window_flag
        let mut sps = vec![0x42u8, 0x01];
        sps.extend_from_slice(&w.bytes());

        let mut config = vec![0, 0, 0, 1, 0x40, 0x01, 0xAA]; // VPS
        config.extend_from_slice(&[0, 0, 0, 1]);
        config.extend_from_slice(&sps);
        config.extend_from_slice(&[0, 0, 0, 1, 0x44, 0x01, 0xBB]); // PPS

        let mut m = Fmp4Muxer::new();
        let init = m
            .init_segment(CodecId::H265, &config)
            .unwrap()
            .expect("init segment");
        assert_eq!(box_types(&init), vec!["ftyp", "moov"]);
        assert!(find(&init, b"hvc1"), "hvc1 sample entry present");
        assert!(find(&init, b"hvcC"), "hvcC decoder config present");
        assert!(
            m.codec_string()
                .as_deref()
                .is_some_and(|s| s.starts_with("hvc1.")),
            "HEVC codec string"
        );
    }

    #[cfg(feature = "codec-av1")]
    #[test]
    fn av1_init_segment_has_av01_and_av1c() {
        use crate::codec::testutil::BitWriter;

        // A sequence-header OBU for 1920x1080, profile 0 (mirrors the AV1 parser
        // test), wrapped as a size-delimited temporal unit.
        let mut w = BitWriter::default();
        w.bits(0, 3); // seq_profile
        w.bit(0); // still_picture
        w.bit(0); // reduced_still_picture_header
        w.bit(0); // timing_info_present_flag
        w.bit(0); // initial_display_delay_present_flag
        w.bits(0, 5); // operating_points_cnt_minus1
        w.bits(0, 12); // operating_point_idc[0]
        w.bits(1, 5); // seq_level_idx[0]
        w.bits(11, 4); // frame_width_bits_minus_1
        w.bits(11, 4); // frame_height_bits_minus_1
        w.bits(1919, 12);
        w.bits(1079, 12);
        w.align();
        let payload = w.bytes();
        let mut config = vec![0x0A, payload.len() as u8];
        config.extend_from_slice(&payload);

        let mut m = Fmp4Muxer::new();
        let init = m
            .init_segment(CodecId::AV1, &config)
            .unwrap()
            .expect("av1 init segment");
        assert_eq!(box_types(&init), vec!["ftyp", "moov"]);
        assert!(find(&init, b"av01"), "av01 sample entry");
        assert!(find(&init, b"av1C"), "av1C config box");
        assert_eq!(m.codec_string().as_deref(), Some("av01.0.01M.08"));

        // AV1 temporal units are muxed verbatim (no Annex-B → AVCC rewrite).
        m.start_segment().unwrap();
        let mut frame = MediaFrame::new_video(
            0,
            0,
            Bytes::from(vec![0x32, 0x02, 0xAA, 0xBB]), // OBU_FRAME
            CodecId::AV1,
            true,
        );
        frame.duration = Some(33);
        m.write(&frame).unwrap();
        let frag = m.finish_segment().unwrap();
        assert_eq!(box_types(&frag), vec!["moof", "mdat"]);
        // mdat carries the OBU bytes unchanged.
        assert!(find(&frag, &[0x32, 0x02, 0xAA, 0xBB]));
    }

    #[test]
    fn fragment_has_moof_mdat_and_correct_sample_count() {
        let mut m = Fmp4Muxer::new();
        m.init_segment(CodecId::H264, &h264_config()).unwrap();
        m.start_segment().unwrap();
        m.write(&annexb_frame(0, true)).unwrap();
        m.write(&annexb_frame(40, false)).unwrap();
        m.write(&annexb_frame(80, false)).unwrap();
        // An audio frame is skipped, not muxed into the video track.
        let audio = MediaFrame::new_audio(80, Bytes::from_static(b"aac"), CodecId::AAC);
        m.write(&audio).unwrap();

        let frag = m.finish_segment().unwrap();
        assert_eq!(box_types(&frag), vec!["moof", "mdat"]);
        assert!(find(&frag, b"tfdt"));
        assert!(find(&frag, b"trun"));

        // sample_count is the u32 right after the trun full-box header.
        let trun_at = position(&frag, b"trun").expect("trun present");
        let count = u32::from_be_bytes(frag[trun_at + 8..trun_at + 12].try_into().unwrap());
        assert_eq!(count, 3, "three video samples, audio skipped");
    }

    #[test]
    fn base_decode_time_advances_across_fragments() {
        let mut m = Fmp4Muxer::new();
        m.init_segment(CodecId::H264, &h264_config()).unwrap();

        m.start_segment().unwrap();
        m.write(&annexb_frame(0, true)).unwrap();
        m.write(&annexb_frame(40, false)).unwrap();
        let _ = m.finish_segment().unwrap();
        // Sample 0 = 40ms (DTS delta); last sample falls back to the prior delta
        // (40ms) → fragment total 80ms.
        assert_eq!(m.base_decode_time, 80);

        m.start_segment().unwrap();
        m.write(&annexb_frame(80, true)).unwrap();
        let frag2 = m.finish_segment().unwrap();
        // tfdt in fragment 2 carries the advanced base decode time.
        let tfdt_at = position(&frag2, b"tfdt").unwrap();
        let base = u64::from_be_bytes(frag2[tfdt_at + 8..tfdt_at + 16].try_into().unwrap());
        assert_eq!(base, 80);
    }

    #[cfg(feature = "codec-vvc")]
    #[test]
    fn vvc_init_segment_has_vvc1_and_vvcc() {
        use crate::codec::testutil::BitWriter;

        // VVC SPS for 1920x1080, 4:2:0, PTL block absent (mirrors the VVC test).
        let mut w = BitWriter::default();
        w.bits(0, 4); // sps_seq_parameter_set_id
        w.bits(0, 4); // sps_video_parameter_set_id
        w.bits(0, 3); // sps_max_sublayers_minus1
        w.bits(1, 2); // chroma_format_idc = 4:2:0
        w.bits(0, 2); // sps_log2_ctu_size_minus5
        w.bit(0); // sps_ptl_dpb_hrd_params_present_flag = 0
        w.bit(0); // sps_gdr_enabled_flag
        w.bit(0); // sps_ref_pic_resampling_enabled_flag
        w.ue(1920);
        w.ue(1080);
        w.bit(0); // sps_conformance_window_flag
        let mut sps = vec![0x00u8, 0x79]; // NAL header: nuh_unit_type 15 (SPS)
        sps.extend_from_slice(&w.bytes());

        let mut config = vec![0, 0, 0, 1];
        config.extend_from_slice(&sps);

        let mut m = Fmp4Muxer::new();
        let init = m
            .init_segment(CodecId::VVC, &config)
            .unwrap()
            .expect("vvc init segment");
        assert_eq!(box_types(&init), vec!["ftyp", "moov"]);
        assert!(find(&init, b"vvc1"), "vvc1 sample entry");
        assert!(find(&init, b"vvcC"), "vvcC config box");
        assert_eq!(m.codec_string().as_deref(), Some("vvc1.0.L0"));

        // VVC samples are length-prefixed (4-byte) like H.264.
        m.start_segment().unwrap();
        let mut frame = MediaFrame::new_video(
            0,
            0,
            Bytes::from(vec![0, 0, 0, 1, 0x00, 0x39, 0xAA]), // IDR_W_RADL Annex-B
            CodecId::VVC,
            true,
        );
        frame.duration = Some(40);
        m.write(&frame).unwrap();
        let frag = m.finish_segment().unwrap();
        assert_eq!(box_types(&frag), vec!["moof", "mdat"]);
        // The NAL was rewritten to a 4-byte length prefix (len 3 → 00 00 00 03).
        assert!(find(&frag, &[0x00, 0x00, 0x00, 0x03, 0x00, 0x39, 0xAA]));
    }

    #[test]
    fn config_without_sps_errors() {
        let mut m = Fmp4Muxer::new();
        // PPS only — no SPS.
        let cfg = vec![0, 0, 0, 1, 0x68, 0xCE, 0x3C, 0x80];
        assert!(m.init_segment(CodecId::H264, &cfg).is_err());
    }

    #[test]
    fn empty_fragment_yields_no_bytes() {
        let mut m = Fmp4Muxer::new();
        m.start_segment().unwrap();
        assert!(m.finish_segment().unwrap().is_empty());
    }

    fn find(buf: &[u8], needle: &[u8]) -> bool {
        position(buf, needle).is_some()
    }
    fn position(buf: &[u8], needle: &[u8]) -> Option<usize> {
        buf.windows(needle.len()).position(|w| w == needle)
    }

    #[test]
    fn integrates_with_segmenter_via_ext_x_map() {
        use crate::packager::{HlsSegmenter, Packager};
        use crate::testing::InMemoryStorage;
        use crate::traits::StorageBackend;

        tokio_test_block(async {
            let store = InMemoryStorage::new();
            let mut seg = HlsSegmenter::new(Fmp4Muxer::new(), store.clone(), "live/fmp4", 2, 5);

            let mut cfg =
                MediaFrame::new_video(0, 0, Bytes::from(h264_config()), CodecId::H264, true);
            cfg.flags |= FrameFlags::CONFIG;
            seg.push(&cfg).await.unwrap();
            for i in 0..6 {
                seg.push(&annexb_frame(i * 1000, true)).await.unwrap();
            }
            seg.finish().await.unwrap();

            assert!(store.get("live/fmp4/init.m4s").await.is_ok());
            let pl = String::from_utf8(store.get("live/fmp4/index.m3u8").await.unwrap().to_vec())
                .unwrap();
            assert!(pl.contains("#EXT-X-MAP:URI=\"init.m4s\""));
        });
    }

    fn tokio_test_block<F: std::future::Future>(fut: F) -> F::Output {
        tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap()
            .block_on(fut)
    }
}