libfreemkv 0.31.2

//! BD Transport Stream muxer — PES frames → 192-byte BD-TS packets.
//!
//! Takes PES frames and writes them as BD-TS (Blu-ray transport stream)
//! packets. Each frame is wrapped in a PES header, split into TS packets,
//! and prepended with the 4-byte TP_extra_header.

use super::hevc::{hvcc_to_annex_b, length_prefixed_to_annex_b};
use std::io::{self, Write};

const SYNC_BYTE: u8 = 0x47;
const TS_PAYLOAD: usize = 184;

/// PID range treated as video (HEVC, triggers Annex-B conversion + RAI
/// on keyframes). Both `write_frame` and `build_pes_header` consult this
/// so a PID's stream_id and its NAL handling can never disagree.
const VIDEO_PID_RANGE: std::ops::RangeInclusive<u16> = 0x1011..=0x101F;

/// Largest PES payload that fits a bounded `PES_packet_length` (u16) on a
/// `0xBD` (private_stream_1) stream after the 8 PES-header bytes. Frames
/// larger than this are split into multiple PES so the length field stays
/// spec-conformant (the unbounded `0` length is only legal for video).
const MAX_BD_PES_PAYLOAD: usize = u16::MAX as usize - 8;

fn is_video_pid(pid: u16) -> bool {
    VIDEO_PID_RANGE.contains(&pid)
}

/// BD-TS muxer: PES frames in, 192-byte BD-TS packets out.
///
/// Constructed over an output writer and a slice of per-track PIDs. The
/// `track` index passed to [`TsMuxer::write_frame`] and
/// [`TsMuxer::set_codec_private`] is the position in that PID slice; all
/// per-track state vectors are sized to `pids.len()`. PIDs in
/// `0x1011..=0x101F` are treated as video (length-prefixed NALUs in,
/// Annex B out, with parameter-set prepend and RAI on keyframes); every
/// other PID is carried as `private_stream_1` (`0xBD`) audio/subtitle.
/// All tracks share one PTS origin seeded from the first video frame, so
/// audio/video PTS offsets are preserved.
pub struct TsMuxer<W: Write> {
    writer: W,
    pids: Vec<u16>,
    continuity: Vec<u8>,                  // per-PID continuity counter (0-15)
    codec_privates: Vec<Option<Vec<u8>>>, // per-track codec_private (for video parameter sets)
    params_written: Vec<bool>,            // per-track: have we written parameter sets?
    /// Global PTS origin (nanoseconds), seeded by the FIRST video frame so
    /// the audio/video offset is preserved. Frames that arrive before it
    /// is set saturate to 0.
    base_pts_ns: Option<i64>,
}

impl<W: Write> TsMuxer<W> {
    pub fn new(writer: W, pids: &[u16]) -> Self {
        let n = pids.len();
        Self {
            writer,
            pids: pids.to_vec(),
            continuity: vec![0u8; n],
            codec_privates: vec![None; n],
            params_written: vec![false; n],
            base_pts_ns: None,
        }
    }

    /// Set codec_private data for a track. Used to prepend VPS/SPS/PPS
    /// as Annex B NALs before the first keyframe in the transport stream.
    ///
    /// `track` is the index into the PID slice passed to [`TsMuxer::new`].
    /// Returns [`Error::MuxTrackRange`](crate::error::Error::MuxTrackRange)
    /// for an out-of-range index.
    pub fn set_codec_private(&mut self, track: usize, data: Vec<u8>) -> io::Result<()> {
        if track >= self.codec_privates.len() {
            return Err(crate::error::Error::MuxTrackRange {
                track,
                tracks: self.codec_privates.len(),
            }
            .into());
        }
        self.codec_privates[track] = Some(data);
        Ok(())
    }

