arcly-stream 0.8.2

//! DASH / LL-DASH manifest (`.mpd`) generation and packaging.
//!
//! Reuses the [`Muxer`](super::Muxer) / [`Packager`](super::Packager) contracts:
//! the media segments are ordinary CMAF (fMP4) — produce them with
//! [`Fmp4Muxer`](super::Fmp4Muxer) — and [`DashManifest`] renders a dynamic
//! (live) MPD using a `SegmentTemplate` + `SegmentTimeline`, which describes
//! variable, keyframe-aligned segment durations exactly.
//!
//! Low-latency DASH is enabled with [`DashManifest::low_latency`] /
//! [`DashPackager::low_latency`]: the MPD carries `availabilityTimeOffset` so a
//! player can fetch the in-progress segment as it is produced (chunked CMAF —
//! the chunks are the same `moof`+`mdat` fragments [`Muxer::take_partial`]
//! yields; the actual chunked HTTP transfer is the delivery layer's job).

use std::collections::VecDeque;
use std::fmt::Write as _;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use super::engine::{SegmentEngine, BANDWIDTH_HINT};
use super::{ChunkSink, Muxer, Packager};
use crate::traits::StorageBackend;
use crate::{MediaFrame, Result};
use async_trait::async_trait;
use bytes::{BufMut, Bytes, BytesMut};

/// One segment's timing in the MPD `SegmentTimeline` (milliseconds).
#[derive(Debug, Clone, Copy)]
struct TimelineEntry {
    start_ms: u64,
    dur_ms: u64,
}

/// A dynamic (live) DASH media-presentation-description generator.
///
/// Holds a sliding window of segment timings and renders a spec-compliant MPD
/// with a number-addressed `SegmentTemplate` and an explicit `SegmentTimeline`.
#[derive(Debug, Clone)]
pub struct DashManifest {
    window: usize,
    timescale: u64,
    low_latency: bool,
    part_target: f64,
    availability_start: u64,                  // unix seconds
    segments: VecDeque<(u64, TimelineEntry)>, // (seq, timing)
    codecs: Option<String>,
    width: u16,
    height: u16,
    bandwidth: u32,
    init_uri: String,
    media_template: String,
    finished: bool,
}

impl DashManifest {
    /// A live manifest retaining `window` segments, timed in milliseconds.
    pub fn new(window: usize) -> Self {
        Self {
            window: window.max(1),
            timescale: 1000,
            low_latency: false,
            part_target: 0.0,
            availability_start: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .map(|d| d.as_secs())
                .unwrap_or(0),
            segments: VecDeque::new(),
            codecs: None,
            width: 0,
            height: 0,
            bandwidth: 0,
            init_uri: "init.m4s".into(),
            media_template: "seg$Number$.m4s".into(),
            finished: false,
        }
    }

    /// Enable LL-DASH: advertise `availabilityTimeOffset` of `part_target`
    /// seconds (one segment minus one part), letting players read the open
    /// segment as its chunks are produced.
    pub fn low_latency(mut self, part_target: f64) -> Self {
        self.low_latency = true;
        self.part_target = part_target.max(0.05);
        self
    }

    /// Set the codec string (`avc1.*` / `hvc1.*`), pixel dimensions, and a
    /// nominal bandwidth (bits/sec) for the `Representation`.
    pub fn set_media_info(
        &mut self,
        codecs: Option<String>,
        width: u16,
        height: u16,
        bandwidth: u32,
    ) {
        self.codecs = codecs;
        self.width = width;
        self.height = height;
        self.bandwidth = bandwidth;
    }

    /// Override the init-segment and media-segment template URIs.
    pub fn set_uris(&mut self, init_uri: impl Into<String>, media_template: impl Into<String>) {
        self.init_uri = init_uri.into();
        self.media_template = media_template.into();
    }

    /// Append a segment timing, evicting the oldest beyond the window.
    pub fn push(&mut self, seq: u64, start_ms: u64, dur_ms: u64) {
        self.segments
            .push_back((seq, TimelineEntry { start_ms, dur_ms }));
        while self.segments.len() > self.window {
            self.segments.pop_front();
        }
    }

    /// Correct the duration of the most recently pushed segment. Used in LL-DASH:
    /// a segment is advertised at *open* time with an estimated duration (so the
    /// player starts fetching it chunked immediately) and its exact duration is
    /// patched in here once the segment is cut.
    pub fn set_last_duration(&mut self, dur_ms: u64) {
        if let Some((_, t)) = self.segments.back_mut() {
            t.dur_ms = dur_ms;
        }
    }

