arcly-stream 0.8.3

//! H.264 (AVC) bitstream helpers: NAL iteration, Annex-B ↔ AVCC conversion,
//! SPS resolution extraction, and the `avcC` decoder-configuration record
//! ([`AvcConfig`]).

use super::bitreader::BitReader;
use super::nal::{iter_nals, unescape_rbsp};
use super::parser::{CodecParser, VideoParams};
use crate::CodecId;
use bytes::{BufMut, Bytes, BytesMut};

/// NAL unit type for a sequence parameter set.
pub const NAL_SPS: u8 = 7;
/// NAL unit type for a picture parameter set.
pub const NAL_PPS: u8 = 8;
/// NAL unit type for a coded slice of an IDR picture (keyframe).
pub const NAL_IDR: u8 = 5;

/// Iterate the NAL units of an Annex-B (`00 00 01` / `00 00 00 01`) bytestream,
/// yielding each NAL payload **without** its start code.
pub fn iter_nals_annexb(data: &[u8]) -> impl Iterator<Item = &[u8]> {
    iter_nals(data)
}

/// Convert an Annex-B bytestream to AVCC (length-prefixed) with a 4-byte length.
pub fn annexb_to_avcc(data: &[u8]) -> Bytes {
    let mut out = BytesMut::with_capacity(data.len());
    for nal in iter_nals(data) {
        out.put_u32(nal.len() as u32);
        out.put_slice(nal);
    }
    out.freeze()
}

/// Convert an AVCC (length-prefixed) bytestream to Annex-B (`00 00 00 01`).
///
/// `nal_length_size` is the prefix width in bytes (1–4), from the
/// AVCDecoderConfigurationRecord (`lengthSizeMinusOne + 1`).
pub fn avcc_to_annexb(data: &[u8], nal_length_size: usize) -> Option<Bytes> {
    if !(1..=4).contains(&nal_length_size) {
        return None;
    }
    let mut out = BytesMut::with_capacity(data.len() + 16);
    let mut i = 0;
    while i + nal_length_size <= data.len() {
        let mut len = 0usize;
        for _ in 0..nal_length_size {
            len = (len << 8) | data[i] as usize;
            i += 1;
        }
        if i + len > data.len() {
            return None; // truncated NAL
        }
        out.put_slice(&[0, 0, 0, 1]);
        out.put_slice(&data[i..i + len]);
        i += len;
    }
    Some(out.freeze())
}

/// Parsed AVC decoder configuration (the `avcC` record from a video sequence
/// header / `AVCDecoderConfigurationRecord`), retaining the parameter sets needed
/// to (a) build Annex-B config for a decoder and (b) self-contain each keyframe.
///
/// Lives in the codec layer (not a protocol module) so every consumer — the RTMP
/// FLV path, the transcode pipeline, and any host — shares one parser.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AvcConfig {
    /// NAL length prefix width in bytes (`lengthSizeMinusOne + 1`).
    pub nal_length_size: usize,
    /// Sequence parameter sets.
    pub sps: Vec<Vec<u8>>,
    /// Picture parameter sets.
    pub pps: Vec<Vec<u8>>,
}

impl AvcConfig {
    /// Parse an `AVCDecoderConfigurationRecord` (`avcC`). Returns `None` if the
    /// buffer is too short or its `configurationVersion` is not 1.
    pub fn parse(rec: &[u8]) -> Option<Self> {
        if rec.len() < 7 || rec[0] != 1 {
            return None;
        }
        let nal_length_size = (rec[4] & 0x03) as usize + 1;
        let mut pos = 5;
        let num_sps = (rec.get(pos)? & 0x1F) as usize;
        pos += 1;
        let mut sps = Vec::with_capacity(num_sps);
        for _ in 0..num_sps {
            let len = u16::from_be_bytes(rec.get(pos..pos + 2)?.try_into().ok()?) as usize;
            pos += 2;
            sps.push(rec.get(pos..pos + len)?.to_vec());
            pos += len;
        }
        let num_pps = *rec.get(pos)? as usize;
        pos += 1;
        let mut pps = Vec::with_capacity(num_pps);
        for _ in 0..num_pps {
            let len = u16::from_be_bytes(rec.get(pos..pos + 2)?.try_into().ok()?) as usize;
            pos += 2;
            pps.push(rec.get(pos..pos + len)?.to_vec());
            pos += len;
        }
        Some(Self {
            nal_length_size,
            sps,
            pps,
        })
    }

