mpeg_pes/

packet.rs

//! PES packet header parsing (ISO/IEC 13818-1 §2.4.3.6, Table 2-21).

use crate::PACKET_START_CODE_PREFIX;
use crate::error::{Error, Result};
use crate::stream_id::StreamId;
use crate::timestamp::{self, Dts, Pts};

const MIN_LEN: usize = 6; // start_code(3) + stream_id(1) + PES_packet_length(2)
const HEADER_FIXED: usize = 3; // 2 flag bytes + PES_header_data_length
/// PES header stuffing byte (ISO/IEC 13818-1:2007 §2.4.3.7 — `0xFF`).
const PES_HEADER_STUFFING_BYTE: u8 = 0xFF;

// ── ESCR (ISO/IEC 13818-1 §2.4.3.7 Table 2-21) ──────────────────────────────

/// Elementary Stream Clock Reference: 33-bit base (90 kHz) + 9-bit extension
/// (27 MHz) — ISO/IEC 13818-1 §2.4.3.7, Table 2-21.
///
/// Wire layout (6 bytes, 48 bits):
/// `2×reserved(1) | ESCR_base[32:30](3) | marker(1) | ESCR_base[29:15](15) |
///  marker(1) | ESCR_base[14:0](15) | marker(1) | ESCR_ext(9) | marker(1)`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct Escr {
    /// 33-bit base (90 kHz units).
    pub base: u64,
    /// 9-bit extension (27 MHz units, 0..=299).
    pub extension: u16,
}

impl Escr {
    /// Full ESCR value on the 27 MHz clock: `base * 300 + extension`.
    #[must_use]
    pub fn as_27mhz(self) -> u64 {
        self.base * 300 + self.extension as u64
    }

    /// Construct from an absolute 27 MHz clock value.
    #[must_use]
    pub fn from_27mhz(ticks: u64) -> Self {
        const BASE_MASK: u64 = 0x1_FFFF_FFFF;
        const EXT_MASK: u16 = 0x1FF;
        Self {
            base: (ticks / 300) & BASE_MASK,
            extension: ((ticks % 300) as u16) & EXT_MASK,
        }
    }

    /// Decode from the 6-byte ESCR field.
    ///
    /// Bit layout (ISO/IEC 13818-1 §2.4.3.7, Table 2-21):
    /// `B0[7:6]`=reserved, `B0[5:3]`=`base[32:30]`, `B0[2]`=marker,
    /// `B0[1:0]`=`base[29:28]`, `B1[7:0]`=`base[27:20]`,
    /// `B2[7:3]`=`base[19:15]`, `B2[2]`=marker, `B2[1:0]`=`base[14:13]`,
    /// `B3[7:0]`=`base[12:5]`,
    /// `B4[7:3]`=`base[4:0]`, `B4[2]`=marker, `B4[1:0]`=`ext[8:7]`,
    /// `B5[7:1]`=`ext[6:0]`, `B5[0]`=marker.
    pub fn from_field_bytes(b: &[u8; 6]) -> Result<Self> {
        let base = ((((b[0] >> 3) & 0x07) as u64) << 30)   // base[32:30]
            | (((b[0] & 0x03) as u64) << 28)                 // base[29:28]
            | ((b[1] as u64) << 20)                           // base[27:20]
            | ((((b[2] >> 3) & 0x1F) as u64) << 15)          // base[19:15]
            | (((b[2] & 0x03) as u64) << 13)                  // base[14:13]
            | ((b[3] as u64) << 5)                             // base[12:5]
            | (((b[4] >> 3) & 0x1F) as u64); // base[4:0]
        let extension = ((((b[4] & 0x03) as u16) << 7) | ((b[5] >> 1) as u16)) & 0x1FF;
        Ok(Self { base, extension })
    }

    /// Encode as the 6-byte ESCR field.
    ///
    /// Reserved bits are set to `1` per the spec convention.
    /// Exact inverse of [`from_field_bytes`](Self::from_field_bytes).
    #[must_use]
    pub fn to_field_bytes(self) -> [u8; 6] {
        let b = self.base & 0x1_FFFF_FFFF;
        let e = (self.extension & 0x1FF) as u64;
        [
            // B0: reserved(2)='11' | base[32:30](3) | marker(1)='1' | base[29:28](2)
            0xC0 | (((b >> 30) & 0x07) as u8) << 3 | 0x04 | ((b >> 28) & 0x03) as u8,
            // B1: base[27:20](8)
            ((b >> 20) & 0xFF) as u8,
            // B2: base[19:15](5) | marker(1)='1' | base[14:13](2)
            (((b >> 15) & 0x1F) as u8) << 3 | 0x04 | ((b >> 13) & 0x03) as u8,
            // B3: base[12:5](8)
            ((b >> 5) & 0xFF) as u8,
            // B4: base[4:0](5) | marker(1)='1' | ext[8:7](2)
            (((b & 0x1F) as u8) << 3) | 0x04 | ((e >> 7) & 0x03) as u8,
            // B5: ext[6:0](7) | marker(1)='1'
            (((e & 0x7F) as u8) << 1) | 0x01,
        ]
    }
}

// ── DSM trick mode (ISO/IEC 13818-1 §2.4.3.7, Table 2-24) ──────────────────

