arcly-stream 0.1.1

//! AV1 bitstream helpers: OBU sequence-header resolution extraction and
//! random-access detection.

use super::bitreader::BitReader;
use super::obu::iter_obus;
use super::parser::{CodecParser, VideoParams};
use crate::CodecId;

/// OBU type for a sequence header.
pub const OBU_SEQUENCE_HEADER: u8 = 1;

/// Parse an AV1 sequence-header OBU payload into [`VideoParams`].
///
/// Handles the common carriage cases (reduced still-picture headers and
/// timing-info-absent sequence headers). Returns `None` for sequence headers
/// carrying timing/decoder-model info, which require fuller parsing.
fn parse_sequence_header(payload: &[u8]) -> Option<VideoParams> {
    let mut r = BitReader::new(payload);
    let seq_profile = r.read_bits(3)? as u8;
    let _still_picture = r.read_bit()?;
    let reduced_still_picture_header = r.read_bit()?;

    let level;
    if reduced_still_picture_header == 1 {
        level = r.read_bits(5)? as u8; // seq_level_idx[0]
    } else {
        let timing_info_present_flag = r.read_bit()?;
        if timing_info_present_flag == 1 {
            return None; // timing_info()/decoder_model not parsed here
        }
        let initial_display_delay_present_flag = r.read_bit()?;
        let operating_points_cnt_minus1 = r.read_bits(5)?;
        let mut level0 = 0u8;
        for i in 0..=operating_points_cnt_minus1 {
            let _operating_point_idc = r.read_bits(12)?;
            let seq_level_idx = r.read_bits(5)?;
            if seq_level_idx > 7 {
                let _seq_tier = r.read_bit()?;
            }
            if initial_display_delay_present_flag == 1 {
                let present = r.read_bit()?;
                if present == 1 {
                    let _initial_display_delay_minus_1 = r.read_bits(4)?;
                }
            }
            if i == 0 {
                level0 = seq_level_idx as u8;
            }
        }
        level = level0;
    }

    let frame_width_bits_minus_1 = r.read_bits(4)?;
    let frame_height_bits_minus_1 = r.read_bits(4)?;
    let max_frame_width_minus_1 = r.read_bits(frame_width_bits_minus_1 + 1)?;
    let max_frame_height_minus_1 = r.read_bits(frame_height_bits_minus_1 + 1)?;

    Some(VideoParams {
        width: max_frame_width_minus_1 + 1,
        height: max_frame_height_minus_1 + 1,
        profile: seq_profile,
        level,
        tier: 0,
        bit_depth: 8,
    })
}

/// [`CodecParser`] implementation for AV1.
pub struct Av1;

impl CodecParser for Av1 {
    const CODEC: CodecId = CodecId::AV1;

    fn parse_config(data: &[u8]) -> Option<VideoParams> {
        iter_obus(data)
            .find(|o| o.obu_type == OBU_SEQUENCE_HEADER)
            .and_then(|o| parse_sequence_header(o.payload))
    }

    fn is_random_access_point(data: &[u8]) -> bool {
        // A sequence-header OBU is only sent at a random-access point; its
        // presence marks a clean decode start (and the keyframe that follows).
        iter_obus(data).any(|o| o.obu_type == OBU_SEQUENCE_HEADER)
    }

    fn carries_config(data: &[u8]) -> bool {
        Self::is_random_access_point(data)
    }

    fn hls_codec_string(p: &VideoParams) -> String {
        // av01.<profile>.<level><tier=M>.<bit_depth>
        format!("av01.{}.{:02}M.{:02}", p.profile, p.level, p.bit_depth)
    }
}

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

    /// Build a low-overhead sequence-header OBU for 1920x1080, profile 0.
    fn seq_header_obu_1920x1080() -> Vec<u8> {
        let mut w = BitWriter::default();
        w.bits(0, 3); // seq_profile
        w.bit(0); // still_picture
        w.bit(0); // reduced_still_picture_header
        w.bit(0); // timing_info_present_flag
        w.bit(0); // initial_display_delay_present_flag
        w.bits(0, 5); // operating_points_cnt_minus1 = 0
        w.bits(0, 12); // operating_point_idc[0]
        w.bits(1, 5); // seq_level_idx[0] = 1 (<= 7 → no seq_tier)
        w.bits(11, 4); // frame_width_bits_minus_1 = 11 → 12-bit width field
        w.bits(11, 4); // frame_height_bits_minus_1 = 11 → 12-bit height field
        w.bits(1919, 12); // max_frame_width_minus_1
        w.bits(1079, 12); // max_frame_height_minus_1
        w.align();
        let payload = w.bytes();

        let mut obu = vec![0x0A]; // header: type=1 (seq header), has_size_field=1
        obu.push(payload.len() as u8); // LEB128 size (< 128)
        obu.extend_from_slice(&payload);
        obu
    }

    #[test]
    fn parse_sequence_header_extracts_resolution() {
        let p = parse_sequence_header(&seq_header_obu_1920x1080()[2..]).expect("parse");
        assert_eq!((p.width, p.height), (1920, 1080));
        assert_eq!((p.profile, p.level), (0, 1));
    }

    #[test]
    fn classifies_and_extracts_via_obus() {
        // Temporal delimiter + sequence header + a frame OBU.
        let mut tu = vec![0x12, 0x00]; // OBU_TEMPORAL_DELIMITER (type 2), size 0
        tu.extend_from_slice(&seq_header_obu_1920x1080());
        tu.extend_from_slice(&[0x32, 0x02, 0xAA, 0xBB]); // OBU_FRAME (type 6), size 2

        assert!(Av1::is_random_access_point(&tu));
        assert!(Av1::carries_config(&tu));
        let p = Av1::parse_config(&tu).expect("params");
        assert_eq!((p.width, p.height), (1920, 1080));
        assert_eq!(Av1::hls_codec_string(&p), "av01.0.01M.08");

        // A temporal unit with no sequence header is not a RAP.
        let no_seq = [0x12, 0x00, 0x32, 0x02, 0xAA, 0xBB];
        assert!(!Av1::is_random_access_point(&no_seq));
    }
}