    /// Mark the presentation static (VOD) — sets `type="static"` and a final
    /// `minimumUpdatePeriod`-free manifest.
    pub fn finish(&mut self) {
        self.finished = true;
    }

    /// The media sequence number of the first segment in the window (the MPD
    /// `startNumber`).
    fn start_number(&self) -> u64 {
        self.segments.front().map(|(seq, _)| *seq).unwrap_or(0)
    }

    /// Largest segment duration in the window (seconds), for `maxSegmentDuration`.
    fn max_seg_secs(&self) -> f64 {
        self.segments
            .iter()
            .map(|(_, t)| t.dur_ms)
            .max()
            .unwrap_or(0) as f64
            / 1000.0
    }

    /// Render the MPD document.
    pub fn render(&self) -> String {
        let mut s = String::with_capacity(512 + self.segments.len() * 48);
        s.push_str("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n");

        let mpd_type = if self.finished { "static" } else { "dynamic" };
        let profile = "urn:mpeg:dash:profile:isoff-live:2011";
        s.push_str("<MPD xmlns=\"urn:mpeg:dash:schema:mpd:2011\" ");
        // `write!` formats straight into the buffer; `format!` would allocate a
        // throwaway String per line, and this renders on every segment cut.
        let _ = write!(s, "type=\"{mpd_type}\" profiles=\"{profile}\" ");
        s.push_str("minBufferTime=\"PT1.5S\" ");
        if !self.finished {
            let _ = write!(
                s,
                "availabilityStartTime=\"{}\" ",
                unix_to_iso8601(self.availability_start)
            );
            s.push_str("minimumUpdatePeriod=\"PT1S\" ");
        }
        let _ = writeln!(s, "maxSegmentDuration=\"PT{:.3}S\">", self.max_seg_secs());

        s.push_str("  <Period id=\"0\" start=\"PT0S\">\n");
        s.push_str("    <AdaptationSet mimeType=\"video/mp4\" segmentAlignment=\"true\" startWithSAP=\"1\">\n");
        let _ = write!(
            s,
            "      <Representation id=\"v0\" bandwidth=\"{}\"",
            self.bandwidth
        );
        if self.width > 0 && self.height > 0 {
            let _ = write!(s, " width=\"{}\" height=\"{}\"", self.width, self.height);
        }
        if let Some(codecs) = &self.codecs {
            let _ = write!(s, " codecs=\"{codecs}\"");
        }
        s.push_str(">\n");

        let _ = write!(
            s,
            "        <SegmentTemplate timescale=\"{}\" initialization=\"{}\" media=\"{}\" startNumber=\"{}\"",
            self.timescale,
            self.init_uri,
            self.media_template,
            self.start_number()
        );
        if self.low_latency {
            let _ = write!(
                s,
                " availabilityTimeOffset=\"{:.3}\" availabilityTimeComplete=\"false\"",
                (self.max_seg_secs() - self.part_target).max(0.0)
            );
        }
        s.push_str(">\n");
        s.push_str("          <SegmentTimeline>\n");
        // Emit an explicit <S> per segment; `t` only needs restating when the
        // timeline is non-contiguous, but stating it on the first entry anchors
        // the window after eviction.
        for (i, (_, t)) in self.segments.iter().enumerate() {
            if i == 0 {
                let _ = writeln!(
                    s,
                    "            <S t=\"{}\" d=\"{}\"/>",
                    t.start_ms, t.dur_ms
                );
            } else {
                let _ = writeln!(s, "            <S d=\"{}\"/>", t.dur_ms);
            }
        }
        s.push_str("          </SegmentTimeline>\n");
        s.push_str("        </SegmentTemplate>\n");
        s.push_str("      </Representation>\n");
        s.push_str("    </AdaptationSet>\n");
        s.push_str("  </Period>\n");
        s.push_str("</MPD>\n");
        s
    }
}

/// Convert Unix seconds to an ISO-8601 UTC timestamp (`YYYY-MM-DDThh:mm:ssZ`),
/// dependency-free (Howard Hinnant's civil-from-days algorithm).
fn unix_to_iso8601(secs: u64) -> String {
    let days = (secs / 86_400) as i64;
    let rem = secs % 86_400;
    let (h, m, sec) = (rem / 3600, (rem % 3600) / 60, rem % 60);

    let z = days + 719_468;
    let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
    let doe = z - era * 146_097;
    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
    let y = yoe + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let month = if mp < 10 { mp + 3 } else { mp - 9 };
    let year = if month <= 2 { y + 1 } else { y };
    format!("{year:04}-{month:02}-{d:02}T{h:02}:{m:02}:{sec:02}Z")
}