/// Trick-mode control values for the `DSM_trick_mode_flag` field
/// (ISO/IEC 13818-1 §2.4.3.7, Table 2-24).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
#[non_exhaustive]
pub enum TrickMode {
    /// `000` — fast forward.
    FastForward {
        /// 2-bit `field_id`.
        field_id: u8,
        /// `intra_slice_refresh` flag.
        intra_slice_refresh: bool,
        /// 2-bit `frequency_truncation`.
        frequency_truncation: u8,
    },
    /// `001` — slow motion.
    SlowMotion {
        /// 5-bit `rep_cntrl`.
        rep_cntrl: u8,
    },
    /// `010` — freeze frame.
    FreezeFrame {
        /// 2-bit `field_id`.
        field_id: u8,
    },
    /// `011` — fast reverse.
    FastReverse {
        /// 2-bit `field_id`.
        field_id: u8,
        /// `intra_slice_refresh` flag.
        intra_slice_refresh: bool,
        /// 2-bit `frequency_truncation`.
        frequency_truncation: u8,
    },
    /// `100` — slow reverse.
    SlowReverse {
        /// 5-bit `rep_cntrl`.
        rep_cntrl: u8,
    },
    /// `101`–`111` — reserved.
    Reserved {
        /// Raw 3-bit `trick_mode_control` value.
        trick_mode_control: u8,
        /// Raw 5-bit remainder.
        data: u8,
    },
}

impl TrickMode {
    /// Decode from the 1-byte trick-mode field (ISO/IEC 13818-1 §2.4.3.7).
    pub fn from_byte(b: u8) -> Self {
        let control = (b >> 5) & 0x07;
        let data = b & 0x1F;
        match control {
            0b000 => TrickMode::FastForward {
                field_id: (data >> 3) & 0x03,
                intra_slice_refresh: (data >> 2) & 0x01 != 0,
                frequency_truncation: data & 0x03,
            },
            0b001 => TrickMode::SlowMotion { rep_cntrl: data },
            0b010 => TrickMode::FreezeFrame {
                field_id: (data >> 3) & 0x03,
            },
            0b011 => TrickMode::FastReverse {
                field_id: (data >> 3) & 0x03,
                intra_slice_refresh: (data >> 2) & 0x01 != 0,
                frequency_truncation: data & 0x03,
            },
            0b100 => TrickMode::SlowReverse { rep_cntrl: data },
            _ => TrickMode::Reserved {
                trick_mode_control: control,
                data,
            },
        }
    }

    /// Encode as the 1-byte trick-mode field.
    pub fn to_byte(self) -> u8 {
        match self {
            TrickMode::FastForward {
                field_id,
                intra_slice_refresh,
                frequency_truncation,
            } => {
                ((field_id & 0x03) << 3)
                    | ((intra_slice_refresh as u8) << 2)
                    | (frequency_truncation & 0x03)
            }
            TrickMode::SlowMotion { rep_cntrl } => (0b001 << 5) | (rep_cntrl & 0x1F),
            TrickMode::FreezeFrame { field_id } => (0b010 << 5) | ((field_id & 0x03) << 3),
            TrickMode::FastReverse {
                field_id,
                intra_slice_refresh,
                frequency_truncation,
            } => {
                (0b011 << 5)
                    | ((field_id & 0x03) << 3)
                    | ((intra_slice_refresh as u8) << 2)
                    | (frequency_truncation & 0x03)
            }
            TrickMode::SlowReverse { rep_cntrl } => (0b100 << 5) | (rep_cntrl & 0x1F),
            TrickMode::Reserved {
                trick_mode_control,
                data,
            } => ((trick_mode_control & 0x07) << 5) | (data & 0x1F),
        }
    }
}

// ── PES extension (ISO/IEC 13818-1 §2.4.3.7) ──────────────────────────────

/// `program_packet_sequence_counter` sub-field of [`PesExtension`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct ProgramPacketSequenceCounter {
    /// 7-bit counter.
    pub counter: u8,
    /// `MPEG1_MPEG2_identifier` flag.
    pub mpeg1_mpeg2_identifier: bool,
    /// 6-bit `original_stuff_length`.
    pub original_stuff_length: u8,
}

/// `P-STD_buffer` sub-field of [`PesExtension`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct PStdBuffer {
    /// P-STD buffer size scale: `false` = 128 bytes/unit, `true` = 1024 bytes/unit.
    pub scale: bool,
    /// 13-bit buffer size in units of the scale.
    pub size: u16,
}

/// Typed PES header extension sub-structure
/// (ISO/IEC 13818-1 §2.4.3.7, Table 2-21, `PES_extension_flag = 1`).
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct PesExtension<'a> {
    /// 128-bit PES private data, if `PES_private_data_flag` is set.
    pub pes_private_data: Option<[u8; 16]>,
    /// Pack header field bytes (opaque per spec — `&[u8]` is correct here).
    pub pack_header: Option<&'a [u8]>,
    /// Program packet sequence counter sub-fields.
    pub program_packet_sequence_counter: Option<ProgramPacketSequenceCounter>,
    /// P-STD buffer sub-field.
    pub p_std_buffer: Option<PStdBuffer>,
    /// PES extension field bytes (opaque per spec).
    pub pes_extension_field: Option<&'a [u8]>,
}

