arcly-stream 0.1.7

//! Egress packaging: turn a live frame stream into HLS/LL-HLS segments.
//!
//! Gated behind `hls`. Provides the [`Muxer`] and [`Packager`] contracts, a
//! sliding-window [`HlsPlaylist`] generator, and an [`HlsSegmenter`] that cuts
//! segments on keyframe boundaries and writes both media segments and the
//! `.m3u8` playlist through any [`StorageBackend`].
//!
//! The container byte-format is pluggable via [`Muxer`]: a [`PassthroughMuxer`]
//! (elementary-stream concatenation) ships here; production MPEG-TS or fMP4/CMAF
//! muxers implement the same trait and drop in unchanged.
//!
//! ```no_run
//! # #[cfg(feature = "storage-fs")]
//! # async fn demo(handle: arcly_stream::StreamHandle) -> arcly_stream::Result<()> {
//! use arcly_stream::packager::{HlsSegmenter, Packager, PassthroughMuxer};
//! use arcly_stream::storage::FsStorage;
//!
//! let storage = FsStorage::new("/var/hls");
//! let mut seg = HlsSegmenter::new(PassthroughMuxer::new("ts"), storage, "live/cam", 6, 5);
//! let mut sub = handle.subscribe_resilient();
//! while let Some(frame) = sub.recv().await {
//!     seg.push(&frame).await?;
//! }
//! seg.finish().await?;
//! # Ok(())
//! # }
//! ```

mod engine;
mod playlist;

use engine::{SegmentEngine, BANDWIDTH_HINT};

#[cfg(feature = "mpegts")]
#[cfg_attr(docsrs, doc(cfg(feature = "mpegts")))]
mod mpegts;

#[cfg(feature = "mpegts")]
#[cfg_attr(docsrs, doc(cfg(feature = "mpegts")))]
pub use mpegts::MpegTsMuxer;

#[cfg(feature = "fmp4")]
#[cfg_attr(docsrs, doc(cfg(feature = "fmp4")))]
mod fmp4;

#[cfg(feature = "fmp4")]
#[cfg_attr(docsrs, doc(cfg(feature = "fmp4")))]
pub use fmp4::Fmp4Muxer;

pub use playlist::{render_master, HlsPlaylist, Part, Segment};

#[cfg(feature = "dash")]
#[cfg_attr(docsrs, doc(cfg(feature = "dash")))]
mod dash;

#[cfg(feature = "dash")]
#[cfg_attr(docsrs, doc(cfg(feature = "dash")))]
pub use dash::{DashManifest, DashPackager};

use crate::traits::StorageBackend;
use crate::{CodecId, FrameFlags, MediaFrame, Result};
use async_trait::async_trait;
use bytes::{BufMut, Bytes, BytesMut};

/// Container muxer: accumulates frames into one segment's bytes.
///
/// Implement this for MPEG-TS, fMP4/CMAF, etc. The segmenter calls
/// [`start_segment`](Self::start_segment) at each boundary, [`write`](Self::write)
/// per frame, and [`finish_segment`](Self::finish_segment) to take the bytes.
pub trait Muxer: Send {
    /// File extension for produced segments (e.g. `"ts"`, `"m4s"`).
    fn extension(&self) -> &'static str;

    /// Begin a fresh segment, discarding any partial state.
    fn start_segment(&mut self) -> Result<()>;

    /// Append one frame to the current segment.
    fn write(&mut self, frame: &MediaFrame) -> Result<()>;

    /// Take the finished segment's bytes, resetting for the next one.
    fn finish_segment(&mut self) -> Result<Bytes>;

    /// Flush the frames buffered since the last part/segment start as an
    /// independently-appendable **partial segment** (LL-HLS `#EXT-X-PART`),
    /// *without* ending the current segment. Returns `None` when nothing is
    /// buffered or the container doesn't support parts.
    ///
    /// The concatenation of a segment's parts must equal the full segment, so a
    /// container's parts are self-contained chunks (fMP4: one `moof`+`mdat`
    /// fragment each). The default returns `None` (no LL-HLS parts).
    fn take_partial(&mut self) -> Result<Option<Bytes>> {
        Ok(None)
    }