    /// The SPS/PPS as an Annex-B access unit (each NAL prefixed with `00 00 00 01`).
    pub fn to_annexb(&self) -> Bytes {
        let mut out = BytesMut::new();
        for nal in self.sps.iter().chain(self.pps.iter()) {
            out.put_slice(&[0, 0, 0, 1]);
            out.put_slice(nal);
        }
        out.freeze()
    }

    /// Rebuild an `AVCDecoderConfigurationRecord` from the stored parameter sets
    /// (used when serving a `play` request).
    pub fn to_avc_record(&self) -> Bytes {
        let sps0 = self
            .sps
            .first()
            .map(|s| s.as_slice())
            .unwrap_or(&[0, 0, 0, 0]);
        let mut out = BytesMut::new();
        out.put_u8(1); // configurationVersion
        out.put_u8(*sps0.get(1).unwrap_or(&0)); // AVCProfileIndication
        out.put_u8(*sps0.get(2).unwrap_or(&0)); // profile_compatibility
        out.put_u8(*sps0.get(3).unwrap_or(&0)); // AVCLevelIndication
        out.put_u8(0xFF); // 6 reserved bits + lengthSizeMinusOne = 3
        out.put_u8(0xE0 | (self.sps.len() as u8 & 0x1F));
        for s in &self.sps {
            out.put_u16(s.len() as u16);
            out.put_slice(s);
        }
        out.put_u8(self.pps.len() as u8);
        for p in &self.pps {
            out.put_u16(p.len() as u16);
            out.put_slice(p);
        }
        out.freeze()
    }

    /// Build an Annex-B access unit from length-prefixed (AVCC) NAL `data`,
    /// optionally prepending SPS/PPS so a keyframe is self-decodable.
    pub fn avcc_to_annexb(&self, avcc: &[u8], prepend_params: bool) -> Option<Bytes> {
        let body = avcc_to_annexb(avcc, self.nal_length_size)?;
        if prepend_params {
            let mut out = BytesMut::with_capacity(body.len() + 64);
            out.put(self.to_annexb());
            out.put(body);
            Some(out.freeze())
        } else {
            Some(body)
        }
    }
}

/// Decoded video dimensions and profile from a sequence parameter set.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SpsInfo {
    /// Luma width in pixels (after cropping).
    pub width: u32,
    /// Luma height in pixels (after cropping).
    pub height: u32,
    /// `profile_idc` (66 = baseline, 77 = main, 100 = high, …).
    pub profile_idc: u8,
    /// `level_idc` × 10 (e.g. 31 = level 3.1).
    pub level_idc: u8,
}

