multimux 0.1.0

//! LL-HLS request handlers: master/media playlists and the init/segment/part
//! byte ranges the media playlist references.
//!
//! Master/media playlist tags are RFC 8216 §4.3.4 (`#EXT-X-STREAM-INF`) and
//! §4.3.3 (`#EXTM3U`/`#EXT-X-VERSION`, rendered by
//! [`crate::store::StreamStore::media_playlist_m3u8`]); the blocking reload
//! query parameters (`_HLS_msn`/`_HLS_part`) are the Blocking Playlist Reload
//! mechanism of RFC 8216bis §6.2.5.2 — the client asks the origin to hold the
//! response open until the requested Media Sequence Number/part is available,
//! bounded so the origin never hangs indefinitely.

use std::sync::Arc;
use std::time::Duration;

use axum::extract::{Path, Query, State};
use axum::http::{StatusCode, header};
use axum::response::{IntoResponse, Response};
use serde::Deserialize;

use crate::origin::AppState;
use crate::store::StreamStore;

/// Track id for the single rendition multimux currently serves per stream
/// (no multi-track/multi-rendition support yet).
pub const DEFAULT_TRACK_ID: u32 = 1;

/// Placeholder `BANDWIDTH` (bits/second) advertised in the master playlist's
/// `#EXT-X-STREAM-INF` — multimux does not yet measure actual encoded
/// bitrate, so a single fixed estimate is used for the single variant served.
const PLACEHOLDER_BANDWIDTH_BPS: u64 = 5_000_000;

/// Upper bound on how long a blocking `media.m3u8` request
/// (`_HLS_msn`/`_HLS_part`) waits for the requested part/segment before
/// falling back to rendering whatever is currently available. RFC 8216bis
/// §6.2.5.2 requires the origin to eventually respond either way — this cap
/// keeps a stalled/slow source from hanging the HTTP response forever.
const BLOCKING_RELOAD_TIMEOUT: Duration = Duration::from_secs(5);

const MEDIA_PLAYLIST_CONTENT_TYPE: &str = "application/vnd.apple.mpegurl";
const MP4_CONTENT_TYPE: &str = "video/mp4";

/// `GET /:stream/master.m3u8` — a minimal single-variant master playlist
/// pointing at `media.m3u8`.
pub async fn master_playlist(
    State(state): State<Arc<AppState>>,
    Path(stream): Path<String>,
) -> Response {
    if !state.streams.contains_key(&stream) {
        return StatusCode::NOT_FOUND.into_response();
    }
    let body =
        format!("#EXTM3U\n#EXT-X-STREAM-INF:BANDWIDTH={PLACEHOLDER_BANDWIDTH_BPS}\nmedia.m3u8\n");
    ([(header::CONTENT_TYPE, MEDIA_PLAYLIST_CONTENT_TYPE)], body).into_response()
}

/// Blocking playlist reload query parameters (RFC 8216bis §6.2.5.2).
#[derive(Debug, Default, Deserialize)]
pub struct BlockingReloadQuery {
    /// The Media Sequence Number the client already has, plus one — the
    /// origin should not respond until a segment/part beyond this is ready.
    #[serde(rename = "_HLS_msn")]
    pub hls_msn: Option<u64>,
    /// The part index (within `_HLS_msn`) the client is waiting for.
    #[serde(rename = "_HLS_part")]
    pub hls_part: Option<u32>,
}

/// Block until `store`'s in-progress segment/part reaches at least
/// `(msn, part)`, or [`BLOCKING_RELOAD_TIMEOUT`] elapses. Never hangs
/// indefinitely and never errors — on timeout (or a closed watch channel) it
/// simply returns, and the caller renders the playlist as it currently is.
async fn wait_for_progress(store: &StreamStore, msn: u64, part: u32) {
    let mut rx = store.subscribe();
    let wait = async {
        loop {
            let (in_progress_seg_seq, part_count) = store.latest_progress();
            let satisfied = u64::from(in_progress_seg_seq) > msn
                || (u64::from(in_progress_seg_seq) == msn && part_count > part);
            if satisfied {
                return;
            }
            if rx.changed().await.is_err() {
                return;
            }
        }
    };
    let _ = tokio::time::timeout(BLOCKING_RELOAD_TIMEOUT, wait).await;
}