    /// Build an fMP4 initialization segment (`ftyp`+`moov` carrying the
    /// `avcC`/`hvcC`/`av1C`/`vvcC` decoder config) from a CONFIG access unit.
    ///
    /// Returning `Some` makes [`HlsSegmenter`] write it once and reference it via
    /// `#EXT-X-MAP`. The default returns `None` — correct for self-initializing
    /// containers like MPEG-TS and the [`PassthroughMuxer`].
    fn init_segment(&mut self, _codec: CodecId, _config_record: &[u8]) -> Result<Option<Bytes>> {
        Ok(None)
    }

    /// Build the initialization segment from **all** the stream's CONFIG access
    /// units seen before the first media sample — e.g. an `(H264, avc)` video
    /// config plus an `(AAC, asc)` audio config — so a container that carries
    /// multiple tracks (fMP4 with muxed audio) can emit a multi-track `moov`.
    ///
    /// The default delegates to [`init_segment`](Self::init_segment) with the
    /// first config, so single-track and self-initializing muxers (MPEG-TS,
    /// passthrough) need no changes. The [`Fmp4Muxer`] overrides this
    /// to add an audio track when an AAC config is present.
    fn build_init_from(&mut self, configs: &[(CodecId, Bytes)]) -> Result<Option<Bytes>> {
        match configs.first() {
            Some((codec, data)) => self.init_segment(*codec, data),
            None => Ok(None),
        }
    }

    /// The HLS `CODECS` attribute for this rendition's master-playlist entry
    /// (e.g. `"hvc1.1.6.L120.B0"`), once known.
    fn codec_string(&self) -> Option<String> {
        None
    }
}

/// A trivial muxer that concatenates frame payloads (elementary stream).
///
/// Useful for testing and raw recording; not a real container. Swap in a TS or
/// fMP4 muxer for player-compatible output.
pub struct PassthroughMuxer {
    ext: &'static str,
    buf: BytesMut,
}

impl PassthroughMuxer {
    /// New passthrough muxer producing segments with the given extension.
    pub fn new(ext: &'static str) -> Self {
        Self {
            ext,
            buf: BytesMut::new(),
        }
    }
}

impl Muxer for PassthroughMuxer {
    fn extension(&self) -> &'static str {
        self.ext
    }
    fn start_segment(&mut self) -> Result<()> {
        self.buf.clear();
        Ok(())
    }
    fn write(&mut self, frame: &MediaFrame) -> Result<()> {
        self.buf.put_slice(&frame.data);
        Ok(())
    }
    fn finish_segment(&mut self) -> Result<Bytes> {
        Ok(std::mem::take(&mut self.buf).freeze())
    }
    fn take_partial(&mut self) -> Result<Option<Bytes>> {
        if self.buf.is_empty() {
            return Ok(None);
        }
        Ok(Some(std::mem::take(&mut self.buf).freeze()))
    }
}

/// Consumes a live frame stream and produces a packaged rendition.
#[async_trait]
pub trait Packager: Send {
    /// Feed one frame; may finalize and emit a segment as a side effect.
    async fn push(&mut self, frame: &MediaFrame) -> Result<()>;
    /// Flush the final segment and close the rendition.
    async fn finish(&mut self) -> Result<()>;
}

/// Keyframe-boundary HLS segmenter writing through a [`StorageBackend`].
///
/// Starts at the first keyframe, cuts a new segment at the first keyframe at or
/// after `target_duration` seconds, and after each cut writes the media segment
/// and the regenerated `index.m3u8` to storage under `prefix`.
pub struct HlsSegmenter<M: Muxer, S: StorageBackend> {
    engine: SegmentEngine<M, S>,
    playlist: HlsPlaylist,
    /// Whether the multivariant (master) playlist has been written yet. It is
    /// emitted once, as soon as the rendition's `CODECS` string is known.
    master_written: bool,
    /// The HLS `CODECS` string, derived from the CONFIG access unit. Cached here
    /// (not just read from the muxer) because the segment clock skips the CONFIG
    /// frame before it reaches a self-initializing muxer like MPEG-TS, so the
    /// muxer would never see it — yet the master playlist still needs it.
    codec_string: Option<String>,
    // ── LL-HLS partial-segment state (inactive when `part_target_ms == 0`) ──
    /// Part target in milliseconds; `0` disables LL-HLS part emission.
    part_target_ms: i64,
    /// Index of the next part within the in-progress segment.
    part_idx: u64,
    /// PTS (ms) at which the current part began, once a segment is open.
    part_anchor_pts: Option<i64>,
    /// PTS (ms) of the most recently written frame (for the final part's length).
    last_pts: i64,
    /// Byte chunks of the in-progress segment's parts, concatenated into the full
    /// segment file at the cut so non-LL clients can fetch the whole segment.
    seg_part_bytes: Vec<Bytes>,
}

