arcly-stream 0.1.6

//! H.264 (AVC) bitstream helpers: NAL iteration, Annex-B ↔ AVCC conversion,
//! and SPS resolution extraction.

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())
}

/// 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]);
    }

    /// 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));
    }
}