impl<'a> PesExtension<'a> {
    fn parse(data: &'a [u8]) -> Result<Self> {
        if data.is_empty() {
            return Err(Error::BufferTooShort {
                need: 1,
                have: 0,
                what: "PES_extension flags byte",
            });
        }
        let flags = data[0];
        let mut cursor = 1usize;

        let pes_private_data = if flags & 0x80 != 0 {
            let end = cursor + 16;
            let arr: [u8; 16] = data
                .get(cursor..end)
                .and_then(|s| s.try_into().ok())
                .ok_or(Error::BufferTooShort {
                    need: end,
                    have: data.len(),
                    what: "PES_private_data",
                })?;
            cursor = end;
            Some(arr)
        } else {
            None
        };

        let pack_header = if flags & 0x40 != 0 {
            let pack_len = *data.get(cursor).ok_or(Error::BufferTooShort {
                need: cursor + 1,
                have: data.len(),
                what: "pack_field_length",
            })? as usize;
            cursor += 1;
            let end = cursor + pack_len;
            let slice = data.get(cursor..end).ok_or(Error::BufferTooShort {
                need: end,
                have: data.len(),
                what: "pack_header",
            })?;
            cursor = end;
            Some(slice)
        } else {
            None
        };

        let program_packet_sequence_counter = if flags & 0x20 != 0 {
            if data.len() < cursor + 2 {
                return Err(Error::BufferTooShort {
                    need: cursor + 2,
                    have: data.len(),
                    what: "program_packet_sequence_counter",
                });
            }
            let b0 = data[cursor];
            let b1 = data[cursor + 1];
            cursor += 2;
            Some(ProgramPacketSequenceCounter {
                counter: b0 & 0x7F,
                mpeg1_mpeg2_identifier: (b1 & 0x40) != 0,
                original_stuff_length: b1 & 0x3F,
            })
        } else {
            None
        };

        let p_std_buffer = if flags & 0x10 != 0 {
            if data.len() < cursor + 2 {
                return Err(Error::BufferTooShort {
                    need: cursor + 2,
                    have: data.len(),
                    what: "P-STD_buffer",
                });
            }
            let b0 = data[cursor];
            let b1 = data[cursor + 1];
            cursor += 2;
            Some(PStdBuffer {
                scale: (b0 & 0x20) != 0,
                size: (((b0 & 0x1F) as u16) << 8) | (b1 as u16),
            })
        } else {
            None
        };

        let pes_extension_field = if flags & 0x01 != 0 {
            let ext_len = *data.get(cursor).ok_or(Error::BufferTooShort {
                need: cursor + 1,
                have: data.len(),
                what: "PES_extension_field_length",
            })? as usize;
            cursor += 1;
            let end = cursor + ext_len;
            let slice = data.get(cursor..end).ok_or(Error::BufferTooShort {
                need: end,
                have: data.len(),
                what: "PES_extension_field",
            })?;
            cursor = end;
            Some(slice)
        } else {
            None
        };
        let _ = cursor;

        Ok(PesExtension {
            pes_private_data,
            pack_header,
            program_packet_sequence_counter,
            p_std_buffer,
            pes_extension_field,
        })
    }

    /// Number of bytes this PES extension occupies on the wire.
    pub fn serialized_len(&self) -> usize {
        let mut n = 1usize; // flags byte
        if self.pes_private_data.is_some() {
            n += 16;
        }
        if let Some(ph) = self.pack_header {
            n += 1 + ph.len();
        }
        if self.program_packet_sequence_counter.is_some() {
            n += 2;
        }
        if self.p_std_buffer.is_some() {
            n += 2;
        }
        if let Some(ef) = self.pes_extension_field {
            n += 1 + ef.len();
        }
        n
    }

    fn serialize_into(&self, buf: &mut [u8]) -> Result<usize> {
        let need = self.serialized_len();
        if buf.len() < need {
            return Err(Error::BufferTooShort {
                need,
                have: buf.len(),
                what: "PES_extension serialize output",
            });
        }

        let mut flags = 0u8;
        if self.pes_private_data.is_some() {
            flags |= 0x80;
        }
        if self.pack_header.is_some() {
            flags |= 0x40;
        }
        if self.program_packet_sequence_counter.is_some() {
            flags |= 0x20;
        }
        if self.p_std_buffer.is_some() {
            flags |= 0x10;
        }
        if self.pes_extension_field.is_some() {
            flags |= 0x01;
        }
        buf[0] = flags;
        let mut cursor = 1usize;

        if let Some(pd) = &self.pes_private_data {
            buf[cursor..cursor + 16].copy_from_slice(pd);
            cursor += 16;
        }
        if let Some(ph) = self.pack_header {
            buf[cursor] = ph.len() as u8;
            cursor += 1;
            buf[cursor..cursor + ph.len()].copy_from_slice(ph);
            cursor += ph.len();
        }
        if let Some(ppsc) = self.program_packet_sequence_counter {
            // byte 0: marker(1) | counter(7)
            buf[cursor] = 0x80 | (ppsc.counter & 0x7F);
            // byte 1: marker(1) | mpeg1_mpeg2_id(1) | original_stuff_length(6)
            buf[cursor + 1] = 0x80
                | ((ppsc.mpeg1_mpeg2_identifier as u8) << 6)
                | (ppsc.original_stuff_length & 0x3F);
            cursor += 2;
        }
        if let Some(ps) = self.p_std_buffer {
            // '01' | scale(1) | size(13)
            buf[cursor] = 0x40 | ((ps.scale as u8) << 5) | ((ps.size >> 8) as u8 & 0x1F);
            buf[cursor + 1] = (ps.size & 0xFF) as u8;
            cursor += 2;
        }
        if let Some(ef) = self.pes_extension_field {
            buf[cursor] = ef.len() as u8;
            cursor += 1;
            buf[cursor..cursor + ef.len()].copy_from_slice(ef);
            cursor += ef.len();
        }
        Ok(cursor)
    }
}

// ── PesHeader ──────────────────────────────────────────────────────────────