impl<M: Muxer, S: StorageBackend> HlsSegmenter<M, S> {
    /// New segmenter writing under `prefix`, targeting `target_duration`-second
    /// segments and a `window`-segment live playlist.
    pub fn new(
        muxer: M,
        storage: S,
        prefix: impl Into<String>,
        target_duration: u64,
        window: usize,
    ) -> Self {
        Self {
            engine: SegmentEngine::new(muxer, storage, prefix, target_duration),
            playlist: HlsPlaylist::new(target_duration, window),
            master_written: false,
            codec_string: None,
            part_target_ms: 0,
            part_idx: 0,
            part_anchor_pts: None,
            last_pts: 0,
            seg_part_bytes: Vec::new(),
        }
    }

    /// The HLS `CODECS` attribute for this rendition's master-playlist entry.
    pub fn codec_string(&self) -> Option<String> {
        self.codec_string
            .clone()
            .or_else(|| self.engine.muxer.codec_string())
    }

    /// On the first CONFIG access unit: write the muxer's fMP4 init segment (if
    /// any) and reference it via `#EXT-X-MAP`, and derive the HLS `CODECS` string
    /// from the parameter sets. The codec string is read straight from the CONFIG
    /// here — not from the muxer — because the segment clock skips the CONFIG
    /// frame before a self-initializing muxer (MPEG-TS) ever sees it.
    async fn ensure_init_segment(&mut self, frame: &MediaFrame) -> Result<()> {
        if self.codec_string.is_none()
            && frame.is_video()
            && frame.flags.contains(FrameFlags::CONFIG)
        {
            self.codec_string = crate::codec::dispatch::parse_config(frame.codec, &frame.data)
                .and_then(|p| crate::codec::dispatch::hls_codec_string(frame.codec, &p));
        }
        if let Some(uri) = self.engine.ensure_init(frame).await? {
            self.playlist.set_map(uri);
        }
        Ok(())
    }

    /// Enable LL-HLS output: the playlist advertises parts and the segmenter
    /// emits a partial segment roughly every `part_target` seconds (writing the
    /// part files and `#EXT-X-PART`/preload-hint lines) in addition to full
    /// segments. Requires a muxer that implements
    /// [`take_partial`](Muxer::take_partial) (fMP4 / passthrough).
    pub fn low_latency(mut self, part_target: f64) -> Self {
        self.playlist = self.playlist.low_latency(part_target);
        self.part_target_ms = (part_target.max(0.05) * 1000.0) as i64;
        self
    }

    /// The storage key of the media playlist.
    pub fn playlist_key(&self) -> String {
        format!("{}/index.m3u8", self.engine.prefix)
    }

    /// The storage key of the multivariant (master) playlist.
    pub fn master_key(&self) -> String {
        format!("{}/master.m3u8", self.engine.prefix)
    }

    /// Once the muxer can report a `CODECS` string, write the multivariant
    /// (master) playlist a single time. This is what advertises `CODECS="hvc1.*"`
    /// for HEVC (and `av01.*`/`vvc1.*`); without it players such as Safari will
    /// not decode H.265, falling back as if the rendition were unplayable.
    async fn ensure_master_playlist(&mut self) -> Result<()> {
        if self.master_written {
            return Ok(());
        }
        let Some(codecs) = self.codec_string() else {
            return Ok(());
        };
        let body = playlist::render_master("index.m3u8", Some(&codecs), BANDWIDTH_HINT);
        self.engine
            .storage
            .put(
                &self.engine.key("master.m3u8"),
                Bytes::from(body.into_bytes()),
            )
            .await?;
        self.master_written = true;
        Ok(())
    }