/// `GET /:stream/media.m3u8` — the LL-HLS media playlist for
/// [`DEFAULT_TRACK_ID`], blocking on `_HLS_msn`/`_HLS_part` when present.
pub async fn media_playlist(
    State(state): State<Arc<AppState>>,
    Path(stream): Path<String>,
    Query(q): Query<BlockingReloadQuery>,
) -> Response {
    let Some(store) = state.streams.get(&stream) else {
        return StatusCode::NOT_FOUND.into_response();
    };
    if let Some(msn) = q.hls_msn {
        let part = q.hls_part.unwrap_or(0);
        wait_for_progress(store, msn, part).await;
    }
    let body = store.media_playlist_m3u8(DEFAULT_TRACK_ID);
    ([(header::CONTENT_TYPE, MEDIA_PLAYLIST_CONTENT_TYPE)], body).into_response()
}

/// `GET /:stream/:file` — catch-all for the dynamic init/segment/part
/// filenames [`crate::store::StreamStore::media_playlist_m3u8`] emits.
///
/// A single catch-all (rather than three routes with per-filename literals)
/// because axum 0.7's `matchit`-based router cannot mix multiple params with
/// literal text in one path segment (e.g. `seg-:track-:seq.m4s`) — only one
/// param per segment is supported, capturing the whole segment. `file` is
/// parsed here instead.
pub async fn dynamic_file(
    State(state): State<Arc<AppState>>,
    Path((stream, file)): Path<(String, String)>,
) -> Response {
    let Some(store) = state.streams.get(&stream) else {
        return StatusCode::NOT_FOUND.into_response();
    };
    match resolve_file(store, &file) {
        Some(bytes) => ([(header::CONTENT_TYPE, MP4_CONTENT_TYPE)], bytes).into_response(),
        None => StatusCode::NOT_FOUND.into_response(),
    }
}