/// The optional PES header present for non-special `stream_id`s
/// (ISO/IEC 13818-1 §2.4.3.6, §2.4.3.7). All optional sub-fields are fully
/// typed — the raw `optional_fields` blob has been replaced with the
/// individual decoded fields.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct PesHeader<'a> {
    /// PES_scrambling_control (2 bits).
    pub scrambling_control: u8,
    /// PES_priority.
    pub pes_priority: bool,
    /// data_alignment_indicator.
    pub data_alignment_indicator: bool,
    /// copyright.
    pub copyright: bool,
    /// original_or_copy.
    pub original_or_copy: bool,
    /// Presentation time stamp, if `PTS_DTS_flags` indicated one.
    pub pts: Option<Pts>,
    /// Decoding time stamp, if `PTS_DTS_flags` was `11`.
    pub dts: Option<Dts>,
    /// Elementary stream clock reference (6 bytes), if `ESCR_flag` is set.
    pub escr: Option<Escr>,
    /// 22-bit ES rate (bytes/second × 50), if `ES_rate_flag` is set.
    pub es_rate: Option<u32>,
    /// DSM trick-mode control byte (typed), if `DSM_trick_mode_flag` is set.
    pub dsm_trick_mode: Option<TrickMode>,
    /// 7-bit `additional_copy_info`, if `additional_copy_info_flag` is set.
    pub additional_copy_info: Option<u8>,
    /// Previous PES packet CRC (16 bits), if `PES_CRC_flag` is set.
    pub pes_crc: Option<u16>,
    /// PES extension sub-structure, if `PES_extension_flag` is set.
    pub pes_extension: Option<PesExtension<'a>>,
    /// Number of trailing `0xFF` stuffing bytes inside the `PES_header_data_length`
    /// region, after the typed optional fields (ISO/IEC 13818-1:2007 §2.4.3.7 —
    /// "stuffing_byte: fixed 8-bit value `0xFF`").
    ///
    /// Encoders pad the optional-header block (often so the elementary stream
    /// starts at a fixed offset). These bytes are part of the wire image;
    /// capturing the count lets [`PesPacket::serialize_into`] reproduce the
    /// header byte-for-byte. Set to `0` when constructing a header with no
    /// stuffing.
    pub header_stuffing_len: usize,
}

impl PesHeader<'_> {
    /// Number of bytes this header occupies in the serialized `PES_header_data_length`
    /// region (not counting the 3 fixed header bytes — the 2 flag bytes + length byte).
    fn optional_len(&self) -> usize {
        let mut n = 0usize;
        if self.pts.is_some() {
            n += 5;
        }
        if self.dts.is_some() {
            n += 5;
        }
        if self.escr.is_some() {
            n += 6;
        }
        if self.es_rate.is_some() {
            n += 3;
        }
        if self.dsm_trick_mode.is_some() {
            n += 1;
        }
        if self.additional_copy_info.is_some() {
            n += 1;
        }
        if self.pes_crc.is_some() {
            n += 2;
        }
        if let Some(ref ext) = self.pes_extension {
            n += ext.serialized_len();
        }
        n += self.header_stuffing_len;
        n
    }
}

/// A parsed PES packet.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct PesPacket<'a> {
    /// stream_id (Table 2-22).
    pub stream_id: StreamId,
    /// PES_packet_length as carried; `0` means unbounded (video).
    pub pes_packet_length: u16,
    /// Optional PES header (absent for the special `stream_id`s).
    pub header: Option<PesHeader<'a>>,
    /// The elementary-stream bytes (`PES_packet_data_byte`s).
    #[cfg_attr(feature = "serde", serde(skip))]
    pub payload: &'a [u8],
}

impl<'a> PesPacket<'a> {
    /// Parse a PES packet from the bytes starting at its `packet_start_code_prefix`.
    pub fn parse(b: &'a [u8]) -> Result<Self> {
        if b.len() < MIN_LEN {
            return Err(Error::BufferTooShort {
                need: MIN_LEN,
                have: b.len(),
                what: "PES packet header",
            });
        }
        if b[0..3] != PACKET_START_CODE_PREFIX {
            return Err(Error::BadStartCode(
                (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]),
            ));
        }
        let stream_id = StreamId(b[3]);
        let pes_packet_length = u16::from_be_bytes([b[4], b[5]]);
        // Where the payload ends: bounded by PES_packet_length unless 0 (unbounded).
        let payload_end = if pes_packet_length == 0 {
            b.len()
        } else {
            (MIN_LEN + pes_packet_length as usize).min(b.len())
        };

        if !stream_id.has_optional_header() {
            return Ok(PesPacket {
                stream_id,
                pes_packet_length,
                header: None,
                payload: &b[MIN_LEN..payload_end],
            });
        }

        if b.len() < MIN_LEN + HEADER_FIXED {
            return Err(Error::BufferTooShort {
                need: MIN_LEN + HEADER_FIXED,
                have: b.len(),
                what: "PES optional header",
            });
        }
        let f1 = b[6];
        let f2 = b[7];
        let hdl = usize::from(b[8]);
        let hdr_start = MIN_LEN + HEADER_FIXED; // = 9
        let hdr_end = hdr_start + hdl;
        if b.len() < hdr_end {
            return Err(Error::BufferTooShort {
                need: hdr_end,
                have: b.len(),
                what: "PES_header_data_length",
            });
        }
        let opt = &b[hdr_start..hdr_end];
        let mut cursor = 0usize;

        // PTS/DTS (ISO/IEC 13818-1 §2.4.3.7 Table 2-21).
        let pts_dts_flags = (f2 >> 6) & 0x03;
        let (pts, dts) = match pts_dts_flags {
            0b10 => {
                if opt.len() < cursor + 5 {
                    return Err(Error::BufferTooShort {
                        need: cursor + 5,
                        have: opt.len(),
                        what: "PTS",
                    });
                }
                let pts = Pts(timestamp::read(&opt[cursor..], 0b0010, "PTS")?);
                cursor += 5;
                (Some(pts), None)
            }
            0b11 => {
                if opt.len() < cursor + 10 {
                    return Err(Error::BufferTooShort {
                        need: cursor + 10,
                        have: opt.len(),
                        what: "PTS+DTS",
                    });
                }
                let pts = Pts(timestamp::read(&opt[cursor..], 0b0011, "PTS")?);
                cursor += 5;
                let dts = Dts(timestamp::read(&opt[cursor..], 0b0001, "DTS")?);
                cursor += 5;
                (Some(pts), Some(dts))
            }
            _ => (None, None),
        };