    /// Flush the buffered samples as one LL-HLS partial segment: write the part
    /// file, append `#EXT-X-PART` + a preload hint for the next part, and stash
    /// the bytes so the full segment file can be assembled at the cut.
    async fn flush_part(&mut self, end_pts: i64) -> Result<()> {
        let Some(bytes) = self.engine.muxer.take_partial()? else {
            return Ok(());
        };
        let anchor = self.part_anchor_pts.unwrap_or(end_pts);
        let duration = (end_pts - anchor).max(0) as f64 / 1000.0;
        let uri = self.engine.part_uri(self.engine.seq, self.part_idx);
        let key = self.engine.key(&uri);
        self.engine.storage.put(&key, bytes.clone()).await?;
        self.seg_part_bytes.push(bytes);
        self.playlist.add_pending_part(Part {
            uri,
            duration,
            // The first part of a segment begins on the anchoring keyframe.
            independent: self.part_idx == 0,
        });
        self.playlist
            .set_preload_hint(self.engine.part_uri(self.engine.seq, self.part_idx + 1));
        self.part_idx += 1;
        self.part_anchor_pts = Some(end_pts);
        self.write_playlist().await
    }

    async fn cut(&mut self, duration: f64) -> Result<()> {
        let uri = self.engine.segment_uri(self.engine.seq);

        if self.part_target_ms > 0 {
            // LL-HLS: flush the segment's final part, then assemble the full
            // segment file from its parts (concatenation == the whole segment).
            self.flush_part(self.last_pts).await?;
            let mut full = BytesMut::new();
            for b in &self.seg_part_bytes {
                full.put_slice(b);
            }
            // Defensive: if the muxer produced no parts (it doesn't implement
            // `take_partial`), fall back to a whole-segment flush so LL-HLS still
            // emits a non-empty segment instead of a 0-byte file.
            let full = if full.is_empty() {
                self.engine.muxer.finish_segment()?
            } else {
                full.freeze()
            };
            let key = self.engine.key(&uri);
            self.engine.storage.put(&key, full).await?;
            self.playlist.commit_segment(Segment {
                seq: self.engine.seq,
                duration,
                uri,
                discontinuity: false,
                parts: Vec::new(), // filled from pending parts by commit_segment
            });
            self.seg_part_bytes.clear();
            self.part_idx = 0;
            self.part_anchor_pts = None;
        } else {
            let bytes = self.engine.muxer.finish_segment()?;
            let key = self.engine.key(&uri);
            self.engine.storage.put(&key, bytes).await?;
            self.playlist.push(Segment {
                seq: self.engine.seq,
                duration,
                uri,
                discontinuity: false,
                parts: Vec::new(),
            });
        }
        self.write_playlist().await?;
        self.engine.seq += 1;
        Ok(())
    }

    async fn write_playlist(&mut self) -> Result<()> {
        let body = self.playlist.render();
        self.engine
            .storage
            .put(
                &self.engine.key("index.m3u8"),
                Bytes::from(body.into_bytes()),
            )
            .await
    }
}

#[async_trait]
impl<M: Muxer, S: StorageBackend> Packager for HlsSegmenter<M, S> {
    async fn push(&mut self, frame: &MediaFrame) -> Result<()> {
        self.ensure_init_segment(frame).await?;
        // Emit the master playlist as soon as the codec string is known — which
        // `ensure_init_segment` derives from the CONFIG frame. This must run
        // before the clock-skip below, since the clock skips that very CONFIG
        // frame (it is not a keyframe), and HEVC/AV1 are unviewable without it.
        self.ensure_master_playlist().await?;
        let decision = self.engine.observe(frame);
        if decision.skip {
            return Ok(());
        }
        if let Some(duration) = decision.cut_previous {
            self.cut(duration).await?;
        }
        if decision.open_new {
            self.engine.muxer.start_segment()?;
            // A fresh segment anchors a new part run at this keyframe.
            self.part_anchor_pts = Some(frame.pts);
            self.part_idx = 0;
            self.seg_part_bytes.clear();
        }
        self.engine.muxer.write(frame)?;
        self.last_pts = frame.pts;

        // LL-HLS: emit a partial segment once the part target has elapsed. Parts
        // are cut at frame boundaries, so a part never splits an access unit.
        if self.part_target_ms > 0 {
            if let Some(anchor) = self.part_anchor_pts {
                if frame.pts - anchor >= self.part_target_ms {
                    self.flush_part(frame.pts).await?;
                }
            }
        }
        Ok(())
    }

