libfreemkv 0.8.2

//! MPEG-2 Program Stream (PS) demuxer.
//!
//! DVDs use MPEG-2 Program Stream, which has:
//! - Pack headers (00 00 01 BA) with SCR timestamps
//! - PES packets (00 00 01 [stream_id]) with variable length
//! - System headers (00 00 01 BB)
//! - Program end code (00 00 01 B9)
//!
//! Stream IDs:
//! - 0xE0-0xEF: video (usually 0xE0)
//! - 0xC0-0xDF: MPEG audio
//! - 0xBD: private stream 1 (AC3, DTS, LPCM, subtitles via sub-stream ID)

/// Pack header start code suffix.
const PACK_HEADER_ID: u8 = 0xBA;

/// System header start code suffix.
const SYSTEM_HEADER_ID: u8 = 0xBB;

/// Program end start code suffix.
const PROGRAM_END_ID: u8 = 0xB9;

/// Private stream 1 (AC3, DTS, LPCM, subtitles).
const PRIVATE_STREAM_1: u8 = 0xBD;

/// A demuxed PES packet from the Program Stream.
#[derive(Debug, Clone)]
pub struct PsPacket {
    /// PES stream ID (0xE0 for video, 0xC0 for audio, 0xBD for private, etc.).
    pub stream_id: u8,
    /// Sub-stream ID for private stream 1 (AC3: 0x80-0x87, DTS: 0x88-0x8F,
    /// LPCM: 0xA0-0xA7, subtitles: 0x20-0x3F).
    pub sub_stream_id: Option<u8>,
    /// Presentation timestamp in 90kHz ticks.
    pub pts: Option<u64>,
    /// Decode timestamp in 90kHz ticks.
    pub dts: Option<u64>,
    /// Elementary stream payload data.
    pub data: Vec<u8>,
}

/// MPEG-2 Program Stream demuxer.
///
/// Accepts raw PS bytes via `feed()` and produces demuxed PES packets.
/// Handles non-aligned input by buffering leftover bytes between calls.
pub struct PsDemuxer {
    buffer: Vec<u8>,
}

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

impl PsDemuxer {
    /// Create a new Program Stream demuxer.
    pub fn new() -> Self {
        Self {
            buffer: Vec::with_capacity(64 * 1024),
        }
    }

    /// Feed raw MPEG-2 PS bytes, returning any completely parsed PES packets.
    pub fn feed(&mut self, data: &[u8]) -> Vec<PsPacket> {
        self.buffer.extend_from_slice(data);
        self.extract_packets()
    }

    /// Flush remaining buffered data, returning any final PES packets.
    pub fn flush(&mut self) -> Vec<PsPacket> {
        // Try to extract whatever remains. If the buffer contains an incomplete
        // PES packet we cannot parse, it will be discarded.
        let packets = self.extract_packets();
        self.buffer.clear();
        packets
    }