/// Parse a sequence parameter set NAL (including its 1-byte NAL header) and
/// extract the coded resolution. Returns `None` on a malformed SPS or one using
/// a scaling matrix (unsupported here).
pub fn parse_sps(nal: &[u8]) -> Option<SpsInfo> {
    if nal.is_empty() || nal[0] & 0x1f != NAL_SPS {
        return None;
    }
    let rbsp = unescape_rbsp(&nal[1..]);
    let mut r = BitReader::new(&rbsp);
    if r.remaining() < 24 {
        return None; // too short to hold profile/constraints/level
    }

    let profile_idc = r.read_bits(8)? as u8;
    let _constraints = r.read_bits(8)?;
    let level_idc = r.read_bits(8)? as u8;
    let _sps_id = r.read_ue()?;

    let mut chroma_format_idc = 1u32; // 4:2:0 default
    if matches!(
        profile_idc,
        100 | 110 | 122 | 244 | 44 | 83 | 86 | 118 | 128 | 138 | 139 | 134 | 135
    ) {
        chroma_format_idc = r.read_ue()?;
        if chroma_format_idc == 3 {
            let _separate_colour_plane = r.read_bit()?;
        }
        let _bit_depth_luma = r.read_ue()?;
        let _bit_depth_chroma = r.read_ue()?;
        let _qpprime = r.read_bit()?;
        let scaling_matrix_present = r.read_bit()?;
        if scaling_matrix_present == 1 {
            return None; // scaling lists not decoded
        }
    }

    let _log2_max_frame_num = r.read_ue()?;
    let pic_order_cnt_type = r.read_ue()?;
    if pic_order_cnt_type == 0 {
        let _log2_max_poc_lsb = r.read_ue()?;
    } else if pic_order_cnt_type == 1 {
        let _delta_pic_order_always_zero = r.read_bit()?;
        let _offset_for_non_ref = r.read_se()?;
        let _offset_for_top_to_bottom = r.read_se()?;
        let cycle = r.read_ue()?;
        for _ in 0..cycle {
            let _ = r.read_se()?;
        }
    }

    let _max_num_ref_frames = r.read_ue()?;
    let _gaps_allowed = r.read_bit()?;
    let pic_width_in_mbs_minus1 = r.read_ue()?;
    let pic_height_in_map_units_minus1 = r.read_ue()?;
    let frame_mbs_only_flag = r.read_bit()?;
    if frame_mbs_only_flag == 0 {
        let _mb_adaptive = r.read_bit()?;
    }
    let _direct_8x8 = r.read_bit()?;

    let frame_cropping_flag = r.read_bit()?;
    let (mut crop_left, mut crop_right, mut crop_top, mut crop_bottom) = (0u32, 0u32, 0u32, 0u32);
    if frame_cropping_flag == 1 {
        crop_left = r.read_ue()?;
        crop_right = r.read_ue()?;
        crop_top = r.read_ue()?;
        crop_bottom = r.read_ue()?;
    }

    let width_mbs = pic_width_in_mbs_minus1 + 1;
    let height_map = pic_height_in_map_units_minus1 + 1;
    let frame_height_mbs = (2 - frame_mbs_only_flag) * height_map;

    // Chroma sub-sampling factors for the crop unit.
    let (sub_w, sub_h) = match chroma_format_idc {
        0 => (1, 1), // monochrome
        1 => (2, 2), // 4:2:0
        2 => (2, 1), // 4:2:2
        _ => (1, 1), // 4:4:4
    };
    let crop_unit_x = if chroma_format_idc == 0 { 1 } else { sub_w };
    let crop_unit_y = (if chroma_format_idc == 0 { 1 } else { sub_h }) * (2 - frame_mbs_only_flag);

    let width = width_mbs * 16 - (crop_left + crop_right) * crop_unit_x;
    let height = frame_height_mbs * 16 - (crop_top + crop_bottom) * crop_unit_y;

    Some(SpsInfo {
        width,
        height,
        profile_idc,
        level_idc,
    })
}

/// The HLS `CODECS` attribute for H.264, e.g. `"avc1.4d401f"` (profile/level
/// encoded as `avc1.PPCCLL`). The middle byte is the constraint-set flags (0).
pub fn hls_codec_string(p: &VideoParams) -> String {
    format!("avc1.{:02x}00{:02x}", p.profile, p.level)
}

/// [`CodecParser`] implementation for H.264 / AVC.
pub struct H264;

impl CodecParser for H264 {
    const CODEC: CodecId = CodecId::H264;

    fn parse_config(data: &[u8]) -> Option<VideoParams> {
        for nal in iter_nals(data) {
            if !nal.is_empty() && nal[0] & 0x1f == NAL_SPS {
                let sps = parse_sps(nal)?;
                return Some(VideoParams {
                    width: sps.width,
                    height: sps.height,
                    profile: sps.profile_idc,
                    level: sps.level_idc,
                    tier: 0,
                    bit_depth: 8,
                });
            }
        }
        None
    }

    fn is_random_access_point(data: &[u8]) -> bool {
        iter_nals(data).any(|n| !n.is_empty() && n[0] & 0x1f == NAL_IDR)
    }

    fn carries_config(data: &[u8]) -> bool {
        iter_nals(data).any(|n| !n.is_empty() && matches!(n[0] & 0x1f, NAL_SPS | NAL_PPS))
    }

