arcly-stream 0.1.3

//! A native, dependency-free **MPEG-TS** ([`MpegTsMuxer`]) implementing
//! [`Muxer`](super::Muxer).
//!
//! Unlike [`PassthroughMuxer`](super::PassthroughMuxer), this produces a real,
//! player-compatible transport stream: 188-byte packets, a Program Association
//! Table (PAT) and Program Map Table (PMT) at every segment boundary, PES
//! packetization with 90 kHz PTS/DTS, and a PCR carried on each keyframe. The
//! output is exactly what an HLS player expects from a `.ts` segment, so it
//! drops straight into [`HlsSegmenter`](super::HlsSegmenter).
//!
//! Gated behind the `mpegts` feature (which implies `hls` + `codec-h264`).
//!
//! ## Expected elementary-stream format
//!
//! The muxer is a pure container — it does not transcode. Frames must arrive in
//! the elementary form TS carries natively:
//!
//! * **Video (H.264)** — Annex-B (`00 00 00 01` start codes). Keyframes should
//!   already carry their SPS/PPS (the [`rtmp`](crate::protocol::rtmp) handler
//!   emits frames in exactly this shape).
//! * **Audio (AAC)** — ADTS framing (each frame prefixed with its 7-byte ADTS
//!   header).
//!
//! ```no_run
//! # #[cfg(all(feature = "mpegts", feature = "storage-fs"))]
//! # async fn demo(handle: arcly_stream::StreamHandle) -> arcly_stream::Result<()> {
//! use arcly_stream::packager::{HlsSegmenter, MpegTsMuxer, Packager};
//! use arcly_stream::storage::FsStorage;
//!
//! let mut seg = HlsSegmenter::new(
//!     MpegTsMuxer::new(),
//!     FsStorage::new("/var/hls"),
//!     "live/cam",
//!     6, // target segment duration (s)
//!     5, // live window (segments)
//! );
//! 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};
use bytes::{BufMut, Bytes, BytesMut};

// ── PID / stream-type assignments ───────────────────────────────────────────

const TS_PACKET_LEN: usize = 188;
const SYNC_BYTE: u8 = 0x47;

const PID_PAT: u16 = 0x0000;
const PID_PMT: u16 = 0x1000;
const PID_VIDEO: u16 = 0x0100;
const PID_AUDIO: u16 = 0x0101;

const STREAM_TYPE_H264: u8 = 0x1B;
const STREAM_TYPE_AAC_ADTS: u8 = 0x0F;

const PES_STREAM_ID_VIDEO: u8 = 0xE0;
const PES_STREAM_ID_AUDIO: u8 = 0xC0;

/// 90 kHz → multiply a millisecond timestamp to reach the MPEG-TS clock.
const TS_CLOCK_HZ_PER_MS: i64 = 90;

/// Native MPEG-TS muxer: packages H.264 + AAC elementary streams into a
/// single-program transport stream.
///
/// One instance packages one rendition. Continuity counters persist across
/// segments while PAT/PMT are re-emitted at each [`start_segment`], so every
/// produced segment is independently decodable.
///
/// [`start_segment`]: Muxer::start_segment
pub struct MpegTsMuxer {
    buf: BytesMut,
    cc_pat: u8,
    cc_pmt: u8,
    cc_video: u8,
    cc_audio: u8,
    /// Set once a video CONFIG (SPS/PPS) frame has been seen.
    codec_string: Option<String>,
    /// Whether the audio elementary stream should be declared in the PMT.
    has_audio: bool,
}

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

impl MpegTsMuxer {
    /// Create a muxer for a video-only or video+audio program.
    ///
    /// The audio elementary stream is declared in the PMT lazily, the first time
    /// an audio frame is written, so a video-only stream produces a clean
    /// single-ES program.
    pub fn new() -> Self {
        Self {
            buf: BytesMut::new(),
            cc_pat: 0,
            cc_pmt: 0,
            cc_video: 0,
            cc_audio: 0,
            codec_string: None,
            has_audio: false,
        }
    }