/// Keyframe-boundary DASH packager writing CMAF segments + a live MPD through a
/// [`StorageBackend`]. The DASH counterpart to
/// [`HlsSegmenter`](super::HlsSegmenter).
pub struct DashPackager<M: Muxer, S: StorageBackend> {
    engine: SegmentEngine<M, S>,
    manifest: DashManifest,
    seg_start_ms: u64,
    // ── LL-DASH chunked-CMAF state (inactive when `part_target_ms == 0`) ──
    /// Part target in milliseconds; `0` disables partial-fragment emission.
    part_target_ms: i64,
    /// DTS (ms) at which the current part began, once a segment is open.
    part_anchor_dts: Option<i64>,
    /// DTS (ms) of the most recently written frame (for the final part's flush).
    last_dts: i64,
    /// Fragments produced for the in-progress segment, concatenated into the full
    /// segment file at the cut so non-LL clients can fetch the whole segment.
    seg_part_bytes: Vec<Bytes>,
    /// Accumulated exact media duration (ms) of the in-progress segment, summed
    /// from each fragment's `last_segment_duration_ms` so the SegmentTimeline `d`
    /// matches the media even though parts are taken incrementally.
    seg_dur_ms: u64,
    /// Last completed segment's duration (ms), used as the estimate for the
    /// provisional timeline entry of the next (open) segment in LL mode.
    last_dur_ms: u64,
    /// Optional live-chunk sink: receives each fragment as it is produced so a
    /// delivery layer can stream the open segment over HTTP chunked transfer.
    sink: Option<Arc<dyn ChunkSink>>,
}