    fn hls_codec_string(params: &VideoParams) -> String {
        hls_codec_string(params)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::codec::testutil::BitWriter;

    #[test]
    fn iterates_and_converts_nals() {
        // Two NALs: [9 0xF0] and [7 0x42], with mixed 4- and 3-byte start codes.
        let annexb = [0, 0, 0, 1, 9, 0xF0, 0, 0, 1, 7, 0x42];
        let nals: Vec<&[u8]> = iter_nals_annexb(&annexb).collect();
        assert_eq!(nals, vec![&[9u8, 0xF0][..], &[7u8, 0x42][..]]);

        // Round-trip Annex-B → AVCC → Annex-B (normalizes to 4-byte codes).
        let avcc = annexb_to_avcc(&annexb);
        assert_eq!(&avcc[..], &[0, 0, 0, 2, 9, 0xF0, 0, 0, 0, 2, 7, 0x42]);
        let back = avcc_to_annexb(&avcc, 4).unwrap();
        assert_eq!(&back[..], &[0, 0, 0, 1, 9, 0xF0, 0, 0, 0, 1, 7, 0x42]);
    }

    #[test]
    fn avc_config_parses_and_round_trips() {
        // avcC: version 1, profile/compat/level bytes, lengthSizeMinusOne=3,
        // one SPS [0x67,0x42,0x00,0x1F], one PPS [0x68,0xCE].
        let avcc = [
            0x01, 0x42, 0x00, 0x1F, 0xFF, 0xE1, 0x00, 0x04, 0x67, 0x42, 0x00, 0x1F, 0x01, 0x00,
            0x02, 0x68, 0xCE,
        ];
        let cfg = AvcConfig::parse(&avcc).expect("parse avcC");
        assert_eq!(cfg.nal_length_size, 4);
        assert_eq!(cfg.sps, vec![vec![0x67, 0x42, 0x00, 0x1F]]);
        assert_eq!(cfg.pps, vec![vec![0x68, 0xCE]]);

        // to_annexb emits each param set with a 4-byte start code.
        assert_eq!(
            &cfg.to_annexb()[..],
            &[0, 0, 0, 1, 0x67, 0x42, 0x00, 0x1F, 0, 0, 0, 1, 0x68, 0xCE]
        );

        // A length-prefixed (AVCC) IDR converts to Annex-B, prepending params.
        let idr_avcc = [0, 0, 0, 2, 0x65, 0xAA];
        let plain = cfg.avcc_to_annexb(&idr_avcc, false).unwrap();
        assert_eq!(&plain[..], &[0, 0, 0, 1, 0x65, 0xAA]);
        let with_params = cfg.avcc_to_annexb(&idr_avcc, true).unwrap();
        assert!(with_params.ends_with(&[0, 0, 0, 1, 0x65, 0xAA]));
        assert!(with_params.len() > plain.len());

        // Rebuilding the record preserves the parameter sets.
        assert_eq!(AvcConfig::parse(&cfg.to_avc_record()), Some(cfg));

        // Not an avcC record (Annex-B start code) → None.
        assert_eq!(AvcConfig::parse(&[0, 0, 0, 1, 0x67]), None);
    }

    /// Synthesize a baseline (profile 66) SPS for 1280x720, no cropping.
    fn sps_1280x720() -> Vec<u8> {
        let mut w = BitWriter::default();
        w.bits(66, 8); // profile_idc (baseline → no chroma block)
        w.bits(0, 8); // constraint flags
        w.bits(31, 8); // level 3.1
        w.ue(0); // seq_parameter_set_id
        w.ue(0); // log2_max_frame_num_minus4
        w.ue(0); // pic_order_cnt_type = 0
        w.ue(0); // log2_max_pic_order_cnt_lsb_minus4
        w.ue(1); // max_num_ref_frames
        w.bit(0); // gaps_in_frame_num_value_allowed
        w.ue(79); // pic_width_in_mbs_minus1 → 80*16 = 1280
        w.ue(44); // pic_height_in_map_units_minus1 → 45*16 = 720
        w.bit(1); // frame_mbs_only_flag
        w.bit(0); // direct_8x8_inference_flag
        w.bit(0); // frame_cropping_flag
        w.bit(0); // vui_parameters_present_flag
        let mut nal = vec![0x67u8]; // NAL header: SPS (type 7), ref_idc 3
        nal.extend_from_slice(&w.bytes());
        nal
    }

    #[test]
    fn parse_sps_extracts_resolution() {
        let sps = parse_sps(&sps_1280x720()).expect("parse");
        assert_eq!((sps.width, sps.height), (1280, 720));
        assert_eq!(sps.profile_idc, 66);
        assert_eq!(sps.level_idc, 31);
    }

    #[test]
    fn codec_parser_classifies_and_extracts() {
        // Annex-B AU: SPS + an IDR slice (type 5).
        let mut au = vec![0, 0, 0, 1];
        au.extend_from_slice(&sps_1280x720());
        au.extend_from_slice(&[0, 0, 0, 1, 0x65, 0xAA]); // IDR NAL (type 5)

        assert!(H264::is_random_access_point(&au));
        assert!(H264::carries_config(&au));
        let p = H264::parse_config(&au).expect("params");
        assert_eq!((p.width, p.height, p.profile, p.level), (1280, 720, 66, 31));
        assert_eq!(H264::hls_codec_string(&p), "avc1.42001f");

        // A lone non-IDR slice (type 1) is not a RAP.
        let delta = [0, 0, 0, 1, 0x41, 0xAA];
        assert!(!H264::is_random_access_point(&delta));
    }
}