    /// Emit the PAT/PMT pair that opens every segment.
    fn write_psi(&mut self) {
        let pat = build_pat();
        push_section_packet(&mut self.buf, PID_PAT, &mut self.cc_pat, &pat);
        let pmt = build_pmt(self.has_audio);
        push_section_packet(&mut self.buf, PID_PMT, &mut self.cc_pmt, &pmt);
    }

    fn write_video(&mut self, frame: &MediaFrame) {
        let pts = frame.pts * TS_CLOCK_HZ_PER_MS;
        let dts = frame.dts * TS_CLOCK_HZ_PER_MS;
        let es = ensure_aud(&frame.data);
        let pes = build_pes(PES_STREAM_ID_VIDEO, pts, Some(dts), &es);
        // Carry a PCR + random-access indicator on keyframes so a player can
        // start cleanly at any segment head.
        let pcr = if frame.is_keyframe() { Some(dts) } else { None };
        let mut cc = self.cc_video;
        packetize(
            &mut self.buf,
            PID_VIDEO,
            &mut cc,
            &pes,
            pcr,
            frame.is_keyframe(),
        );
        self.cc_video = cc;
    }

    fn write_audio(&mut self, frame: &MediaFrame) {
        let pts = frame.pts * TS_CLOCK_HZ_PER_MS;
        let pes = build_pes(PES_STREAM_ID_AUDIO, pts, None, &frame.data);
        let mut cc = self.cc_audio;
        // Audio packets carry their own PCR-less adaptation only when stuffing
        // is required; `packetize` handles that.
        packetize(&mut self.buf, PID_AUDIO, &mut cc, &pes, None, false);
        self.cc_audio = cc;
    }
}

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

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

    fn write(&mut self, frame: &MediaFrame) -> Result<()> {
        match frame.codec {
            CodecId::H264 => {
                if frame.flags.contains(crate::FrameFlags::CONFIG) {
                    // Capture the codec string for the master playlist; the SPS/PPS
                    // bytes still flow through so keyframes remain self-contained.
                    self.codec_string = h264_codec_string(&frame.data);
                }
                self.write_video(frame);
            }
            CodecId::AAC => {
                if frame.flags.contains(crate::FrameFlags::CONFIG) {
                    // AudioSpecificConfig is not carried inline in TS (ADTS is
                    // self-describing); record presence and drop the config AU.
                    self.has_audio = true;
                    return Ok(());
                }
                self.has_audio = true;
                self.write_audio(frame);
            }
            other => {
                return Err(crate::StreamError::UnsupportedCodec(format!(
                    "MpegTsMuxer carries H.264 + AAC only, got {other:?}"
                )));
            }
        }
        Ok(())
    }

    fn finish_segment(&mut self) -> Result<Bytes> {
        Ok(std::mem::take(&mut self.buf).freeze())
    }

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

// ── PSI (PAT / PMT) construction ────────────────────────────────────────────

/// Build the PAT section body (without the leading pointer field).
fn build_pat() -> Vec<u8> {
    let mut s = Vec::with_capacity(16);
    s.push(0x00); // table_id = program_association_section
                  // section_syntax_indicator(1)=1, '0', reserved '11', section_length(12).
                  // section_length spans from transport_stream_id through CRC = 5 + 4 = 9.
    let section_len = 5 + 4 + 4; // ts_id..current_next(5) + program loop(4) + CRC(4)
    s.push(0xB0 | (((section_len >> 8) & 0x0F) as u8));
    s.push((section_len & 0xFF) as u8);
    s.extend_from_slice(&[0x00, 0x01]); // transport_stream_id = 1
    s.push(0xC1); // reserved '11', version 0, current_next_indicator 1
    s.push(0x00); // section_number
    s.push(0x00); // last_section_number
                  // One program: program_number = 1 → program_map_PID.
    s.extend_from_slice(&[0x00, 0x01]);
    s.push(0xE0 | (((PID_PMT >> 8) & 0x1F) as u8));
    s.push((PID_PMT & 0xFF) as u8);
    append_crc32(&mut s);
    s
}