impl<M: Muxer, S: StorageBackend> DashPackager<M, S> {
    /// New packager writing under `prefix`, `target_duration`-second segments,
    /// a `window`-segment live MPD.
    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),
            manifest: DashManifest::new(window),
            seg_start_ms: 0,
            part_target_ms: 0,
            part_anchor_dts: None,
            last_dts: 0,
            seg_part_bytes: Vec::new(),
            seg_dur_ms: 0,
            last_dur_ms: target_duration.max(1) * 1000,
            sink: None,
        }
    }

    /// Enable LL-DASH: advertise `availabilityTimeOffset` in the MPD *and* emit a
    /// CMAF fragment roughly every `part_target` seconds (via
    /// [`Muxer::take_partial`]) so an attached [`ChunkSink`] can stream the
    /// in-progress segment. Requires a muxer that implements `take_partial`.
    pub fn low_latency(mut self, part_target: f64) -> Self {
        self.manifest = self.manifest.low_latency(part_target);
        self.part_target_ms = (part_target.max(0.05) * 1000.0) as i64;
        self
    }

    /// Attach a [`ChunkSink`] that receives each segment's fragments as produced.
    /// Only meaningful together with [`low_latency`](Self::low_latency).
    pub fn with_chunk_sink(mut self, sink: Arc<dyn ChunkSink>) -> Self {
        self.sink = Some(sink);
        self
    }

    /// The storage key of the MPD.
    pub fn manifest_key(&self) -> String {
        format!("{}/manifest.mpd", self.engine.prefix)
    }

    async fn ensure_init(&mut self, frame: &MediaFrame) -> Result<()> {
        if let Some(uri) = self.engine.ensure_init(frame).await? {
            self.manifest.set_uris(
                uri,
                format!("seg$Number$.{}", self.engine.muxer.extension()),
            );
            // Media info for the Representation, once the muxer can report it.
            self.manifest
                .set_media_info(self.engine.muxer.codec_string(), 0, 0, BANDWIDTH_HINT);
        }
        Ok(())
    }

    /// LL-DASH: take the buffered samples as one CMAF fragment, hand it to the
    /// sink (for chunked streaming) and stash it for the full-segment file. Sums
    /// the fragment's exact duration into the segment total. No-op if the muxer
    /// has nothing buffered.
    async fn flush_part(&mut self) -> Result<()> {
        let Some(bytes) = self.engine.muxer.take_partial()? else {
            return Ok(());
        };
        if let Some(d) = self.engine.muxer.last_segment_duration_ms() {
            self.seg_dur_ms += d;
        }
        if let Some(sink) = &self.sink {
            sink.chunk(self.engine.seq, bytes.clone());
        }
        self.seg_part_bytes.push(bytes);
        self.part_anchor_dts = Some(self.last_dts);
        Ok(())
    }

    async fn cut(&mut self, duration: f64) -> Result<()> {
        if self.part_target_ms > 0 {
            // LL-DASH: flush the segment's final fragment, then assemble the full
            // segment file from its fragments (concatenation == the whole segment).
            self.flush_part().await?;
            let mut full = BytesMut::new();
            for b in &self.seg_part_bytes {
                full.put_slice(b);
            }
            // Defensive: if no fragments were produced, fall back to a whole-segment
            // flush so we never write a 0-byte file.
            let full = if full.is_empty() {
                self.engine.muxer.finish_segment()?
            } else {
                full.freeze()
            };
            if full.is_empty() {
                return Ok(());
            }
            let uri = self.engine.segment_uri(self.engine.seq);
            let key = self.engine.key(&uri);
            self.engine.storage.put(&key, full).await?;
            // Summed fragment durations track the media exactly; fall back to the
            // keyframe-span when (defensively) no fragment reported a duration.
            let dur_ms = if self.seg_dur_ms > 0 {
                self.seg_dur_ms
            } else {
                (duration * 1000.0) as u64
            };
            // The segment was already advertised (provisional duration) when it
            // opened; patch in its exact duration now that it is complete.
            self.manifest.set_last_duration(dur_ms);
            self.seg_start_ms += dur_ms;
            self.last_dur_ms = dur_ms;
            self.write_manifest().await?;
            if let Some(sink) = &self.sink {
                sink.close(self.engine.seq);
            }
            self.seg_part_bytes.clear();
            self.seg_dur_ms = 0;
            self.engine.seq += 1;
            return Ok(());
        }

        let bytes = self.engine.muxer.finish_segment()?;
        if bytes.is_empty() {
            return Ok(());
        }
        let uri = self.engine.segment_uri(self.engine.seq);
        let key = self.engine.key(&uri);
        self.engine.storage.put(&key, bytes).await?;
        // Prefer the muxer's exact fragment duration so the SegmentTimeline `t`/`d`
        // tracks the media's `tfdt` precisely. Falling back to the keyframe-span
        // `duration` drifts by the last sample's fallback duration each segment,
        // which dash.js renders as a stutter at every segment boundary.
        let dur_ms = self
            .engine
            .muxer
            .last_segment_duration_ms()
            .filter(|&d| d > 0)
            .unwrap_or((duration * 1000.0) as u64);
        self.manifest
            .push(self.engine.seq, self.seg_start_ms, dur_ms);
        self.seg_start_ms += dur_ms;
        self.write_manifest().await?;
        self.engine.seq += 1;
        Ok(())
    }

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

#[async_trait]
impl<M: Muxer, S: StorageBackend> Packager for DashPackager<M, S> {
    async fn push(&mut self, frame: &MediaFrame) -> Result<()> {
        self.ensure_init(frame).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 fragment run at this keyframe. Use DTS
            // (decode/arrival order) so part timing stays monotonic with B-frames.
            self.part_anchor_dts = Some(frame.dts);
            self.seg_part_bytes.clear();
            self.seg_dur_ms = 0;
            if let Some(sink) = &self.sink {
                sink.open(self.engine.seq);
            }
            // LL-DASH: advertise the segment in the SegmentTimeline as soon as it
            // opens (with the previous segment's duration as an estimate, patched
            // exactly at the cut), so the player discovers and starts fetching it
            // chunked right away — instead of stalling at every segment boundary
            // until the next manifest refresh reveals the completed segment.
            if self.part_target_ms > 0 {
                self.manifest
                    .push(self.engine.seq, self.seg_start_ms, self.last_dur_ms);
                self.write_manifest().await?;
            }
        }
        self.engine.muxer.write(frame)?;
        self.last_dts = frame.dts;