    /// Scan the buffer for complete start-code-delimited units and parse them.
    fn extract_packets(&mut self) -> Vec<PsPacket> {
        let mut packets = Vec::with_capacity(4);
        let mut pos = 0;

        while let Some(sc) = find_start_code(&self.buffer, pos) {

            if sc + 3 >= self.buffer.len() {
                // Not enough bytes to read the start code ID.
                break;
            }

            let code = self.buffer[sc + 3];

            match code {
                PROGRAM_END_ID => {
                    // 00 00 01 B9 — 4 bytes, no payload.
                    pos = sc + 4;
                }
                PACK_HEADER_ID => {
                    // Pack header: need at least 14 bytes for MPEG-2 pack.
                    if sc + 14 > self.buffer.len() {
                        break; // wait for more data
                    }
                    // MPEG-2 packs have bit pattern 01 in bits 7-6 of byte 4.
                    let stuffing = (self.buffer[sc + 13] & 0x07) as usize;
                    let pack_len = 14 + stuffing;
                    if sc + pack_len > self.buffer.len() {
                        break;
                    }
                    pos = sc + pack_len;
                }
                SYSTEM_HEADER_ID => {
                    // System header: 00 00 01 BB [length:2] ...
                    if sc + 6 > self.buffer.len() {
                        break;
                    }
                    let header_len =
                        ((self.buffer[sc + 4] as usize) << 8) | self.buffer[sc + 5] as usize;
                    let total = 6 + header_len;
                    if sc + total > self.buffer.len() {
                        break;
                    }
                    pos = sc + total;
                }
                id if is_pes_stream_id(id) => {
                    // PES packet: 00 00 01 [stream_id] [length:2] ...
                    if sc + 6 > self.buffer.len() {
                        break;
                    }
                    let pes_packet_len =
                        ((self.buffer[sc + 4] as usize) << 8) | self.buffer[sc + 5] as usize;

                    // Total bytes = 6 (start code + stream_id + length) + pes_packet_len.
                    // A length of 0 means unbounded (video streams); in that case we need
                    // to find the next start code to delimit the packet.
                    let end = if pes_packet_len == 0 {
                        // Find the next start code after this one.
                        match find_start_code(&self.buffer, sc + 4) {
                            Some(next_sc) => next_sc,
                            None => break, // wait for more data
                        }
                    } else {
                        let e = sc + 6 + pes_packet_len;
                        if e > self.buffer.len() {
                            break; // wait for more data
                        }
                        e
                    };

                    if let Some(pkt) = parse_pes_packet(&self.buffer[sc..end]) {
                        packets.push(pkt);
                    }
                    pos = end;
                }
                _ => {
                    // Unknown start code — skip past it.
                    pos = sc + 4;
                }
            }
        }

        if pos > 0 {
            self.buffer.drain(..pos);
        }

        packets
    }
}

/// Check whether a start code byte is a valid PES stream ID that carries payload.
fn is_pes_stream_id(id: u8) -> bool {
    // Video: 0xE0-0xEF, MPEG audio: 0xC0-0xDF, private stream 1: 0xBD,
    // private stream 2: 0xBF, padding: 0xBE, ECM/EMM etc.
    // We parse anything in the PES range.
    matches!(id, 0xBD..=0xEF)
}

/// Parse a single PES packet from a byte slice that starts at the start code.
fn parse_pes_packet(data: &[u8]) -> Option<PsPacket> {
    // Minimum: 00 00 01 [id] [len:2] = 6 bytes
    if data.len() < 6 {
        return None;
    }
    if data[0] != 0x00 || data[1] != 0x00 || data[2] != 0x01 {
        return None;
    }

    let stream_id = data[3];

    // Padding stream — skip entirely.
    if stream_id == 0xBE {
        return None;
    }

    // Streams without standard PES header extension.
    if stream_id == 0xBF {
        let payload = if data.len() > 6 { &data[6..] } else { &[] };
        return Some(PsPacket {
            stream_id,
            sub_stream_id: None,
            pts: None,
            dts: None,
            data: payload.to_vec(),
        });
    }

    // Standard PES header: [6]=flags1, [7]=flags2, [8]=header_data_length
    if data.len() < 9 {
        return None;
    }

    let pts_dts_flags = (data[7] >> 6) & 0x03;
    let header_data_len = data[8] as usize;
    let header_end = 9 + header_data_len;

    if header_end > data.len() {
        return None;
    }

    let mut pts = None;
    let mut dts = None;

    if pts_dts_flags >= 2 && data.len() >= 14 {
        pts = Some(parse_pts(&data[9..14]));
    }
    if pts_dts_flags == 3 && data.len() >= 19 {
        dts = Some(parse_pts(&data[14..19]));
    }

    let payload = &data[header_end..];

    // For private stream 1, the first payload byte is the sub-stream ID,
    // followed by a sub-header whose length depends on the sub-stream type.
    let (sub_stream_id, es_data) = if stream_id == PRIVATE_STREAM_1 && !payload.is_empty() {
        let sub_id = payload[0];
        let skip = match sub_id {
            0x80..=0x8F => 4, // AC3/DTS: sub_id + frame_count + access_unit_ptr(2)
            0xA0..=0xA7 => 7, // LPCM: sub_id + frames + ptr(2) + emphasis + quant_freq + channels
            _ => 1,
        };
        let start = skip.min(payload.len());
        (Some(sub_id), payload[start..].to_vec())
    } else {
        (None, payload.to_vec())
    };

    Some(PsPacket {
        stream_id,
        sub_stream_id,
        pts,
        dts,
        data: es_data,
    })
}