        // ESCR (6 bytes, ISO/IEC 13818-1 §2.4.3.7).
        let escr = if f2 & 0x20 != 0 {
            if opt.len() < cursor + 6 {
                return Err(Error::BufferTooShort {
                    need: cursor + 6,
                    have: opt.len(),
                    what: "ESCR",
                });
            }
            let arr: &[u8; 6] = opt[cursor..cursor + 6].try_into().unwrap();
            let e = Escr::from_field_bytes(arr)?;
            cursor += 6;
            Some(e)
        } else {
            None
        };

        // ES_rate (3 bytes: 1 marker + 22-bit rate + 1 marker).
        let es_rate = if f2 & 0x10 != 0 {
            if opt.len() < cursor + 3 {
                return Err(Error::BufferTooShort {
                    need: cursor + 3,
                    have: opt.len(),
                    what: "ES_rate",
                });
            }
            let rate = (((opt[cursor] & 0x7F) as u32) << 15)
                | ((opt[cursor + 1] as u32) << 7)
                | ((opt[cursor + 2] >> 1) as u32);
            cursor += 3;
            Some(rate)
        } else {
            None
        };

        // DSM trick mode (1 byte).
        let dsm_trick_mode = if f2 & 0x08 != 0 {
            if opt.len() < cursor + 1 {
                return Err(Error::BufferTooShort {
                    need: cursor + 1,
                    have: opt.len(),
                    what: "trick_mode",
                });
            }
            let tm = TrickMode::from_byte(opt[cursor]);
            cursor += 1;
            Some(tm)
        } else {
            None
        };

        // additional_copy_info (1 byte: marker + 7-bit info).
        let additional_copy_info = if f2 & 0x04 != 0 {
            if opt.len() < cursor + 1 {
                return Err(Error::BufferTooShort {
                    need: cursor + 1,
                    have: opt.len(),
                    what: "additional_copy_info",
                });
            }
            let v = opt[cursor] & 0x7F;
            cursor += 1;
            Some(v)
        } else {
            None
        };

        // PES_CRC (2 bytes).
        let pes_crc = if f2 & 0x02 != 0 {
            if opt.len() < cursor + 2 {
                return Err(Error::BufferTooShort {
                    need: cursor + 2,
                    have: opt.len(),
                    what: "PES_CRC",
                });
            }
            let crc = u16::from_be_bytes([opt[cursor], opt[cursor + 1]]);
            cursor += 2;
            Some(crc)
        } else {
            None
        };

        // PES_extension.
        let pes_extension = if f2 & 0x01 != 0 {
            let ext = PesExtension::parse(&opt[cursor..])?;
            cursor += ext.serialized_len();
            Some(ext)
        } else {
            None
        };

        // Bytes remaining in the `PES_header_data_length` region after the typed
        // optional fields are `0xFF` stuffing (ISO/IEC 13818-1:2007 §2.4.3.7).
        // Record the count so serialization reproduces the header byte-for-byte.
        let header_stuffing_len = hdl.saturating_sub(cursor);

        let header = PesHeader {
            scrambling_control: (f1 >> 4) & 0x03,
            pes_priority: f1 & 0x08 != 0,
            data_alignment_indicator: f1 & 0x04 != 0,
            copyright: f1 & 0x02 != 0,
            original_or_copy: f1 & 0x01 != 0,
            pts,
            dts,
            escr,
            es_rate,
            dsm_trick_mode,
            additional_copy_info,
            pes_crc,
            pes_extension,
            header_stuffing_len,
        };