    /// Write a PES frame as BD-TS packets.
    /// Video frame data is expected as length-prefixed NALUs (MKV/PES format)
    /// and is converted to Annex B for transport stream.
    ///
    /// `track` is the index into the PID slice passed to [`TsMuxer::new`].
    /// Returns [`Error::MuxTrackRange`](crate::error::Error::MuxTrackRange)
    /// for an out-of-range index.
    pub fn write_frame(
        &mut self,
        track: usize,
        pts_ns: i64,
        keyframe: bool,
        data: &[u8],
    ) -> io::Result<()> {
        if track >= self.pids.len() {
            return Err(crate::error::Error::MuxTrackRange {
                track,
                tracks: self.pids.len(),
            }
            .into());
        }
        let pid = self.pids[track];
        let is_video = is_video_pid(pid);

        // Drop non-key video before any keyframe — decoder has no IDR or
        // parameter sets to anchor on.
        if is_video && !keyframe && !self.params_written[track] {
            return Ok(());
        }

        // Seed the global PTS origin from the FIRST video frame only, so the
        // audio/video offset is preserved. A leading audio frame must not
        // pull the base up and collapse the first video IDR to t=0.
        if is_video {
            self.base_pts_ns.get_or_insert(pts_ns);
        }
        let base = self.base_pts_ns.unwrap_or(pts_ns);
        let pts_ns = pts_ns.saturating_sub(base);

        // For video: convert length-prefixed NALUs to Annex B (start codes).
        // Prepend codec_private parameter sets on the FIRST keyframe only.
        //
        // Arm `params_written` on the first video keyframe regardless of
        // whether it carries data: an empty-data keyframe still anchors
        // the stream, and leaving the flag unset would make every later
        // non-key frame fail the drop guard above and silently vanish.
        // For non-video the ES bytes pass through unchanged, so borrow
        // `data` directly rather than copying it; only video needs an
        // owned Annex-B conversion buffer.
        let es_data: std::borrow::Cow<'_, [u8]> = if is_video {
            let mut annex_b = Vec::new();
            if keyframe && !self.params_written[track] {
                if let Some(ref cp) = self.codec_privates[track] {
                    if let Some(params) = hvcc_to_annex_b(cp) {
                        annex_b.extend_from_slice(&params);
                    }
                }
                self.params_written[track] = true;
            }
            annex_b.extend_from_slice(&length_prefixed_to_annex_b(data));
            std::borrow::Cow::Owned(annex_b)
        } else {
            std::borrow::Cow::Borrowed(data)
        };

        let pts_90k = if pts_ns >= 0 {
            (pts_ns as u64).saturating_mul(9) / 100_000
        } else {
            0
        };