    async fn finish(&mut self) -> Result<()> {
        if let Some(duration) = self.engine.flush() {
            self.cut(duration).await?;
        }
        self.playlist.finish();
        self.write_playlist().await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::testing::{video_frame, InMemoryStorage};

    #[tokio::test]
    async fn segments_on_keyframe_after_target_and_writes_playlist() {
        let store = InMemoryStorage::new();
        let mut seg =
            HlsSegmenter::new(PassthroughMuxer::new("ts"), store.clone(), "live/cam", 2, 5);

        // Keyframes every 1s; with a 2s target, a cut happens at the keyframe
        // at/after 2s. pts in ms.
        for i in 0..5 {
            let pts = i * 1000;
            seg.push(&video_frame(pts, true)).await.unwrap();
            // a delta in between
            seg.push(&video_frame(pts + 500, false)).await.unwrap();
        }
        seg.finish().await.unwrap();

        let playlist = store.get("live/cam/index.m3u8").await.unwrap();
        let text = String::from_utf8(playlist.to_vec()).unwrap();
        assert!(text.contains("#EXTM3U"));
        assert!(text.contains("#EXT-X-ENDLIST"));
        // At least the first full segment must have been written to storage.
        assert!(store.get("live/cam/seg0.ts").await.is_ok());
        // Segment payload is the concatenation of its frames' bytes.
        assert!(!store.get("live/cam/seg0.ts").await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn ll_hls_emits_part_files_and_ext_x_part() {
        let store = InMemoryStorage::new();
        // 2s segments, parts every ~0.5s.
        let mut seg =
            HlsSegmenter::new(PassthroughMuxer::new("m4s"), store.clone(), "live/ll", 2, 5)
                .low_latency(0.5);

        // Keyframe at 0, deltas every 250ms; next keyframe at 2000 cuts segment 0.
        seg.push(&video_frame(0, true)).await.unwrap();
        for i in 1..8 {
            seg.push(&video_frame(i * 250, false)).await.unwrap();
        }
        seg.push(&video_frame(2000, true)).await.unwrap(); // cut seg 0
        seg.push(&video_frame(2250, false)).await.unwrap();
        seg.finish().await.unwrap();

        let text =
            String::from_utf8(store.get("live/ll/index.m3u8").await.unwrap().to_vec()).unwrap();
        // LL-HLS headers + actual part lines for segment 0.
        assert!(text.contains("#EXT-X-PART-INF:PART-TARGET=0.500"));
        assert!(
            text.contains("#EXT-X-PART:DURATION="),
            "renders EXT-X-PART lines"
        );
        assert!(
            text.contains("INDEPENDENT=YES"),
            "first part is independent"
        );
        // Part files were written, and the first full segment too.
        assert!(
            store.get("live/ll/seg0.0.m4s").await.is_ok(),
            "part 0 written"
        );
        assert!(
            store.get("live/ll/seg0.1.m4s").await.is_ok(),
            "part 1 written"
        );
        assert!(
            store.get("live/ll/seg0.m4s").await.is_ok(),
            "full segment 0 written"
        );

        // The full segment equals the concatenation of its parts.
        let mut concat = Vec::new();
        let mut p = 0;
        while let Ok(part) = store.get(&format!("live/ll/seg0.{p}.m4s")).await {
            concat.extend_from_slice(&part);
            p += 1;
        }
        assert!(p >= 2, "segment 0 had multiple parts");
        let full = store.get("live/ll/seg0.m4s").await.unwrap();
        assert_eq!(&concat[..], &full[..], "full segment == concat of parts");
    }

    #[tokio::test]
    async fn ll_hls_falls_back_to_full_segment_when_muxer_has_no_parts() {
        // InitMuxer does NOT implement `take_partial`; in LL-HLS mode the cut path
        // must still write a non-empty segment (regression: 0-byte .ts segments).
        let store = InMemoryStorage::new();
        let mut seg = HlsSegmenter::new(
            InitMuxer {
                buf: BytesMut::new(),
            },
            store.clone(),
            "live/nopart",
            2,
            5,
        )
        .low_latency(0.5);

        seg.push(&video_frame(0, true)).await.unwrap();
        for i in 1..6 {
            seg.push(&video_frame(i * 1000, i % 2 == 0)).await.unwrap();
        }
        seg.finish().await.unwrap();

        let seg0 = store.get("live/nopart/seg0.m4s").await.unwrap();
        assert!(!seg0.is_empty(), "LL-HLS segment must not be empty");
    }

    #[tokio::test]
    async fn writes_master_playlist_with_hevc_codecs() {
        use crate::FrameFlags;
        let store = InMemoryStorage::new();
        let mut seg = HlsSegmenter::new(
            InitMuxer {
                buf: BytesMut::new(),
            },
            store.clone(),
            "live/hevc",
            2,
            5,
        );

        let mut cfg = video_frame(0, true);
        cfg.codec = CodecId::H265;
        cfg.flags |= FrameFlags::CONFIG;
        seg.push(&cfg).await.unwrap();
        for i in 1..4 {
            seg.push(&video_frame(i * 1000, true)).await.unwrap();
        }
        seg.finish().await.unwrap();

        let master =
            String::from_utf8(store.get("live/hevc/master.m3u8").await.unwrap().to_vec()).unwrap();
        assert!(master.contains("#EXT-X-STREAM-INF:"));
        assert!(
            master.contains("CODECS=\"hvc1.1.6.L120.B0\""),
            "master advertises HEVC codec string: {master}"
        );
        assert!(master.contains("index.m3u8"));
    }

    /// A muxer that emits an fMP4-style init segment, to exercise EXT-X-MAP.
    struct InitMuxer {
        buf: BytesMut,
    }
    impl Muxer for InitMuxer {
        fn extension(&self) -> &'static str {
            "m4s"
        }
        fn start_segment(&mut self) -> Result<()> {
            self.buf.clear();
            Ok(())
        }
        fn write(&mut self, frame: &MediaFrame) -> Result<()> {
            self.buf.put_slice(&frame.data);
            Ok(())
        }
        fn finish_segment(&mut self) -> Result<Bytes> {
            Ok(std::mem::take(&mut self.buf).freeze())
        }
        fn init_segment(&mut self, _codec: CodecId, config_record: &[u8]) -> Result<Option<Bytes>> {
            // A real muxer would wrap this in ftyp+moov; here we just echo it.
            Ok(Some(Bytes::copy_from_slice(config_record)))
        }
        fn codec_string(&self) -> Option<String> {
            Some("hvc1.1.6.L120.B0".into())
        }
    }

    #[tokio::test]
    async fn fmp4_muxer_writes_init_segment_and_ext_x_map() {
        use crate::FrameFlags;
        let store = InMemoryStorage::new();
        let mut seg = HlsSegmenter::new(
            InitMuxer {
                buf: BytesMut::new(),
            },
            store.clone(),
            "live/hevc",
            2,
            5,
        );
        assert_eq!(seg.codec_string().as_deref(), Some("hvc1.1.6.L120.B0"));

        // First frame carries CONFIG → triggers the init segment.
        let mut cfg = video_frame(0, true);
        cfg.codec = CodecId::H265;
        cfg.flags |= FrameFlags::CONFIG;
        seg.push(&cfg).await.unwrap();
        for i in 1..6 {
            seg.push(&video_frame(i * 1000, true)).await.unwrap();
        }
        seg.finish().await.unwrap();

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