/// Build the PMT section body (without the leading pointer field).
fn build_pmt(has_audio: bool) -> Vec<u8> {
    let mut es = Vec::with_capacity(10);
    // Video ES entry.
    es.push(STREAM_TYPE_H264);
    es.push(0xE0 | (((PID_VIDEO >> 8) & 0x1F) as u8));
    es.push((PID_VIDEO & 0xFF) as u8);
    es.push(0xF0); // reserved '1111', ES_info_length hi
    es.push(0x00); // ES_info_length lo = 0
    if has_audio {
        es.push(STREAM_TYPE_AAC_ADTS);
        es.push(0xE0 | (((PID_AUDIO >> 8) & 0x1F) as u8));
        es.push((PID_AUDIO & 0xFF) as u8);
        es.push(0xF0);
        es.push(0x00);
    }

    let mut s = Vec::with_capacity(20 + es.len());
    s.push(0x02); // table_id = TS_program_map_section
                  // section_length: program_number..program_info(9) + es + CRC(4).
    let section_len = 9 + es.len() + 4;
    s.push(0xB0 | (((section_len >> 8) & 0x0F) as u8));
    s.push((section_len & 0xFF) as u8);
    s.extend_from_slice(&[0x00, 0x01]); // program_number = 1
    s.push(0xC1); // version 0, current_next 1
    s.push(0x00); // section_number
    s.push(0x00); // last_section_number
    s.push(0xE0 | (((PID_VIDEO >> 8) & 0x1F) as u8)); // reserved '111', PCR_PID hi
    s.push((PID_VIDEO & 0xFF) as u8); // PCR_PID lo (video carries the PCR)
    s.push(0xF0); // reserved '1111', program_info_length hi
    s.push(0x00); // program_info_length lo = 0
    s.extend_from_slice(&es);
    append_crc32(&mut s);
    s
}

/// Wrap a PSI section in a single TS packet (PUSI set, pointer_field = 0,
/// payload-only, stuffed to 188 with `0xFF`). PSI sections here are always small
/// enough to fit one packet.
fn push_section_packet(out: &mut BytesMut, pid: u16, cc: &mut u8, section: &[u8]) {
    let mut pkt = Vec::with_capacity(TS_PACKET_LEN);
    pkt.push(SYNC_BYTE);
    pkt.push(0x40 | ((pid >> 8) & 0x1F) as u8); // PUSI = 1
    pkt.push((pid & 0xFF) as u8);
    pkt.push(0x10 | (*cc & 0x0F)); // AFC = payload only
    *cc = (*cc + 1) & 0x0F;
    pkt.push(0x00); // pointer_field
    pkt.extend_from_slice(section);
    pkt.resize(TS_PACKET_LEN, 0xFF);
    out.put_slice(&pkt);
}

// ── PES construction ────────────────────────────────────────────────────────

/// Build a PES packet for an elementary-stream access unit.
fn build_pes(stream_id: u8, pts: i64, dts: Option<i64>, payload: &[u8]) -> Vec<u8> {
    let pts_dts_flags = if dts.is_some() { 0b11 } else { 0b10 };
    let header_data_len = if dts.is_some() { 10 } else { 5 };

    let mut pes = Vec::with_capacity(payload.len() + 14);
    pes.extend_from_slice(&[0x00, 0x00, 0x01]); // packet_start_code_prefix
    pes.push(stream_id);

    // PES_packet_length: 0 = unbounded (only legal for video). For audio we set
    // the real length so the demuxer can frame it without a following start code.
    let pes_payload_len = 3 + header_data_len + payload.len();
    let length_field = if stream_id == PES_STREAM_ID_VIDEO {
        0
    } else {
        pes_payload_len.min(0xFFFF) as u16
    };
    pes.put_u16(length_field);

    pes.push(0x80); // '10', scrambling 0, priority 0, alignment 0, copyright 0, original 0
    pes.push((pts_dts_flags << 6) as u8); // PTS_DTS_flags, other flags 0
    pes.push(header_data_len as u8);

    write_timestamp(&mut pes, if dts.is_some() { 0b0011 } else { 0b0010 }, pts);
    if let Some(dts) = dts {
        write_timestamp(&mut pes, 0b0001, dts);
    }
    pes.extend_from_slice(payload);
    pes
}