/// Parse a 5-byte PTS/DTS timestamp field (33 bits at 90kHz).
///
/// Layout:
/// ```text
/// byte0: [marker_4bits][bit32][marker_1]
/// byte1: [bits 31..24]
/// byte2: [bits 23..15][marker_1]
/// byte3: [bits 14..7]
/// byte4: [bits 6..0][marker_1]
/// ```
fn parse_pts(buf: &[u8]) -> u64 {
    debug_assert!(buf.len() >= 5);
    let b0 = buf[0] as u64;
    let b1 = buf[1] as u64;
    let b2 = buf[2] as u64;
    let b3 = buf[3] as u64;
    let b4 = buf[4] as u64;

    ((b0 >> 1) & 0x07) << 30 | b1 << 22 | (b2 >> 1) << 15 | b3 << 7 | b4 >> 1
}

/// Find the position of the next start code (00 00 01) at or after `from`.
fn find_start_code(data: &[u8], from: usize) -> Option<usize> {
    if data.len() < from + 3 {
        return None;
    }
    (from..data.len() - 2).find(|&i| data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x01)
}

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

    // --- Pack header detection ---

    #[test]
    fn detect_pack_header() {
        let mut demuxer = PsDemuxer::new();

        // MPEG-2 pack header: 14 bytes, stuffing_length = 0
        let mut pack = vec![
            0x00, 0x00, 0x01, 0xBA, // start code
            0x44, 0x00, 0x04, 0x00, 0x04, 0x01, // SCR (6 bytes)
            0x01, 0x89, 0xC3, // mux_rate (3 bytes)
            0xF8, // stuffing_length = 0 (lower 3 bits)
        ];

        // Follow with a PES packet so we have a delimiter
        pack.extend_from_slice(&[
            0x00, 0x00, 0x01, 0xE0, // video stream
            0x00, 0x08, // length = 8
            0x80, 0x00, 0x00, // flags: no PTS/DTS, header_data_length = 0
            0xAA, 0xBB, 0xCC, 0xDD, 0xEE, // payload (5 bytes)
        ]);

        let packets = demuxer.feed(&pack);
        assert_eq!(packets.len(), 1);
        assert_eq!(packets[0].stream_id, 0xE0);
        assert_eq!(packets[0].data, vec![0xAA, 0xBB, 0xCC, 0xDD, 0xEE]);
    }

    #[test]
    fn pack_header_with_stuffing() {
        let mut demuxer = PsDemuxer::new();

        // Pack header with 3 stuffing bytes
        let mut data = vec![
            0x00, 0x00, 0x01, 0xBA, 0x44, 0x00, 0x04, 0x00, 0x04, 0x01, 0x01, 0x89, 0xC3,
            0xFB, // stuffing_length = 3
            0xFF, 0xFF, 0xFF, // stuffing bytes
        ];

        // Followed by a PES packet
        data.extend_from_slice(&[
            0x00, 0x00, 0x01, 0xC0, // audio stream
            0x00, 0x05, // length = 5
            0x80, 0x00, 0x00, // flags: no PTS, header_data_len=0
            0x11, 0x22, // payload
        ]);

        let packets = demuxer.feed(&data);
        assert_eq!(packets.len(), 1);
        assert_eq!(packets[0].stream_id, 0xC0);
        assert_eq!(packets[0].data, vec![0x11, 0x22]);
    }

    // --- PES header + PTS parsing ---

    #[test]
    fn pes_header_with_pts() {
        let mut demuxer = PsDemuxer::new();

        // PTS = 90000 (1 second at 90kHz)
        // 90000 = 0x15F90
        // bit32=0, bits 29-15 = 0x0002BF, bits 14-0 = 0x1F90
        // byte0: 0010_0_1 = 0x21 ... actually let's encode properly:
        //
        // pts = 90000
        // byte0: (0010 << 4) | ((pts >> 29) & 0x0E) | 1
        //      = 0x20 | ((90000 >> 29) & 0x0E) | 1 = 0x20 | 0 | 1 = 0x21
        // byte1: (pts >> 22) & 0xFF = (90000 >> 22) & 0xFF = 0
        // byte2: ((pts >> 14) & 0xFE) | 1 = ((90000 >> 14) & 0xFE) | 1 = (0x0A & 0xFE) | 1 = 0x0B
        // byte3: (pts >> 7) & 0xFF = (90000 >> 7) & 0xFF = (703) & 0xFF = 0xBF
        // byte4: ((pts & 0x7F) << 1) | 1 = ((90000 & 0x7F) << 1) | 1 = (0x10 << 1) | 1 = 0x21

        let pts_bytes = encode_pts(90000, 0x20);

        let mut data = vec![
            0x00, 0x00, 0x01, 0xE0, // video stream
            0x00, 0x0D, // length = 13
            0x80, 0x80, 0x05, // flags: PTS only, header_data_len=5
        ];
        data.extend_from_slice(&pts_bytes);
        data.extend_from_slice(&[0xDE, 0xAD, 0xBE, 0xEF, 0x00]); // payload

        // Add a delimiter
        data.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]); // program end

        let packets = demuxer.feed(&data);
        assert_eq!(packets.len(), 1);
        assert_eq!(packets[0].stream_id, 0xE0);
        assert_eq!(packets[0].pts, Some(90000));
        assert!(packets[0].dts.is_none());
        assert_eq!(packets[0].data, vec![0xDE, 0xAD, 0xBE, 0xEF, 0x00]);
    }

    #[test]
    fn pes_header_with_pts_and_dts() {
        let mut demuxer = PsDemuxer::new();

        let pts_bytes = encode_pts(180000, 0x30); // PTS marker = 0x30
        let dts_bytes = encode_pts(90000, 0x10); // DTS marker = 0x10

        let mut data = vec![
            0x00, 0x00, 0x01, 0xE0, 0x00, 0x11, // length = 17
            0x80, 0xC0, 0x0A, // flags: PTS+DTS, header_data_len=10
        ];
        data.extend_from_slice(&pts_bytes);
        data.extend_from_slice(&dts_bytes);
        data.extend_from_slice(&[0xCA, 0xFE]); // payload

        data.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]);

        let packets = demuxer.feed(&data);
        assert_eq!(packets.len(), 1);
        assert_eq!(packets[0].pts, Some(180000));
        assert_eq!(packets[0].dts, Some(90000));
    }

    // --- Private stream 1 sub-stream extraction ---

    #[test]
    fn private_stream_1_ac3_substream() {
        let mut demuxer = PsDemuxer::new();

        // AC3 sub-header: sub_id(1) + frame_count(1) + access_unit_ptr(2) = 4 bytes
        let mut data = vec![
            0x00, 0x00, 0x01, 0xBD, // private stream 1
            0x00, 0x0B, // length = 11
            0x80, 0x00, 0x00, // no PTS, header_data_len=0
            0x80, // sub-stream ID: AC3 stream 0
            0x01, 0x00, 0x02, // frame_count + access_unit_ptr (sub-header bytes)
            0xAA, 0xBB, 0xCC, 0xDD, // AC3 payload
        ];

        data.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]);

        let packets = demuxer.feed(&data);
        assert_eq!(packets.len(), 1);
        assert_eq!(packets[0].stream_id, 0xBD);
        assert_eq!(packets[0].sub_stream_id, Some(0x80));
        assert_eq!(packets[0].data, vec![0xAA, 0xBB, 0xCC, 0xDD]);
    }

    #[test]
    fn private_stream_1_dts_substream() {
        let mut demuxer = PsDemuxer::new();

        // DTS sub-header: sub_id(1) + frame_count(1) + access_unit_ptr(2) = 4 bytes
        let mut data = vec![
            0x00, 0x00, 0x01, 0xBD, 0x00, 0x09, // length = 9
            0x80, 0x00, 0x00, // no PTS, header_data_len=0
            0x88, // sub-stream ID: DTS stream 0
            0x01, 0x00, 0x00, // sub-header (frame_count + access_unit_ptr)
            0x11, 0x22,
        ];
        data.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]);

        let packets = demuxer.feed(&data);
        assert_eq!(packets.len(), 1);
        assert_eq!(packets[0].sub_stream_id, Some(0x88));
        assert_eq!(packets[0].data, vec![0x11, 0x22]);
    }

    #[test]
    fn private_stream_1_subtitle_substream() {
        let mut demuxer = PsDemuxer::new();

        let mut data = vec![
            0x00, 0x00, 0x01, 0xBD, 0x00, 0x06, 0x80, 0x00, 0x00,
            0x20, // sub-stream ID: subtitle stream 0
            0xFF, 0xFE,
        ];
        data.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]);

        let packets = demuxer.feed(&data);
        assert_eq!(packets.len(), 1);
        assert_eq!(packets[0].sub_stream_id, Some(0x20));
    }

    #[test]
    fn private_stream_1_lpcm_substream() {
        let mut demuxer = PsDemuxer::new();

        // LPCM sub-header: sub_id(1) + frames(1) + ptr(2) + emphasis(1) + quant_freq(1) + channels(1) = 7 bytes
        let mut data = vec![
            0x00, 0x00, 0x01, 0xBD, 0x00, 0x0C, // length = 12
            0x80, 0x00, 0x00, // no PTS, header_data_len=0
            0xA0, // sub-stream ID: LPCM stream 0
            0x01, 0x00, 0x00, 0x00, 0x00, 0x00, // LPCM sub-header (6 bytes after sub_id)
            0x01, 0x02, // LPCM payload
        ];
        data.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]);

        let packets = demuxer.feed(&data);
        assert_eq!(packets.len(), 1);
        assert_eq!(packets[0].sub_stream_id, Some(0xA0));
        assert_eq!(packets[0].data, vec![0x01, 0x02]);
    }

    // --- Incremental feeding ---

    #[test]
    fn incremental_feed() {
        let mut demuxer = PsDemuxer::new();

        let mut full = vec![
            0x00, 0x00, 0x01, 0xE0, 0x00, 0x06, // length = 6
            0x80, 0x00, 0x00, // no PTS, header_data_len=0
            0xAA, 0xBB, 0xCC,
        ];
        full.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]);

        // Feed in two halves
        let mid = full.len() / 2;
        let p1 = demuxer.feed(&full[..mid]);
        assert!(p1.is_empty(), "first half should not produce packets");

        let p2 = demuxer.feed(&full[mid..]);
        assert_eq!(p2.len(), 1);
        assert_eq!(p2[0].data, vec![0xAA, 0xBB, 0xCC]);
    }

    // --- Multiple PES packets ---

    #[test]
    fn multiple_pes_packets() {
        let mut demuxer = PsDemuxer::new();

        let mut data = Vec::new();

        // First PES: video
        data.extend_from_slice(&[
            0x00, 0x00, 0x01, 0xE0, 0x00, 0x05, 0x80, 0x00, 0x00, 0x11, 0x22,
        ]);

        // Second PES: audio
        data.extend_from_slice(&[
            0x00, 0x00, 0x01, 0xC0, 0x00, 0x05, 0x80, 0x00, 0x00, 0x33, 0x44,
        ]);

        // Delimiter
        data.extend_from_slice(&[0x00, 0x00, 0x01, 0xB9]);

        let packets = demuxer.feed(&data);
        assert_eq!(packets.len(), 2);
        assert_eq!(packets[0].stream_id, 0xE0);
        assert_eq!(packets[1].stream_id, 0xC0);
    }

    // --- PTS parsing edge cases ---

    #[test]
    fn pts_zero() {
        // PTS = 0 encoded
        let pts = parse_pts(&encode_pts(0, 0x20));
        assert_eq!(pts, 0);
    }

    #[test]
    fn pts_large_value() {
        // Test a large PTS value (close to 33-bit max)
        let val: u64 = (1 << 32) - 1; // 0xFFFFFFFF
        let encoded = encode_pts(val, 0x20);
        let decoded = parse_pts(&encoded);
        assert_eq!(decoded, val);
    }

    // --- Helper: encode PTS for tests ---

    fn encode_pts(pts: u64, marker_prefix: u8) -> [u8; 5] {
        let mut buf = [0u8; 5];
        buf[0] = marker_prefix | (((pts >> 30) as u8) & 0x07) << 1 | 1;
        buf[1] = ((pts >> 22) & 0xFF) as u8;
        buf[2] = (((pts >> 15) & 0x7F) as u8) << 1 | 1;
        buf[3] = ((pts >> 7) & 0xFF) as u8;
        buf[4] = (((pts) & 0x7F) as u8) << 1 | 1;
        buf
    }
}