        // Video PES may be unbounded (length 0); a 0xBD private_stream_1
        // PES must carry a bounded length, so split oversized audio/sub
        // access units into multiple PES packets. Each emitted PES carries
        // the same PTS and starts on its own PUSI packet (only the keyframe
        // RAI rides the first packet of the first PES).
        if is_video || es_data.len() <= MAX_BD_PES_PAYLOAD {
            self.write_pes_chain(track, pid, pts_90k, is_video, keyframe, &es_data)?;
        } else {
            let mut first_pes = true;
            for chunk in es_data.chunks(MAX_BD_PES_PAYLOAD) {
                self.write_pes_chain(track, pid, pts_90k, is_video, keyframe && first_pes, chunk)?;
                first_pes = false;
            }
        }
        Ok(())
    }

    /// Wrap `es_data` in a PES header and split it into 192-byte BD-TS
    /// packets. `keyframe` drives the RAI bit on the first packet (video
    /// only). The PES header and ES bytes are sliced in place — no second
    /// full-frame copy.
    fn write_pes_chain(
        &mut self,
        track: usize,
        pid: u16,
        pts_90k: u64,
        is_video: bool,
        keyframe: bool,
        es_data: &[u8],
    ) -> io::Result<()> {
        let pes_header = build_pes_header(pid, pts_90k, es_data.len());

        // Logical PES packet = header bytes followed by es_data. It is
        // indexed (and written) in place, without materializing the
        // concatenation, to avoid a second full-frame copy on the hot path.
        let pes_len = pes_header.len() + es_data.len();

        let mut offset = 0;
        let mut first = true;
        while offset < pes_len {
            let remaining = pes_len - offset;

            // Invariant: TP_extra(4) + TS_header(4) + AF(af_bytes) + payload(payload_len) = 192,
            // i.e. af_bytes + payload_len = TS_PAYLOAD (184).
            // RAI on first packet of a keyframe video PES requires AF with flags=0x40.
            let want_rai = first && keyframe && is_video;

            // Pick payload_len and af_bytes per case.
            let (af_bytes, payload_len): (usize, usize) = if want_rai {
                // Minimum AF = 2 bytes (length=1, flags=0x40). Payload caps at 182.
                let max_payload = TS_PAYLOAD - 2;
                let p = remaining.min(max_payload);
                (TS_PAYLOAD - p, p)
            } else if remaining >= TS_PAYLOAD {
                (0, TS_PAYLOAD) // no AF, full payload
            } else {
                // Stuffing-only AF, payload = remaining.
                (TS_PAYLOAD - remaining, remaining)
            };

            // TP_extra_header (4 bytes — arrival time, set to 0)
            let tp_extra = [0u8; 4];

            // TS header (4 bytes)
            let cc = self.continuity[track];
            self.continuity[track] = (cc + 1) & 0x0F;

            let mut ts_header = [0u8; 4];
            ts_header[0] = SYNC_BYTE;
            ts_header[1] = ((pid >> 8) as u8) & 0x1F;
            if first {
                ts_header[1] |= 0x40; // PUSI
            }
            ts_header[2] = pid as u8;
            ts_header[3] = if af_bytes > 0 {
                0x30 | cc // AF + payload
            } else {
                0x10 | cc // payload only
            };

            self.writer.write_all(&tp_extra)?;
            self.writer.write_all(&ts_header)?;

            if af_bytes > 0 {
                static STUFF_FF: [u8; 184] = [0xFF; 184];
                if want_rai {
                    // RAI AF: length byte + flags(0x40) + (af_bytes - 2) stuffing.
                    let af_len_field = (af_bytes - 1) as u8;
                    self.writer.write_all(&[af_len_field])?;
                    self.writer.write_all(&[0x40u8])?;
                    let stuff = af_bytes - 2;
                    if stuff > 0 {
                        self.writer.write_all(&STUFF_FF[..stuff])?;
                    }
                } else {
                    // Stuffing-only AF.
                    // af_bytes == 1: length=0, no flags.
                    // af_bytes >= 2: length = af_bytes-1, flags=0, rest 0xFF.
                    if af_bytes == 1 {
                        self.writer.write_all(&[0u8])?;
                    } else {
                        self.writer.write_all(&[(af_bytes - 1) as u8])?;
                        self.writer.write_all(&[0u8])?;
                        if af_bytes > 2 {
                            self.writer.write_all(&STUFF_FF[..af_bytes - 2])?;
                        }
                    }
                }
            }

            // Write the payload span [offset, offset+payload_len), which may
            // straddle the header/es_data boundary — emit each side in one
            // write_all rather than copying the whole frame again.
            let end = offset + payload_len;
            let hdr_len = pes_header.len();
            if offset < hdr_len {
                let hdr_end = end.min(hdr_len);
                self.writer.write_all(&pes_header[offset..hdr_end])?;
            }
            if end > hdr_len {
                let es_start = offset.max(hdr_len) - hdr_len;
                let es_end = end - hdr_len;
                self.writer.write_all(&es_data[es_start..es_end])?;
            }

            offset += payload_len;
            first = false;
        }

        Ok(())
    }

    /// Flush the underlying writer. BD-TS needs no stream trailer, so this
    /// only drains buffering; the muxer remains usable afterwards.
    pub fn finish(&mut self) -> io::Result<()> {
        self.writer.flush()
    }
}

/// Build a PES packet header for a BD stream.
fn build_pes_header(pid: u16, pts_90k: u64, data_len: usize) -> Vec<u8> {
    // Determine stream_id from PID range
    let stream_id: u8 = if is_video_pid(pid) {
        0xE0 // video
    } else {
        0xBD // audio, PGS subtitle, or default (private stream 1)
    };

    let pes_data_len = data_len + 8; // 3 header bytes + 5 PTS bytes + data
    let mut header = Vec::with_capacity(14);

    // Start code: 00 00 01 stream_id
    header.push(0x00);
    header.push(0x00);
    header.push(0x01);
    header.push(stream_id);

    // PES packet length. The unbounded form (0) is only spec-legal for
    // video; `write_frame` splits oversized 0xBD access units so a private
    // stream always fits a bounded u16 length here. The `> 65535` arm
    // remains a defensive fallback for video only.
    if stream_id == 0xE0 || pes_data_len > 65535 {
        header.push(0x00);
        header.push(0x00);
    } else {
        let len = pes_data_len as u16;
        header.push((len >> 8) as u8);
        header.push(len as u8);
    }

    // Flags: 10xx xxxx — MPEG-2, PTS present
    header.push(0x80); // marker bits
    header.push(0x80); // PTS present

    // PES header data length
    header.push(5); // 5 bytes of PTS

    // PTS (5 bytes, 33-bit timestamp with markers)
    let pts = pts_90k & 0x1_FFFF_FFFF;
    header.push(0x21 | (((pts >> 29) & 0x0E) as u8));
    header.push(((pts >> 22) & 0xFF) as u8);
    header.push(0x01 | (((pts >> 14) & 0xFE) as u8));
    header.push(((pts >> 7) & 0xFF) as u8);
    header.push(0x01 | (((pts << 1) & 0xFE) as u8));

    header
}