        // LL-DASH: emit a CMAF fragment once the part target has elapsed. Cut at
        // frame boundaries so a fragment never splits an access unit.
        if self.part_target_ms > 0 {
            if let Some(anchor) = self.part_anchor_dts {
                if frame.dts - anchor >= self.part_target_ms {
                    self.flush_part().await?;
                }
            }
        }
        Ok(())
    }

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

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

    #[test]
    fn iso8601_known_vector() {
        // 2026-06-15T11:35:43Z → seconds since epoch.
        assert_eq!(unix_to_iso8601(0), "1970-01-01T00:00:00Z");
        assert_eq!(unix_to_iso8601(1_609_459_200), "2021-01-01T00:00:00Z");
        assert_eq!(
            unix_to_iso8601(1_609_459_200 + 3661),
            "2021-01-01T01:01:01Z"
        );
    }

    #[test]
    fn manifest_renders_dynamic_mpd_with_timeline() {
        let mut m = DashManifest::new(3);
        m.set_media_info(Some("avc1.42001f".into()), 1280, 720, 2_000_000);
        m.push(0, 0, 2000);
        m.push(1, 2000, 1980);
        let out = m.render();
        assert!(out.contains("type=\"dynamic\""));
        assert!(out.contains("availabilityStartTime="));
        assert!(out.contains("<SegmentTimeline>"));
        assert!(out.contains("<S t=\"0\" d=\"2000\"/>"));
        assert!(out.contains("<S d=\"1980\"/>"));
        assert!(out.contains("codecs=\"avc1.42001f\""));
        assert!(out.contains("width=\"1280\" height=\"720\""));
    }

    #[test]
    fn low_latency_emits_availability_time_offset() {
        let mut m = DashManifest::new(3).low_latency(0.5);
        m.push(0, 0, 2000);
        let out = m.render();
        assert!(out.contains("availabilityTimeOffset="));
        assert!(out.contains("availabilityTimeComplete=\"false\""));
    }

    #[test]
    fn finish_marks_static() {
        let mut m = DashManifest::new(3);
        m.push(0, 0, 2000);
        m.finish();
        let out = m.render();
        assert!(out.contains("type=\"static\""));
        assert!(!out.contains("minimumUpdatePeriod"));
    }

    #[tokio::test]
    async fn ll_dash_streams_fragments_through_sink() {
        use std::sync::Mutex;

        #[derive(Default)]
        struct CollectSink {
            chunks: Mutex<Vec<(u64, Vec<u8>)>>,
            closed: Mutex<Vec<u64>>,
        }
        impl ChunkSink for CollectSink {
            fn open(&self, _seq: u64) {}
            fn chunk(&self, seq: u64, bytes: Bytes) {
                self.chunks.lock().unwrap().push((seq, bytes.to_vec()));
            }
            fn close(&self, seq: u64) {
                self.closed.lock().unwrap().push(seq);
            }
        }

        let store = InMemoryStorage::new();
        let sink = Arc::new(CollectSink::default());
        let mut dash = DashPackager::new(
            PassthroughMuxer::new("m4s"),
            store.clone(),
            "live/dash",
            2,
            5,
        )
        .low_latency(0.5)
        .with_chunk_sink(sink.clone());

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

        // Segment 0 was streamed as multiple fragments and then closed. Collect the
        // owned bytes (and drop the lock) before the await below.
        let (seg0_count, concat) = {
            let chunks = sink.chunks.lock().unwrap();
            let seg0: Vec<&Vec<u8>> = chunks
                .iter()
                .filter(|(s, _)| *s == 0)
                .map(|(_, b)| b)
                .collect();
            let mut concat = Vec::new();
            for b in &seg0 {
                concat.extend_from_slice(b);
            }
            (seg0.len(), concat)
        };
        assert!(seg0_count >= 2, "segment 0 streamed in multiple fragments");
        assert!(sink.closed.lock().unwrap().contains(&0), "segment 0 closed");

        // The concatenation of the streamed fragments equals the full segment file.
        let full = store.get("live/dash/seg0.m4s").await.unwrap();
        assert_eq!(
            &concat[..],
            &full[..],
            "streamed fragments == full segment file"
        );
    }

    #[tokio::test]
    async fn packager_writes_segments_and_mpd() {
        let store = InMemoryStorage::new();
        let mut dash = DashPackager::new(
            PassthroughMuxer::new("m4s"),
            store.clone(),
            "live/dash",
            2,
            5,
        );

        for i in 0..4 {
            let pts = i * 1000;
            dash.push(&video_frame(pts, true)).await.unwrap();
            dash.push(&video_frame(pts + 500, false)).await.unwrap();
        }
        dash.finish().await.unwrap();

        let mpd =
            String::from_utf8(store.get("live/dash/manifest.mpd").await.unwrap().to_vec()).unwrap();
        assert!(mpd.contains("<MPD"));
        assert!(mpd.contains("type=\"static\"")); // finished
        assert!(mpd.contains("<SegmentTimeline>"));
        assert!(store.get("live/dash/seg0.m4s").await.is_ok());
    }
}