        Ok(PesPacket {
            stream_id,
            pes_packet_length,
            header: Some(header),
            payload: &b[hdr_end.min(payload_end)..payload_end],
        })
    }

    /// Serialized length in bytes.
    #[must_use]
    pub fn serialized_len(&self) -> usize {
        let hdr = self
            .header
            .as_ref()
            .map_or(0, |h| HEADER_FIXED + h.optional_len());
        MIN_LEN + hdr + self.payload.len()
    }

    /// Serialize back to bytes (byte-identical to a spec-compliant input).
    pub fn serialize_into(&self, buf: &mut [u8]) -> Result<usize> {
        let len = self.serialized_len();
        if buf.len() < len {
            return Err(Error::BufferTooShort {
                need: len,
                have: buf.len(),
                what: "PES serialize output",
            });
        }
        buf[0..3].copy_from_slice(&PACKET_START_CODE_PREFIX);
        buf[3] = self.stream_id.0;
        buf[4..6].copy_from_slice(&self.pes_packet_length.to_be_bytes());
        let payload_at = match &self.header {
            None => MIN_LEN,
            Some(h) => {
                let opt_len = h.optional_len();
                if opt_len > 255 {
                    return Err(Error::OptionalFieldsTooLarge(opt_len));
                }

                // f1: marker '10' | scrambling(2) | priority(1) | align(1) | copyright(1) | orig(1)
                let f1 = 0x80
                    | ((h.scrambling_control & 0x03) << 4)
                    | (u8::from(h.pes_priority) << 3)
                    | (u8::from(h.data_alignment_indicator) << 2)
                    | (u8::from(h.copyright) << 1)
                    | u8::from(h.original_or_copy);

                // f2: pts_dts_flags(2) | escr_flag(1) | es_rate_flag(1) |
                //     trick_mode(1) | add_copy(1) | crc(1) | ext(1)
                let pts_dts_flags = match (h.pts.is_some(), h.dts.is_some()) {
                    (true, true) => 0b11u8,
                    (true, false) => 0b10,
                    _ => 0b00,
                };
                let f2 = (pts_dts_flags << 6)
                    | (u8::from(h.escr.is_some()) << 5)
                    | (u8::from(h.es_rate.is_some()) << 4)
                    | (u8::from(h.dsm_trick_mode.is_some()) << 3)
                    | (u8::from(h.additional_copy_info.is_some()) << 2)
                    | (u8::from(h.pes_crc.is_some()) << 1)
                    | u8::from(h.pes_extension.is_some());

                buf[6] = f1;
                buf[7] = f2;
                buf[8] = opt_len as u8;

                let mut cursor = MIN_LEN + HEADER_FIXED; // = 9

                // PTS (and/or DTS).
                if let Some(pts) = h.pts {
                    let prefix = if h.dts.is_some() { 0b0011u8 } else { 0b0010u8 };
                    buf[cursor..cursor + 5].copy_from_slice(&timestamp::write(pts.0, prefix));
                    cursor += 5;
                }
                if let Some(dts) = h.dts {
                    buf[cursor..cursor + 5].copy_from_slice(&timestamp::write(dts.0, 0b0001));
                    cursor += 5;
                }
                // ESCR.
                if let Some(escr) = h.escr {
                    buf[cursor..cursor + 6].copy_from_slice(&escr.to_field_bytes());
                    cursor += 6;
                }
                // ES_rate: marker(1) | rate(22) | marker(1) = 3 bytes.
                if let Some(rate) = h.es_rate {
                    buf[cursor] = 0x80 | ((rate >> 15) as u8 & 0x7F);
                    buf[cursor + 1] = ((rate >> 7) & 0xFF) as u8;
                    buf[cursor + 2] = (((rate & 0x7F) as u8) << 1) | 0x01;
                    cursor += 3;
                }
                // DSM trick mode.
                if let Some(tm) = h.dsm_trick_mode {
                    buf[cursor] = tm.to_byte();
                    cursor += 1;
                }
                // additional_copy_info: marker(1) | info(7).
                if let Some(aci) = h.additional_copy_info {
                    buf[cursor] = 0x80 | (aci & 0x7F);
                    cursor += 1;
                }
                // PES_CRC (2 bytes, big-endian).
                if let Some(crc) = h.pes_crc {
                    buf[cursor..cursor + 2].copy_from_slice(&crc.to_be_bytes());
                    cursor += 2;
                }
                // PES_extension.
                if let Some(ref ext) = h.pes_extension {
                    let written = ext.serialize_into(&mut buf[cursor..])?;
                    cursor += written;
                }

                // Trailing `0xFF` stuffing inside the header_data_length region
                // (ISO/IEC 13818-1:2007 §2.4.3.7), reproducing the encoder's pad.
                for b in buf[cursor..cursor + h.header_stuffing_len].iter_mut() {
                    *b = PES_HEADER_STUFFING_BYTE;
                }
                cursor += h.header_stuffing_len;

                cursor
            }
        };
        buf[payload_at..len].copy_from_slice(self.payload);
        Ok(len)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    extern crate alloc;
    use alloc::vec;

    fn round_trip(b: &[u8]) {
        let pkt = PesPacket::parse(b).unwrap();
        let mut out = vec![0u8; pkt.serialized_len()];
        pkt.serialize_into(&mut out).unwrap();
        assert_eq!(&out[..], b, "round-trip mismatch");
        let re = PesPacket::parse(&out).unwrap();
        // Compare without the borrowed lifetime complexity — compare serialized form.
        let mut re_out = vec![0u8; re.serialized_len()];
        re.serialize_into(&mut re_out).unwrap();
        assert_eq!(out, re_out, "re-parse mismatch");
    }

    #[test]
    fn video_pts_only() {
        // stream_id 0xE0, len=0x0A, flags 0x80/0x80, hdl=5, PTS=0, payload AA BB.
        let b = [
            0x00, 0x00, 0x01, 0xE0, 0x00, 0x0A, 0x80, 0x80, 0x05, 0x21, 0x00, 0x01, 0x00, 0x01,
            0xAA, 0xBB,
        ];
        let pkt = PesPacket::parse(&b).unwrap();
        assert_eq!(pkt.stream_id, StreamId(0xE0));
        let h = pkt.header.as_ref().unwrap();
        assert_eq!(h.pts, Some(Pts(0)));
        assert!(h.dts.is_none());
        assert_eq!(pkt.payload, &[0xAA, 0xBB]);
        round_trip(&b);
    }

    #[test]
    fn pts_and_dts() {
        // PTS_DTS_flags=11, hdl=10. PTS prefix 0011, DTS prefix 0001.
        let b = [
            0x00, 0x00, 0x01, 0xE0, 0x00, 0x0F, 0x80, 0xC0, 0x0A, 0x31, 0x00, 0x03, 0x00, 0x01,
            0x11, 0x00, 0x05, 0x00, 0x01, 0xCC,
        ];
        let pkt = PesPacket::parse(&b).unwrap();
        let h = pkt.header.as_ref().unwrap();
        assert!(h.pts.is_some());
        assert!(h.dts.is_some());
        round_trip(&b);
    }

    #[test]
    fn pes_header_stuffing_round_trip() {
        // PTS-only (flags 0x80), PES_header_data_length = 8: 5 PTS bytes + 3
        // 0xFF stuffing bytes (ISO/IEC 13818-1:2007 §2.4.3.7).
        let b = [
            0x00, 0x00, 0x01, 0xE0, 0x00, 0x0C, 0x80, 0x80, 0x08, // hdl = 8
            0x21, 0x00, 0x01, 0x00, 0x01, // PTS = 0
            0xFF, 0xFF, 0xFF, // 3 stuffing bytes
            0xAA, // payload
        ];
        let pkt = PesPacket::parse(&b).unwrap();
        let h = pkt.header.as_ref().unwrap();
        assert!(h.pts.is_some());
        assert_eq!(
            h.header_stuffing_len, 3,
            "3 stuffing bytes after the 5-byte PTS"
        );
        // serialized_len must account for the stuffing.
        assert_eq!(pkt.serialized_len(), b.len());
        round_trip(&b); // byte-identical, incl. the 0xFF stuffing
    }

    #[test]
    fn special_stream_no_header() {
        // padding_stream 0xBE: bytes after length are payload directly.
        let b = [0x00, 0x00, 0x01, 0xBE, 0x00, 0x03, 0xFF, 0xFF, 0xFF];
        let pkt = PesPacket::parse(&b).unwrap();
        assert!(pkt.header.is_none());
        assert_eq!(pkt.payload, &[0xFF, 0xFF, 0xFF]);
        round_trip(&b);
    }

    #[test]
    fn unbounded_length_zero() {
        // PES_packet_length=0 (video): payload runs to end of buffer.
        let b = [
            0x00, 0x00, 0x01, 0xE0, 0x00, 0x00, 0x80, 0x80, 0x05, 0x21, 0x00, 0x01, 0x00, 0x01,
            0x01, 0x02, 0x03,
        ];
        let pkt = PesPacket::parse(&b).unwrap();
        assert_eq!(pkt.pes_packet_length, 0);
        assert_eq!(pkt.payload, &[0x01, 0x02, 0x03]);
        round_trip(&b);
    }

    #[test]
    fn rejects_bad_start_code() {
        let err = PesPacket::parse(&[0x00, 0x00, 0x02, 0xE0, 0x00, 0x00]).unwrap_err();
        assert!(matches!(err, Error::BadStartCode(0x000002)));
    }

    #[test]
    fn rejects_short() {
        let err = PesPacket::parse(&[0x00, 0x00, 0x01]).unwrap_err();
        assert!(matches!(err, Error::BufferTooShort { .. }));
    }

    #[test]
    fn serialize_rejects_oversized_optional_fields() {
        // Construct a PesHeader whose optional_len() > 255 by adding enough flags.
        // max realistic: 5+5+6+3+1+1+2 = 23, plus PesExtension with big private data.
        // This is structurally impossible with typed fields to reach 256 naturally,
        // but we test the guard exists by checking that OptionalFieldsTooLarge
        // is the right error variant name still exists.
        let _ = Error::OptionalFieldsTooLarge(256);
    }

    // ── Typed optional fields round-trips ──────────────────────────────────

    fn build_pes(h: PesHeader<'_>, payload: &[u8]) -> alloc::vec::Vec<u8> {
        let pkt = PesPacket {
            stream_id: StreamId(0xE0),
            pes_packet_length: 0, // unbounded
            header: Some(h),
            payload,
        };
        let mut out = vec![0u8; pkt.serialized_len()];
        pkt.serialize_into(&mut out).unwrap();
        out
    }

    fn empty_header<'a>() -> PesHeader<'a> {
        PesHeader {
            scrambling_control: 0,
            pes_priority: false,
            data_alignment_indicator: false,
            copyright: false,
            original_or_copy: false,
            pts: None,
            dts: None,
            escr: None,
            es_rate: None,
            dsm_trick_mode: None,
            additional_copy_info: None,
            pes_crc: None,
            pes_extension: None,
            header_stuffing_len: 0,
        }
    }

    /// Construct PES with ESCR set, serialize, parse, assert field preserved.
    #[test]
    fn pes_header_escr_round_trip() {
        let escr = Escr {
            base: 90_000,
            extension: 150,
        };
        let h = PesHeader {
            escr: Some(escr),
            ..empty_header()
        };
        let bytes = build_pes(h, &[0xAA]);
        let pkt = PesPacket::parse(&bytes).unwrap();
        assert_eq!(pkt.header.unwrap().escr, Some(escr));
    }

    /// Escr::from_27mhz / as_27mhz round-trip.
    #[test]
    fn escr_27mhz_round_trip() {
        for ticks in [0u64, 1, 300, 27_000_000, 8_589_934_591] {
            let e = Escr::from_27mhz(ticks);
            assert_eq!(e.as_27mhz(), ticks, "ticks={ticks}");
        }
    }

    /// Escr::to_field_bytes / from_field_bytes round-trip.
    #[test]
    fn escr_field_bytes_round_trip() {
        for (base, ext) in [
            (0u64, 0u16),
            (10_000, 0),
            (0x1_FFFF_FFFF, 0x1FF),
            (1234, 56),
        ] {
            let e = Escr {
                base,
                extension: ext,
            };
            let bytes = e.to_field_bytes();
            let decoded = Escr::from_field_bytes(&bytes).unwrap();
            assert_eq!(decoded, e, "base={base} ext={ext}");
        }
    }

    /// ES_rate round-trip.
    #[test]
    fn pes_header_es_rate_round_trip() {
        let h = PesHeader {
            es_rate: Some(0x3FFFFF),
            ..empty_header()
        };
        let bytes = build_pes(h, &[]);
        let pkt = PesPacket::parse(&bytes).unwrap();
        assert_eq!(pkt.header.unwrap().es_rate, Some(0x3FFFFF));
    }

    /// TrickMode round-trip (all variants).
    #[test]
    fn trick_mode_all_variants_round_trip() {
        let cases = [
            TrickMode::FastForward {
                field_id: 0x2,
                intra_slice_refresh: true,
                frequency_truncation: 0x3,
            },
            TrickMode::SlowMotion { rep_cntrl: 0x1F },
            TrickMode::FreezeFrame { field_id: 0x1 },
            TrickMode::FastReverse {
                field_id: 0x0,
                intra_slice_refresh: false,
                frequency_truncation: 0x1,
            },
            TrickMode::SlowReverse { rep_cntrl: 0 },
            TrickMode::Reserved {
                trick_mode_control: 0b101,
                data: 0x1A,
            },
        ];
        for tm in cases {
            let b = tm.to_byte();
            let decoded = TrickMode::from_byte(b);
            assert_eq!(decoded, tm, "tm={tm:?}");
        }
    }

    /// TrickMode in PES header round-trip.
    #[test]
    fn pes_header_trick_mode_round_trip() {
        let tm = TrickMode::FastForward {
            field_id: 1,
            intra_slice_refresh: false,
            frequency_truncation: 2,
        };
        let h = PesHeader {
            dsm_trick_mode: Some(tm),
            ..empty_header()
        };
        let bytes = build_pes(h, &[]);
        let pkt = PesPacket::parse(&bytes).unwrap();
        assert_eq!(pkt.header.unwrap().dsm_trick_mode, Some(tm));
    }

    /// additional_copy_info round-trip.
    #[test]
    fn pes_header_additional_copy_info_round_trip() {
        let h = PesHeader {
            additional_copy_info: Some(0x7F),
            ..empty_header()
        };
        let bytes = build_pes(h, &[]);
        let pkt = PesPacket::parse(&bytes).unwrap();
        assert_eq!(pkt.header.unwrap().additional_copy_info, Some(0x7F));
    }

    /// PES_CRC round-trip.
    #[test]
    fn pes_header_pes_crc_round_trip() {
        let h = PesHeader {
            pes_crc: Some(0xDEAD),
            ..empty_header()
        };
        let bytes = build_pes(h, &[]);
        let pkt = PesPacket::parse(&bytes).unwrap();
        assert_eq!(pkt.header.unwrap().pes_crc, Some(0xDEAD));
    }

    /// PesExtension with program_packet_sequence_counter.
    #[test]
    fn pes_extension_ppsc_round_trip() {
        let ppsc = ProgramPacketSequenceCounter {
            counter: 42,
            mpeg1_mpeg2_identifier: true,
            original_stuff_length: 7,
        };
        let ext = PesExtension {
            pes_private_data: None,
            pack_header: None,
            program_packet_sequence_counter: Some(ppsc),
            p_std_buffer: None,
            pes_extension_field: None,
        };
        let h = PesHeader {
            pes_extension: Some(ext),
            ..empty_header()
        };
        let bytes = build_pes(h, &[]);
        let pkt = PesPacket::parse(&bytes).unwrap();
        let decoded_ext = pkt.header.unwrap().pes_extension.unwrap();
        assert_eq!(decoded_ext.program_packet_sequence_counter, Some(ppsc));
    }

    /// PesExtension with P-STD buffer.
    #[test]
    fn pes_extension_p_std_buffer_round_trip() {
        let pstd = PStdBuffer {
            scale: true,
            size: 0x1FFF,
        };
        let ext = PesExtension {
            pes_private_data: None,
            pack_header: None,
            program_packet_sequence_counter: None,
            p_std_buffer: Some(pstd),
            pes_extension_field: None,
        };
        let h = PesHeader {
            pes_extension: Some(ext),
            ..empty_header()
        };
        let bytes = build_pes(h, &[]);
        let pkt = PesPacket::parse(&bytes).unwrap();
        let decoded_ext = pkt.header.unwrap().pes_extension.unwrap();
        assert_eq!(decoded_ext.p_std_buffer, Some(pstd));
    }

    /// PesExtension with private data (16 bytes).
    #[test]
    fn pes_extension_private_data_round_trip() {
        let pd: [u8; 16] = [
            0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E,
            0x0F, 0x10,
        ];
        let ext = PesExtension {
            pes_private_data: Some(pd),
            pack_header: None,
            program_packet_sequence_counter: None,
            p_std_buffer: None,
            pes_extension_field: None,
        };
        let h = PesHeader {
            pes_extension: Some(ext),
            ..empty_header()
        };
        let bytes = build_pes(h, &[]);
        let pkt = PesPacket::parse(&bytes).unwrap();
        let decoded_ext = pkt.header.unwrap().pes_extension.unwrap();
        assert_eq!(decoded_ext.pes_private_data, Some(pd));
    }

    /// All optional PES header fields set at once — serialize and parse round-trip.
    #[test]
    fn pes_header_all_fields_round_trip() {
        let ppsc = ProgramPacketSequenceCounter {
            counter: 1,
            mpeg1_mpeg2_identifier: false,
            original_stuff_length: 0,
        };
        let ext = PesExtension {
            pes_private_data: None,
            pack_header: None,
            program_packet_sequence_counter: Some(ppsc),
            p_std_buffer: Some(PStdBuffer {
                scale: false,
                size: 100,
            }),
            pes_extension_field: None,
        };
        let h = PesHeader {
            scrambling_control: 0,
            pes_priority: true,
            data_alignment_indicator: false,
            copyright: false,
            original_or_copy: true,
            pts: Some(Pts(90_000)),
            dts: Some(Dts(85_000)),
            escr: Some(Escr {
                base: 1000,
                extension: 0,
            }),
            es_rate: Some(50_000),
            dsm_trick_mode: Some(TrickMode::SlowMotion { rep_cntrl: 3 }),
            additional_copy_info: Some(5),
            pes_crc: Some(0xCAFE),
            pes_extension: Some(ext),
            header_stuffing_len: 0,
        };
        let bytes = build_pes(h, &[0xFF]);
        let pkt = PesPacket::parse(&bytes).unwrap();
        let dh = pkt.header.unwrap();
        assert_eq!(dh.pts, Some(Pts(90_000)));
        assert_eq!(dh.dts, Some(Dts(85_000)));
        assert!(dh.escr.is_some());
        assert_eq!(dh.es_rate, Some(50_000));
        assert_eq!(
            dh.dsm_trick_mode,
            Some(TrickMode::SlowMotion { rep_cntrl: 3 })
        );
        assert_eq!(dh.additional_copy_info, Some(5));
        assert_eq!(dh.pes_crc, Some(0xCAFE));
        assert!(dh.pes_extension.is_some());
    }
}