/// Encode a 33-bit 90 kHz timestamp in the 5-byte PES form with the given
/// 4-bit prefix (`0b0011` PTS-with-DTS, `0b0010` PTS-only, `0b0001` DTS).
fn write_timestamp(out: &mut Vec<u8>, prefix: u8, ts: i64) {
    let ts = ts as u64 & 0x1_FFFF_FFFF; // 33 bits
    out.push((prefix << 4) | (((ts >> 30) & 0x07) as u8) << 1 | 0x01);
    out.push(((ts >> 22) & 0xFF) as u8);
    out.push((((ts >> 15) & 0x7F) as u8) << 1 | 0x01);
    out.push(((ts >> 7) & 0xFF) as u8);
    out.push(((ts & 0x7F) as u8) << 1 | 0x01);
}

/// Split a PES packet across 188-byte TS packets for `pid`, optionally carrying
/// a PCR (and random-access indicator) in the first packet's adaptation field.
fn packetize(out: &mut BytesMut, pid: u16, cc: &mut u8, pes: &[u8], pcr: Option<i64>, rai: bool) {
    let mut pos = 0;
    let mut first = true;
    while pos < pes.len() {
        let remaining = pes.len() - pos;
        let want_af = first && (pcr.is_some() || rai);

        // Fixed adaptation-field overhead: length byte + flags byte + 6-byte PCR.
        let af_fixed = if want_af {
            2 + if pcr.is_some() { 6 } else { 0 }
        } else {
            2 // length byte + flags byte, used only when stuffing is needed
        };
        let max_payload_with_af = TS_PACKET_LEN - 4 - af_fixed;
        let use_af = want_af || remaining < (TS_PACKET_LEN - 4);

        let (payload_len, stuffing) = if use_af {
            let take = remaining.min(max_payload_with_af);
            (take, max_payload_with_af - take)
        } else {
            (TS_PACKET_LEN - 4, 0)
        };

        let mut pkt = Vec::with_capacity(TS_PACKET_LEN);
        pkt.push(SYNC_BYTE);
        let pusi = if first { 0x40 } else { 0x00 };
        pkt.push(pusi | ((pid >> 8) & 0x1F) as u8);
        pkt.push((pid & 0xFF) as u8);
        let afc = if use_af { 0b11 } else { 0b01 };
        pkt.push((afc << 4) | (*cc & 0x0F));
        *cc = (*cc + 1) & 0x0F;

        if use_af {
            let mut flags = 0u8;
            if want_af && rai {
                flags |= 0x40; // random_access_indicator
            }
            let pcr_present = want_af && pcr.is_some();
            if pcr_present {
                flags |= 0x10; // PCR_flag
            }
            // adaptation_field_length = flags(1) + pcr(6?) + stuffing.
            let af_len = 1 + if pcr_present { 6 } else { 0 } + stuffing;
            pkt.push(af_len as u8);
            pkt.push(flags);
            if pcr_present {
                write_pcr(&mut pkt, pcr.unwrap());
            }
            pkt.resize(pkt.len() + stuffing, 0xFF);
        }

        pkt.extend_from_slice(&pes[pos..pos + payload_len]);
        debug_assert_eq!(pkt.len(), TS_PACKET_LEN);
        out.put_slice(&pkt);

        pos += payload_len;
        first = false;
    }
}