#[cfg(test)]
mod tests {
    use super::*;

    const BD_PACKET_SIZE: usize = 192;
    const VIDEO_PID: u16 = 0x1011;

    /// Parsed BD-TS packet (192 bytes total: 4 TP_extra + 4 TS header + 184 body).
    struct TsPacket {
        pid: u16,
        pusi: bool,
        #[allow(dead_code)]
        cc: u8,
        /// Adaptation field body (length byte stripped) when present.
        af: Option<Vec<u8>>,
        /// Payload bytes (after AF, if any).
        payload: Vec<u8>,
    }

    /// Walk 192-byte BD-TS packets.
    fn parse_bd_ts(buf: &[u8]) -> Vec<TsPacket> {
        let mut out = Vec::new();
        for chunk in buf.chunks(BD_PACKET_SIZE) {
            if chunk.len() != BD_PACKET_SIZE {
                break;
            }
            // Skip TP_extra_header (4 bytes), parse TS header.
            let h = &chunk[4..];
            assert_eq!(h[0], 0x47, "bad sync byte");
            let pusi = (h[1] & 0x40) != 0;
            let pid = (((h[1] & 0x1F) as u16) << 8) | h[2] as u16;
            let afc = (h[3] >> 4) & 0x03;
            let cc = h[3] & 0x0F;
            let body = &h[4..]; // 184 bytes

            let (af, payload) = match afc {
                0b01 => (None, body.to_vec()),
                0b11 => {
                    let af_len = body[0] as usize;
                    let af_body = body[1..1 + af_len].to_vec();
                    let payload = body[1 + af_len..].to_vec();
                    (Some(af_body), payload)
                }
                0b10 => {
                    let af_len = body[0] as usize;
                    (Some(body[1..1 + af_len].to_vec()), Vec::new())
                }
                _ => (None, Vec::new()),
            };

            out.push(TsPacket {
                pid,
                pusi,
                cc,
                af,
                payload,
            });
        }
        out
    }

    /// Build a fake HEVC NAL with a 4-byte length prefix.
    /// nal_type=19/20 are IDR; 1 is non-key (TRAIL_N/R).
    fn fake_hevc_nal(nal_type: u8, body_len: usize) -> Vec<u8> {
        let mut nal = Vec::with_capacity(2 + body_len);
        // 2-byte NAL header: forbidden_zero(1)=0 | nal_unit_type(6) | layer_id(6)=0 | tid_plus1(3)=1
        nal.push((nal_type & 0x3F) << 1);
        nal.push(0x01);
        for i in 0..body_len {
            nal.push((i & 0xFF) as u8);
        }
        let mut framed = Vec::with_capacity(4 + nal.len());
        framed.extend_from_slice(&(nal.len() as u32).to_be_bytes());
        framed.extend_from_slice(&nal);
        framed
    }