/// Parse a dynamic origin filename and fetch its bytes from `store`:
/// - `init-{track}.mp4` -> [`StreamStore::init_bytes`]
/// - `seg-{track}-{seq}.m4s` -> [`StreamStore::segment_bytes`]
/// - `part-{track}-{seq}.{idx}.m4s` -> [`StreamStore::part_bytes`]
///
/// `{track}` is validated as a number but otherwise unused: `store` holds a
/// single track's data (see [`DEFAULT_TRACK_ID`]). Returns `None` (-> 404)
/// for any filename that doesn't match one of the three shapes, or whose
/// numeric fields don't parse.
fn resolve_file(store: &StreamStore, file: &str) -> Option<Vec<u8>> {
    if let Some(rest) = file.strip_prefix("init-") {
        let track = rest.strip_suffix(".mp4")?;
        track.parse::<u32>().ok()?;
        return store.init_bytes();
    }
    if let Some(rest) = file.strip_prefix("seg-") {
        let rest = rest.strip_suffix(".m4s")?;
        let (track, seq) = rest.split_once('-')?;
        track.parse::<u32>().ok()?;
        let seq: u32 = seq.parse().ok()?;
        return store.segment_bytes(seq);
    }
    if let Some(rest) = file.strip_prefix("part-") {
        let rest = rest.strip_suffix(".m4s")?;
        let (track_seq, idx) = rest.rsplit_once('.')?;
        let (track, seq) = track_seq.split_once('-')?;
        track.parse::<u32>().ok()?;
        let seq: u32 = seq.parse().ok()?;
        let idx: u32 = idx.parse().ok()?;
        return store.part_bytes(seq, idx);
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::origin::AppState;
    use std::collections::HashMap;
    use transmux::ll_hls::{PartInfo, SegmentInfo};

    fn part(seq: u32, idx: u32) -> PartInfo {
        PartInfo {
            bytes: vec![0x10 + idx as u8; 4],
            duration: 0.5,
            independent: idx == 0,
            segment_seq: seq,
            part_index: idx,
        }
    }

    fn seg(seq: u32) -> SegmentInfo {
        SegmentInfo {
            bytes: vec![0x20 + seq as u8; 8],
            duration: 4.0,
            segment_seq: seq,
            part_count: 2,
        }
    }

    /// One populated stream ("cam1"): a closed segment 1, plus two live parts
    /// of in-progress segment 2 -- so `latest_progress()` is `(2, 2)`.
    fn make_state() -> Arc<AppState> {
        let store = Arc::new(StreamStore::new(4.0, 500, 4));
        store.set_init(vec![0xAA; 8]);
        store.add_segment(seg(1));
        store.add_part(part(2, 0));
        store.add_part(part(2, 1));
        let mut streams = HashMap::new();
        streams.insert("cam1".to_string(), store);
        Arc::new(AppState { streams })
    }

    async fn body_string(resp: Response) -> String {
        let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        String::from_utf8(bytes.to_vec()).unwrap()
    }

    async fn body_bytes(resp: Response) -> Vec<u8> {
        axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap()
            .to_vec()
    }

    #[tokio::test]
    async fn master_playlist_known_stream_ok() {
        let state = make_state();
        let resp = master_playlist(State(state), Path("cam1".to_string())).await;
        assert_eq!(resp.status(), StatusCode::OK);
        let body = body_string(resp).await;
        assert!(body.contains("#EXTM3U"));
        assert!(body.contains("#EXT-X-STREAM-INF"));
        assert!(body.contains("media.m3u8"));
    }

    #[tokio::test]
    async fn master_playlist_unknown_stream_404() {
        let state = make_state();
        let resp = master_playlist(State(state), Path("nope".to_string())).await;
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);
    }

    #[tokio::test]
    async fn media_playlist_no_query_renders_now() {
        let state = make_state();
        let resp = media_playlist(
            State(state),
            Path("cam1".to_string()),
            Query(BlockingReloadQuery::default()),
        )
        .await;
        assert_eq!(resp.status(), StatusCode::OK);
        let body = body_string(resp).await;
        assert!(body.contains("#EXT-X-PART"), "body: {body}");
    }

    #[tokio::test]
    async fn media_playlist_unknown_stream_404() {
        let state = make_state();
        let resp = media_playlist(
            State(state),
            Path("nope".to_string()),
            Query(BlockingReloadQuery::default()),
        )
        .await;
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);
    }

    #[tokio::test]
    async fn media_playlist_already_satisfied_blocking_request_resolves_immediately() {
        // latest_progress() for "cam1" is (2, 2): asking for msn=1 (an
        // earlier segment) is already satisfied and must not wait.
        let state = make_state();
        let resp = media_playlist(
            State(state),
            Path("cam1".to_string()),
            Query(BlockingReloadQuery {
                hls_msn: Some(1),
                hls_part: Some(0),
            }),
        )
        .await;
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn media_playlist_already_satisfied_same_msn_lower_part() {
        // in_progress_seg_seq == msn and part_count(2) > part(1): satisfied.
        let state = make_state();
        let resp = media_playlist(
            State(state),
            Path("cam1".to_string()),
            Query(BlockingReloadQuery {
                hls_msn: Some(2),
                hls_part: Some(1),
            }),
        )
        .await;
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn dynamic_file_init_present() {
        let state = make_state();
        let resp = dynamic_file(
            State(state),
            Path(("cam1".to_string(), "init-1.mp4".to_string())),
        )
        .await;
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(body_bytes(resp).await, vec![0xAA; 8]);
    }

    #[tokio::test]
    async fn dynamic_file_segment_present_and_absent() {
        let state = make_state();
        let ok = dynamic_file(
            State(state.clone()),
            Path(("cam1".to_string(), "seg-1-1.m4s".to_string())),
        )
        .await;
        assert_eq!(ok.status(), StatusCode::OK);
        assert_eq!(body_bytes(ok).await, vec![0x21; 8]);

        let missing = dynamic_file(
            State(state),
            Path(("cam1".to_string(), "seg-1-99.m4s".to_string())),
        )
        .await;
        assert_eq!(missing.status(), StatusCode::NOT_FOUND);
    }

    #[tokio::test]
    async fn dynamic_file_part_present() {
        let state = make_state();
        let resp = dynamic_file(
            State(state),
            Path(("cam1".to_string(), "part-1-2.0.m4s".to_string())),
        )
        .await;
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(body_bytes(resp).await, vec![0x10; 4]);
    }

    #[tokio::test]
    async fn dynamic_file_unknown_stream_404() {
        let state = make_state();
        let resp = dynamic_file(
            State(state),
            Path(("nope".to_string(), "init-1.mp4".to_string())),
        )
        .await;
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);
    }

    #[tokio::test]
    async fn dynamic_file_unmatched_filename_404() {
        let state = make_state();
        let resp = dynamic_file(
            State(state),
            Path(("cam1".to_string(), "not-a-thing.txt".to_string())),
        )
        .await;
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);
    }
}