/// Encode a 42-bit Program Clock Reference (33-bit base @ 90 kHz, 9-bit
/// extension fixed at 0) into the 6-byte adaptation-field form.
fn write_pcr(out: &mut Vec<u8>, base90: i64) {
    let base = base90 as u64 & 0x1_FFFF_FFFF;
    out.push(((base >> 25) & 0xFF) as u8);
    out.push(((base >> 17) & 0xFF) as u8);
    out.push(((base >> 9) & 0xFF) as u8);
    out.push(((base >> 1) & 0xFF) as u8);
    out.push((((base & 0x01) as u8) << 7) | 0x7E); // marker '111111', ext bit 8 = 0
    out.push(0x00); // ext bits 7..0 = 0
}

/// Prepend an Access Unit Delimiter NAL if the Annex-B AU does not already start
/// with one, improving demuxer/player compatibility.
fn ensure_aud(annexb: &[u8]) -> Bytes {
    // First NAL type after the leading start code.
    let starts_with_aud = match annexb {
        [0, 0, 0, 1, b, ..] | [0, 0, 1, b, ..] => (b & 0x1F) == 9,
        _ => false,
    };
    if starts_with_aud {
        return Bytes::copy_from_slice(annexb);
    }
    let mut out = BytesMut::with_capacity(annexb.len() + 6);
    out.put_slice(&[0x00, 0x00, 0x00, 0x01, 0x09, 0xF0]); // AUD, primary_pic_type 7
    out.put_slice(annexb);
    out.freeze()
}

/// Derive the HLS `CODECS` string (`avc1.PPCCLL`) from an Annex-B SPS/PPS AU.
fn h264_codec_string(annexb: &[u8]) -> Option<String> {
    use crate::codec::parser::CodecParser;
    let params = crate::codec::h264::H264::parse_config(annexb)?;
    Some(crate::codec::h264::hls_codec_string(&params))
}

// ── CRC-32/MPEG-2 ───────────────────────────────────────────────────────────