    #[test]
    fn keyframe_param_threads_through() {
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
            let idr = fake_hevc_nal(19, 100);
            mux.write_frame(0, 0, true, &idr).unwrap();
            let p = fake_hevc_nal(1, 80);
            mux.write_frame(0, 41_000_000, false, &p).unwrap();
            mux.finish().unwrap();
        }
        assert!(!sink.is_empty());
        let packets = parse_bd_ts(&sink);
        assert!(packets.iter().any(|p| p.pid == VIDEO_PID && p.pusi));
    }

    #[test]
    fn rai_set_on_first_packet_of_keyframe_pes() {
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
            let idr = fake_hevc_nal(19, 200);
            mux.write_frame(0, 0, true, &idr).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        let first_pusi = packets
            .iter()
            .find(|p| p.pid == VIDEO_PID && p.pusi)
            .expect("video PUSI packet exists");
        let af = first_pusi.af.as_ref().expect("AF present on keyframe PES");
        assert!(!af.is_empty(), "AF body has flags byte");
        assert_eq!(af[0] & 0x40, 0x40, "RAI bit set");
    }

    #[test]
    fn rai_clear_on_non_keyframe_pes() {
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
            let idr = fake_hevc_nal(19, 100);
            mux.write_frame(0, 0, true, &idr).unwrap();
            let p = fake_hevc_nal(1, 100);
            mux.write_frame(0, 41_000_000, false, &p).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        // Second PUSI packet on the video PID belongs to the non-key frame.
        let pusi_video: Vec<&TsPacket> = packets
            .iter()
            .filter(|p| p.pid == VIDEO_PID && p.pusi)
            .collect();
        assert!(pusi_video.len() >= 2, "two PUSI packets expected");
        let second = pusi_video[1];
        match &second.af {
            None => {}
            Some(af) if af.is_empty() => {} // length=0 case
            Some(af) => assert_eq!(af[0] & 0x40, 0, "RAI must be clear on non-key PES"),
        }
    }

    #[test]
    fn codec_private_prepended_only_on_first_keyframe() {
        // Build a minimal hvcC with one recognizable NAL.
        let marker: &[u8] = &[0xDE, 0xAD, 0xBE, 0xEF, 0xCA, 0xFE];
        let mut hvcc = vec![0u8; 22];
        hvcc.push(1); // numArrays
        hvcc.push(32); // VPS NAL type byte (high bits arbitrary)
        hvcc.extend_from_slice(&1u16.to_be_bytes()); // numNalus
        hvcc.extend_from_slice(&(marker.len() as u16).to_be_bytes());
        hvcc.extend_from_slice(marker);

        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
            mux.set_codec_private(0, hvcc).unwrap();
            // Non-IDR before any IDR: should be dropped.
            let p = fake_hevc_nal(1, 50);
            mux.write_frame(0, 0, false, &p).unwrap();
            // IDR: should carry codec_private NALs prepended.
            let idr = fake_hevc_nal(19, 50);
            mux.write_frame(0, 41_000_000, true, &idr).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        // Concatenate all video PID payloads in emission order.
        let video_bytes: Vec<u8> = packets
            .iter()
            .filter(|p| p.pid == VIDEO_PID)
            .flat_map(|p| p.payload.clone())
            .collect();
        // marker bytes must appear in the stream (codec_private was prepended).
        let pos_marker = video_bytes
            .windows(marker.len())
            .position(|w| w == marker)
            .expect("codec_private marker bytes present in TS payload");
        // Find IDR body byte (0x26 = (19<<1)). pos_idr must be AFTER marker.
        let idr_header = (19u8 << 1) & 0x7E;
        let pos_idr = video_bytes
            .iter()
            .position(|&b| b == idr_header)
            .expect("IDR NAL header present in TS payload");
        assert!(
            pos_marker < pos_idr,
            "codec_private must precede IDR in TS payload"
        );
    }

    #[test]
    fn empty_data_keyframe_arms_params_so_later_frames_survive() {
        // An empty-data keyframe must still arm params_written; otherwise
        // every subsequent non-key frame would be dropped by the
        // pre-keyframe guard and the track would emit no real frames.
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
            // Keyframe with empty payload (e.g. a frame whose NALs were
            // all stripped upstream) — anchors the stream.
            mux.write_frame(0, 0, true, &[]).unwrap();
            // Now a real non-key frame; it must NOT be dropped.
            let p = fake_hevc_nal(1, 80);
            mux.write_frame(0, 41_000_000, false, &p).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        // The non-key frame's NAL body byte (0x02 = (1<<1)) must appear in
        // a video payload — proof it wasn't dropped.
        let video_bytes: Vec<u8> = packets
            .iter()
            .filter(|p| p.pid == VIDEO_PID)
            .flat_map(|p| p.payload.clone())
            .collect();
        assert!(
            video_bytes
                .windows(4)
                .any(|w| w == [0x00, 0x00, 0x00, 0x01]),
            "later non-key frame must survive after an empty-data keyframe"
        );
    }

    #[test]
    fn non_key_before_first_keyframe_dropped() {
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
            let p = fake_hevc_nal(1, 80);
            mux.write_frame(0, 0, false, &p).unwrap();
            mux.finish().unwrap();
        }
        // Nothing should be emitted for that PID.
        let packets = parse_bd_ts(&sink);
        assert!(
            !packets.iter().any(|p| p.pid == VIDEO_PID),
            "non-key before first keyframe must be dropped"
        );
    }

    const AUDIO_PID: u16 = 0x1100;

    /// Decode the 33-bit PTS from the first PUSI packet on `pid`. Assumes
    /// the PES header carries PTS (flags 0x80 at PES byte 7).
    fn first_pts_90k(packets: &[TsPacket], pid: u16) -> u64 {
        let pkt = packets
            .iter()
            .find(|p| p.pid == pid && p.pusi)
            .expect("PUSI packet present");
        // PES payload starts the packet payload: 00 00 01 stream_id len len
        // flags1 flags2 hdr_len then 5 PTS bytes.
        let p = &pkt.payload;
        let pts = &p[9..14];
        ((((pts[0] >> 1) & 0x07) as u64) << 30)
            | ((pts[1] as u64) << 22)
            | (((pts[2] >> 1) as u64) << 15)
            | ((pts[3] as u64) << 7)
            | ((pts[4] >> 1) as u64)
    }

    #[test]
    fn av_offset_preserved_with_audio_before_first_video() {
        // Audio at t=0 arrives BEFORE the first video keyframe at t=1s.
        // The global base must be seeded from the VIDEO frame so the
        // audio/video PTS offset is preserved (audio earlier ⇒ saturates to
        // 0, video lands at +1s = 90000 ticks), not both collapsed to 0.
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID, AUDIO_PID]);
            // Audio frame first, at PTS 0.
            mux.write_frame(1, 0, false, &[0x0B, 0x77, 0x00, 0x00])
                .unwrap();
            // Video keyframe at PTS 1s — seeds the base.
            let idr = fake_hevc_nal(19, 100);
            mux.write_frame(0, 1_000_000_000, true, &idr).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        let video_pts = first_pts_90k(&packets, VIDEO_PID);
        let audio_pts = first_pts_90k(&packets, AUDIO_PID);
        // Video keyframe is the base ⇒ its relative PTS is 0.
        assert_eq!(video_pts, 0, "video keyframe seeds the base at t=0");
        // Audio arrived 1s earlier ⇒ saturates to 0, NOT lifted past video.
        assert_eq!(audio_pts, 0, "earlier audio saturates to 0");
        assert!(
            audio_pts <= video_pts,
            "audio must not be pulled ahead of the video base"
        );
    }

    #[test]
    fn out_of_range_track_errors() {
        let mut sink: Vec<u8> = Vec::new();
        let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
        let err = mux.write_frame(5, 0, true, &[0xAA]).unwrap_err();
        assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
        let err2 = mux.set_codec_private(5, vec![0u8; 4]).unwrap_err();
        assert_eq!(err2.kind(), std::io::ErrorKind::InvalidInput);
    }

    #[test]
    fn oversized_bd_audio_pes_is_split_and_bounded() {
        // A private_stream_1 (0xBD) audio frame larger than the bounded PES
        // limit must be split into multiple PES, each with a non-zero
        // PES_packet_length (never the unbounded 0 form, which is illegal
        // for 0xBD).
        let mut sink: Vec<u8> = Vec::new();
        let big: Vec<u8> = (0..(MAX_BD_PES_PAYLOAD + 5000))
            .map(|i| (i & 0xFF) as u8)
            .collect();
        {
            let mut mux = TsMuxer::new(&mut sink, &[AUDIO_PID]);
            mux.write_frame(0, 0, false, &big).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        let pusi: Vec<&TsPacket> = packets
            .iter()
            .filter(|p| p.pid == AUDIO_PID && p.pusi)
            .collect();
        assert!(
            pusi.len() >= 2,
            "oversized audio must span ≥2 PES, got {}",
            pusi.len()
        );
        for p in pusi {
            // PES length field at payload bytes [4..6] must be non-zero.
            let len = u16::from_be_bytes([p.payload[4], p.payload[5]]);
            assert_ne!(len, 0, "0xBD PES must carry a bounded length");
        }
    }

    // ════════════════════════════════════════════════════════════════════
    // Added hardening tests
    // ════════════════════════════════════════════════════════════════════

    /// Concatenate the ES payloads of all packets on `pid`, stripping the
    /// PES header off each PUSI packet. A PUSI packet starts a PES whose
    /// header is `00 00 01 stream_id len len 80 80 05` + 5 PTS bytes = 14
    /// bytes for our muxer (always PTS-present, header_data_length 5).
    fn reassemble_es(packets: &[TsPacket], pid: u16) -> Vec<u8> {
        let mut out = Vec::new();
        for p in packets.iter().filter(|p| p.pid == pid) {
            if p.pusi {
                // Skip the 14-byte PES header (3 startcode + 1 stream_id +
                // 2 length + 2 flags + 1 hdr_len + 5 PTS).
                assert!(p.payload.len() >= 14, "PUSI payload holds a PES header");
                out.extend_from_slice(&p.payload[14..]);
            } else {
                out.extend_from_slice(&p.payload);
            }
        }
        out
    }

    #[test]
    fn every_packet_is_exactly_192_bytes() {
        // BD-TS packets are 192 bytes (4 TP_extra + 188 TS). The muxer must
        // never emit a short or long packet — that would desync any reader.
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
            let idr = fake_hevc_nal(19, 500); // spans several packets
            mux.write_frame(0, 0, true, &idr).unwrap();
            mux.finish().unwrap();
        }
        assert!(!sink.is_empty());
        assert_eq!(sink.len() % BD_PACKET_SIZE, 0, "output must be 192-aligned");
        for chunk in sink.chunks(BD_PACKET_SIZE) {
            assert_eq!(chunk.len(), BD_PACKET_SIZE);
            assert_eq!(chunk[4], SYNC_BYTE, "TS sync byte at offset 4");
        }
    }

    #[test]
    fn audio_es_round_trips_byte_for_byte_through_demuxer() {
        // The mux→demux round trip must preserve every audio ES byte. A
        // muxer that dropped/duplicated payload on a packet boundary would
        // silently corrupt the audio. Use a payload spanning many packets.
        let es: Vec<u8> = (0..1000u32).map(|i| (i & 0xFF) as u8).collect();
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[AUDIO_PID]);
            mux.write_frame(0, 0, false, &es).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        let got = reassemble_es(&packets, AUDIO_PID);
        assert_eq!(got, es, "audio ES must survive mux→demux unchanged");
    }

    #[test]
    fn continuity_counter_wraps_modulo_16() {
        // ISO 13818-1: continuity_counter is 4 bits, incrementing per packet
        // on a PID and wrapping 15→0. A frame spanning >16 packets exercises
        // the wrap.
        let es: Vec<u8> = vec![0xAB; 20 * 184]; // 20 packets of audio payload
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[AUDIO_PID]);
            mux.write_frame(0, 0, false, &es).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        let ccs: Vec<u8> = packets
            .iter()
            .filter(|p| p.pid == AUDIO_PID)
            .map(|p| p.cc)
            .collect();
        assert!(ccs.len() > 16, "need >16 packets to test the wrap");
        for w in ccs.windows(2) {
            assert_eq!(w[1], (w[0] + 1) & 0x0F, "CC increments mod 16");
        }
        // Prove a wrap actually occurred (a 15→0 transition exists).
        assert!(
            ccs.windows(2).any(|w| w[0] == 0x0F && w[1] == 0x00),
            "CC must wrap 15→0 across >16 packets"
        );
    }

    #[test]
    fn pts_encoded_at_90khz_decodes_correctly() {
        // pts_ns → 90 kHz ticks = pts_ns * 9 / 100_000. 1 second (1e9 ns)
        // = 90_000 ticks. The first (base) video frame rebases to 0, so use
        // a second frame at a known offset and check its encoded PTS.
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
            let idr = fake_hevc_nal(19, 50);
            mux.write_frame(0, 0, true, &idr).unwrap(); // base = 0
            let p = fake_hevc_nal(1, 50);
            // +1 second relative to base.
            mux.write_frame(0, 1_000_000_000, false, &p).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        let video_pusi: Vec<&TsPacket> = packets
            .iter()
            .filter(|p| p.pid == VIDEO_PID && p.pusi)
            .collect();
        assert!(video_pusi.len() >= 2);
        // Decode PTS of the SECOND video PES (the +1s frame).
        let p = &video_pusi[1].payload;
        let pts = ((((p[9] >> 1) & 0x07) as u64) << 30)
            | ((p[10] as u64) << 22)
            | (((p[11] >> 1) as u64) << 15)
            | ((p[12] as u64) << 7)
            | ((p[13] >> 1) as u64);
        assert_eq!(pts, 90_000, "1s offset encodes to 90000 ticks @ 90 kHz");
    }

    #[test]
    fn video_pes_uses_unbounded_length_field() {
        // build_pes_header: video (stream_id 0xE0) always uses the unbounded
        // (0x0000) PES_packet_length form — video PES can exceed u16.
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
            let idr = fake_hevc_nal(19, 50);
            mux.write_frame(0, 0, true, &idr).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        let pusi = packets
            .iter()
            .find(|p| p.pid == VIDEO_PID && p.pusi)
            .unwrap();
        // PES length field at payload[4..6].
        let len = u16::from_be_bytes([pusi.payload[4], pusi.payload[5]]);
        assert_eq!(len, 0, "video PES length field is the unbounded 0 form");
        // stream_id (payload[3]) is 0xE0 for video.
        assert_eq!(pusi.payload[3], 0xE0, "video stream_id 0xE0");
    }

    #[test]
    fn audio_pes_stream_id_is_private_stream_1() {
        // Non-video PIDs are carried as private_stream_1 (0xBD).
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[AUDIO_PID]);
            mux.write_frame(0, 0, false, &[0x0B, 0x77, 0x01, 0x02])
                .unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        let pusi = packets
            .iter()
            .find(|p| p.pid == AUDIO_PID && p.pusi)
            .unwrap();
        assert_eq!(pusi.payload[3], 0xBD, "audio carried as private_stream_1");
    }

    #[test]
    fn negative_relative_pts_saturates_to_zero() {
        // A frame earlier than the base (negative relative PTS) must encode
        // PTS 0, never an underflowed huge value. Audio at t=0 before a
        // video keyframe at t=2s: base=video, audio relative = -2s → 0.
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID, AUDIO_PID]);
            mux.write_frame(1, 0, false, &[0x0B, 0x77, 0x00, 0x00])
                .unwrap();
            let idr = fake_hevc_nal(19, 50);
            mux.write_frame(0, 2_000_000_000, true, &idr).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        assert_eq!(
            first_pts_90k(&packets, AUDIO_PID),
            0,
            "earlier audio saturates to 0"
        );
    }

    #[test]
    fn no_base_seeded_by_audio_only_stream() {
        // If only audio frames are written (no video), base_pts_ns is never
        // seeded by them; each frame rebases to itself via unwrap_or(pts_ns),
        // so the first audio frame lands at relative 0. Proves audio never
        // seeds the global base (which would corrupt later A/V offsets).
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[AUDIO_PID]);
            // First audio frame at 5s.
            mux.write_frame(0, 5_000_000_000, false, &[0x01, 0x02])
                .unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        // With no video base, base = unwrap_or(pts_ns) = this frame's pts,
        // so relative PTS is 0.
        assert_eq!(first_pts_90k(&packets, AUDIO_PID), 0);
    }

    #[test]
    fn oversized_audio_split_preserves_all_bytes() {
        // The oversized-0xBD split must not lose or reorder ES bytes across
        // the multiple PES it produces. Reassembling all audio packets must
        // reproduce the original frame exactly.
        let big: Vec<u8> = (0..(MAX_BD_PES_PAYLOAD + 3000))
            .map(|i| (i & 0xFF) as u8)
            .collect();
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[AUDIO_PID]);
            mux.write_frame(0, 0, false, &big).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        let got = reassemble_es(&packets, AUDIO_PID);
        assert_eq!(got.len(), big.len(), "no bytes lost in the PES split");
        assert_eq!(got, big, "split audio reassembles byte-for-byte");
    }

    #[test]
    fn af_plus_payload_always_fills_184() {
        // Invariant from write_pes_chain: af_bytes + payload_len == 184 on
        // every packet (so the 192-byte frame is exact). Verify for a video
        // keyframe (which forces an RAI adaptation field on packet 1).
        let mut sink: Vec<u8> = Vec::new();
        {
            let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
            let idr = fake_hevc_nal(19, 400);
            mux.write_frame(0, 0, true, &idr).unwrap();
            mux.finish().unwrap();
        }
        let packets = parse_bd_ts(&sink);
        for p in packets.iter().filter(|p| p.pid == VIDEO_PID) {
            let af_total = p.af.as_ref().map(|a| a.len() + 1).unwrap_or(0); // +1 length byte
            assert_eq!(
                af_total + p.payload.len(),
                184,
                "AF area + payload must fill the 184-byte TS body"
            );
        }
    }
}