/// Append the MPEG-2 systems CRC-32 (poly `0x04C11DB7`, init `0xFFFFFFFF`, no
/// reflection, no final XOR) over `data`'s current contents.
fn append_crc32(data: &mut Vec<u8>) {
    let mut crc: u32 = 0xFFFF_FFFF;
    for &b in data.iter() {
        crc ^= (b as u32) << 24;
        for _ in 0..8 {
            crc = if crc & 0x8000_0000 != 0 {
                (crc << 1) ^ 0x04C1_1DB7
            } else {
                crc << 1
            };
        }
    }
    data.extend_from_slice(&crc.to_be_bytes());
}

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

    fn annexb_keyframe() -> Bytes {
        // SPS (valid baseline 1280x720) + PPS + IDR slice, all Annex-B.
        Bytes::from_static(&[
            0, 0, 0, 1, 0x67, 0x42, 0x00, 0x1F, 0xF4, 0x02, 0x80, 0x2D, 0x80, // SPS
            0, 0, 0, 1, 0x68, 0xCE, 0x3C, 0x80, // PPS
            0, 0, 0, 1, 0x65, 0x88, 0x84, 0x00, // IDR slice
        ])
    }

    fn all_packets_aligned(ts: &[u8]) {
        assert_eq!(ts.len() % TS_PACKET_LEN, 0, "stream is packet-aligned");
        for chunk in ts.chunks(TS_PACKET_LEN) {
            assert_eq!(chunk[0], SYNC_BYTE, "every packet starts with 0x47");
        }
    }

    fn pids(ts: &[u8]) -> Vec<u16> {
        ts.chunks(TS_PACKET_LEN)
            .map(|p| (((p[1] & 0x1F) as u16) << 8) | p[2] as u16)
            .collect()
    }

    #[test]
    fn segment_opens_with_pat_pmt_and_is_packet_aligned() {
        let mut m = MpegTsMuxer::new();
        m.start_segment().unwrap();
        let kf = MediaFrame::new_video(0, 0, annexb_keyframe(), CodecId::H264, true);
        m.write(&kf).unwrap();
        let seg = m.finish_segment().unwrap();

        all_packets_aligned(&seg);
        let pids = pids(&seg);
        assert_eq!(pids[0], PID_PAT, "first packet is the PAT");
        assert_eq!(pids[1], PID_PMT, "second packet is the PMT");
        assert!(pids[2..].contains(&PID_VIDEO), "video PID present");
    }

    #[test]
    fn keyframe_packet_carries_pcr_and_random_access() {
        let mut m = MpegTsMuxer::new();
        m.start_segment().unwrap();
        m.write(&MediaFrame::new_video(
            0,
            0,
            annexb_keyframe(),
            CodecId::H264,
            true,
        ))
        .unwrap();
        let seg = m.finish_segment().unwrap();

        // Find the first video packet; it must have an adaptation field (AFC=11)
        // with the random-access + PCR flags set.
        let vpkt = seg
            .chunks(TS_PACKET_LEN)
            .find(|p| (((p[1] & 0x1F) as u16) << 8 | p[2] as u16) == PID_VIDEO)
            .expect("video packet");
        let afc = (vpkt[3] >> 4) & 0x03;
        assert_eq!(afc, 0b11, "adaptation field + payload present");
        let flags = vpkt[5];
        assert_ne!(flags & 0x40, 0, "random_access_indicator set");
        assert_ne!(flags & 0x10, 0, "PCR_flag set");
    }

    #[test]
    fn audio_is_declared_in_pmt_only_after_first_audio_frame() {
        let mut m = MpegTsMuxer::new();
        m.start_segment().unwrap();
        // ADTS AAC frame (syncword 0xFFF…) — 7-byte header + 1 payload byte.
        let adts = Bytes::from_static(&[0xFF, 0xF1, 0x50, 0x80, 0x01, 0xA0, 0x01, 0xCC]);
        m.write(&MediaFrame::new_audio(0, adts, CodecId::AAC))
            .unwrap();
        let seg = m.finish_segment().unwrap();
        assert!(pids(&seg).contains(&PID_AUDIO), "audio PID present");
    }

    #[test]
    fn long_payload_spans_multiple_packets_and_stays_aligned() {
        let mut m = MpegTsMuxer::new();
        m.start_segment().unwrap();
        let mut big = vec![0, 0, 0, 1, 0x65]; // IDR start
        big.extend(std::iter::repeat_n(0xAB, 1000)); // force several TS packets
        m.write(&MediaFrame::new_video(
            33,
            33,
            Bytes::from(big),
            CodecId::H264,
            true,
        ))
        .unwrap();
        let seg = m.finish_segment().unwrap();
        all_packets_aligned(&seg);
        let video_packets = pids(&seg).iter().filter(|&&p| p == PID_VIDEO).count();
        assert!(video_packets > 1, "payload spread over multiple packets");
    }

    #[test]
    fn config_frame_sets_codec_string() {
        let mut m = MpegTsMuxer::new();
        m.start_segment().unwrap();
        let mut cfg = MediaFrame::new_video(0, 0, annexb_keyframe(), CodecId::H264, false);
        cfg.flags |= FrameFlags::CONFIG;
        m.write(&cfg).unwrap();
        assert!(
            m.codec_string().is_some_and(|s| s.starts_with("avc1.")),
            "codec string derived from SPS"
        );
    }

    #[test]
    fn crc32_matches_known_vector() {
        // CRC-32/MPEG-2 of "123456789" is 0x0376E6E7.
        let mut data = b"123456789".to_vec();
        append_crc32(&mut data);
        let crc = &data[data.len() - 4..];
        assert_eq!(crc, &[0x03, 0x76, 0xE6, 0xE7]);
    }
}