//! Reading an AVC or HEVC elementary stream: the Annex B byte stream (ISO/IEC
//! 14496-10 Annex B, ISO/IEC 23008-2 Annex B), its NAL units, access units, sequence
//! parameter set and the start of each slice header. This reads; it codes nothing.
//!
//! [`check_stream`] holds a stream to [platform-encode] §4.3: one access unit per
//! picture, an IDR access unit at exactly the planned pictures and no other random
//! access point, the parameter sets in every IDR access unit and the same bytes each
//! time, progressive 8-bit 4:2:0, a single layer, and no reordering.
//!
//! [platform-encode]: ../../../../extensions/platform-encode.md

use super::Codec;
use std::ops::Range;

/// Why a stream is refused, with the reason code of [platform-encode] §9.
///
/// [platform-encode]: ../../../../extensions/platform-encode.md
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Reason {
    pub code: &'static str,
    pub detail: String,
}

fn refuse<T>(code: &'static str, detail: impl Into<String>) -> Result<T, Reason> {
    Err(Reason { code, detail: detail.into() })
}

const MALFORMED: &str = "encode-stream-malformed";

/// One NAL unit: its byte range in the stream (the NAL unit itself, without the start
/// code prefix or trailing zero bytes).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Nal {
    pub range: Range<usize>,
    /// `nal_unit_type`.
    pub kind: u8,
    /// AVC: `nal_ref_idc`. HEVC: `nuh_layer_id`.
    pub aux: u8,
}

/// Split an Annex B byte stream into its NAL units (Annex B.2). The stream begins with
/// zero bytes and a start code prefix `00 00 01`; a NAL unit ends where the next three
/// bytes are `00 00 00` or `00 00 01`; zero bytes after it are trailing. Refused: a
/// stream that does not begin that way, a NAL unit of no bytes, or one holding `00 00 02`.
pub fn nal_units(codec: Codec, s: &[u8]) -> Result<Vec<Nal>, Reason> {
    let mut i = 0usize;
    while i < s.len() && s[i] == 0 {
        i += 1;
    }
    if i < 2 || i >= s.len() || s[i] != 1 {
        return refuse(MALFORMED, "the stream does not begin with a start code prefix");
    }
    let header = match codec {
        Codec::Avc => 1,
        Codec::Hevc => 2,
    };
    let mut out = Vec::new();
    let mut at = i + 1;
    loop {
        // The NAL unit runs to the next 00 00 00 or 00 00 01, or the end of the stream.
        let mut end = at;
        let mut next = None;
        while end + 2 < s.len() {
            if s[end] == 0 && s[end + 1] == 0 {
                match s[end + 2] {
                    0 | 1 => break,
                    2 => return refuse(MALFORMED, format!("00 00 02 inside the NAL unit at byte {at}")),
                    _ => {}
                }
            }
            end += 1;
        }
        if end + 2 >= s.len() {
            end = s.len();
        }
        // Skip trailing zeros to the next start code, if any.
        let mut j = end;
        while j < s.len() && s[j] == 0 {
            j += 1;
        }
        if j < s.len() {
            if s[j] != 1 || j - end < 2 {
                return refuse(MALFORMED, format!("bytes after the NAL unit at byte {at} are not a start code prefix"));
            }
            next = Some(j + 1);
        }
        if end - at < header {
            return refuse(MALFORMED, format!("a NAL unit of {} bytes at byte {at}", end - at));
        }
        if s[at] & 0x80 != 0 {
            return refuse(MALFORMED, format!("forbidden_zero_bit set in the NAL unit at byte {at}"));
        }
        let (kind, aux) = match codec {
            Codec::Avc => (s[at] & 0x1f, (s[at] >> 5) & 3),
            Codec::Hevc => ((s[at] >> 1) & 0x3f, ((s[at] & 1) << 5) | (s[at + 1] >> 3)),
        };
        if codec == Codec::Hevc && s[at + 1] & 7 == 0 {
            return refuse(MALFORMED, format!("nuh_temporal_id_plus1 of 0 in the NAL unit at byte {at}"));
        }
        out.push(Nal { range: at..end, kind, aux });
        match next {
            Some(n) if n < s.len() => at = n,
            Some(_) => return refuse(MALFORMED, "a start code prefix at the end of the stream"),
            None => break,
        }
    }
    Ok(out)
}

/// The RBSP of a NAL unit's payload: `00 00 03` read as `00 00` (§7.4.1, emulation
/// prevention), after the NAL unit header.
pub fn rbsp(codec: Codec, nal: &[u8]) -> Vec<u8> {
    let skip = match codec {
        Codec::Avc => 1,
        Codec::Hevc => 2,
    };
    let mut out = Vec::with_capacity(nal.len());
    let mut zeros = 0;
    for &b in &nal[skip.min(nal.len())..] {
        if zeros >= 2 && b == 3 {
            zeros = 0;
            continue;
        }
        zeros = if b == 0 { zeros + 1 } else { 0 };
        out.push(b);
    }
    out
}

/// A big-endian bit reader over an RBSP, with Exp-Golomb codes (§9.1). Reading past the
/// end is an error.
pub struct Bits<'a> {
    b: &'a [u8],
    pos: usize,
}

impl<'a> Bits<'a> {
    pub fn new(b: &'a [u8]) -> Bits<'a> {
        Bits { b, pos: 0 }
    }

    pub fn u(&mut self, n: u32) -> Result<u64, Reason> {
        let mut v = 0u64;
        for _ in 0..n {
            let byte = *self.b.get(self.pos / 8).ok_or_else(|| Reason { code: MALFORMED, detail: "a syntax element runs past its NAL unit".into() })?;
            v = (v << 1) | u64::from((byte >> (7 - (self.pos % 8))) & 1);
            self.pos += 1;
        }
        Ok(v)
    }

    pub fn flag(&mut self) -> Result<bool, Reason> {
        Ok(self.u(1)? == 1)
    }

    /// `ue(v)`, at most 32 leading zero bits.
    pub fn ue(&mut self) -> Result<u64, Reason> {
        let mut zeros = 0;
        while !self.flag()? {
            zeros += 1;
            if zeros > 32 {
                return refuse(MALFORMED, "an Exp-Golomb code of more than 32 leading zeros");
            }
        }
        Ok((1u64 << zeros) - 1 + self.u(zeros)?)
    }

    /// `se(v)`.
    pub fn se(&mut self) -> Result<i64, Reason> {
        let k = self.ue()?;
        Ok(if k % 2 == 1 { k.div_ceil(2) as i64 } else { -((k / 2) as i64) })
    }

    pub fn skip(&mut self, n: u32) -> Result<(), Reason> {
        self.u(n).map(|_| ())
    }
}

/// What a sequence parameter set says that a record states and a check uses.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SpsInfo {
    /// `profile_idc` (AVC) or `general_profile_idc` (HEVC).
    pub profile_idc: u8,
    /// `level_idc` (AVC) or `general_level_idc` (HEVC).
    pub level_idc: u8,
    /// HEVC `general_tier_flag`; false for AVC.
    pub high_tier: bool,
    pub chroma_format_idc: u8,
    pub bit_depth_luma: u8,
    pub bit_depth_chroma: u8,
    /// The cropped picture's width and height in luma samples.
    pub width: u32,
    pub height: u32,
    /// Progressive frames only: AVC `frame_mbs_only_flag`; true for HEVC.
    pub frame_mbs_only: bool,
    /// AVC `pic_order_cnt_type`; 0 for HEVC.
    pub poc_type: u8,
    /// `log2_max_frame_num_minus4 + 4` (AVC); 0 for HEVC.
    pub log2_max_frame_num: u8,
    /// `log2_max_pic_order_cnt_lsb_minus4 + 4`.
    pub log2_max_poc_lsb: u8,
    /// HEVC `sps_max_num_reorder_pics` of the highest sub-layer; AVC: the VUI's
    /// `max_num_reorder_frames` where the SPS has one.
    pub max_num_reorder: Option<u64>,
    /// AVC: the VUI's `video_full_range_flag` where the SPS signals one.
    pub full_range: Option<bool>,
    pub separate_colour_plane: bool,
}

fn scaling_list(r: &mut Bits<'_>, size: usize) -> Result<(), Reason> {
    let (mut last, mut next) = (8i64, 8i64);
    for _ in 0..size {
        if next != 0 {
            let delta = r.se()?;
            next = (last + delta + 256).rem_euclid(256);
        }
        last = if next == 0 { last } else { next };
    }
    Ok(())
}

fn avc_hrd(r: &mut Bits<'_>) -> Result<(), Reason> {
    let cnt = r.ue()?;
    if cnt > 31 {
        return refuse(MALFORMED, "cpb_cnt_minus1 above 31");
    }
    r.skip(8)?;
    for _ in 0..=cnt {
        r.ue()?;
        r.ue()?;
        r.skip(1)?;
    }
    r.skip(20)
}

/// Read an AVC sequence parameter set's RBSP (ISO/IEC 14496-10 §7.3.2.1.1, Annex E.1.1).
pub fn avc_sps(rbsp: &[u8]) -> Result<SpsInfo, Reason> {
    let mut r = Bits::new(rbsp);
    let profile_idc = r.u(8)? as u8;
    r.skip(8)?;
    let level_idc = r.u(8)? as u8;
    r.ue()?;
    let (mut chroma, mut sep, mut bdl, mut bdc) = (1u64, false, 8u64, 8u64);
    if matches!(profile_idc, 100 | 110 | 122 | 244 | 44 | 83 | 86 | 118 | 128 | 138 | 139 | 134 | 135) {
        chroma = r.ue()?;
        if chroma > 3 {
            return refuse(MALFORMED, "chroma_format_idc above 3");
        }
        if chroma == 3 {
            sep = r.flag()?;
        }
        bdl = r.ue()? + 8;
        bdc = r.ue()? + 8;
        r.skip(1)?;
        if r.flag()? {
            for i in 0..if chroma != 3 { 8 } else { 12 } {
                if r.flag()? {
                    scaling_list(&mut r, if i < 6 { 16 } else { 64 })?;
                }
            }
        }
    }
    let log2_max_frame_num = r.ue()? + 4;
    let poc_type = r.ue()?;
    let mut log2_max_poc_lsb = 0;
    match poc_type {
        0 => log2_max_poc_lsb = r.ue()? + 4,
        1 => {
            r.skip(1)?;
            r.se()?;
            r.se()?;
            let n = r.ue()?;
            if n > 255 {
                return refuse(MALFORMED, "num_ref_frames_in_pic_order_cnt_cycle above 255");
            }
            for _ in 0..n {
                r.se()?;
            }
        }
        2 => {}
        _ => return refuse(MALFORMED, "pic_order_cnt_type above 2"),
    }
    if log2_max_frame_num > 16 || log2_max_poc_lsb > 16 {
        return refuse(MALFORMED, "log2_max_frame_num or log2_max_pic_order_cnt_lsb above 16");
    }
    r.ue()?;
    r.skip(1)?;
    let wmbs = r.ue()? + 1;
    let hmap = r.ue()? + 1;
    let frame_mbs_only = r.flag()?;
    if !frame_mbs_only {
        r.skip(1)?;
    }
    r.skip(1)?;
    let (mut cl, mut cr, mut ct, mut cb) = (0u64, 0u64, 0u64, 0u64);
    if r.flag()? {
        cl = r.ue()?;
        cr = r.ue()?;
        ct = r.ue()?;
        cb = r.ue()?;
    }
    let (sub_w, sub_h) = match (chroma, sep) {
        (1, _) => (2, 2),
        (2, _) => (2, 1),
        _ => (1, 1),
    };
    let (crop_x, crop_y) = if chroma == 0 || sep { (1, 2 - u64::from(frame_mbs_only)) } else { (sub_w, sub_h * (2 - u64::from(frame_mbs_only))) };
    let full_w = wmbs * 16;
    let full_h = (2 - u64::from(frame_mbs_only)) * hmap * 16;
    let (cw, ch) = (crop_x * (cl + cr), crop_y * (ct + cb));
    if cw >= full_w || ch >= full_h || full_w > 1 << 16 || full_h > 1 << 16 {
        return refuse(MALFORMED, "a cropping window larger than the picture");
    }
    let (mut full_range, mut max_num_reorder) = (None, None);
    if r.flag()? {
        // vui_parameters (E.1.1)
        if r.flag()? && r.u(8)? == 255 {
            r.skip(32)?;
        }
        if r.flag()? {
            r.skip(1)?;
        }
        if r.flag()? {
            r.skip(3)?;
            full_range = Some(r.flag()?);
            if r.flag()? {
                r.skip(24)?;
            }
        }
        if r.flag()? {
            r.ue()?;
            r.ue()?;
        }
        if r.flag()? {
            r.skip(65)?;
        }
        let nal_hrd = r.flag()?;
        if nal_hrd {
            avc_hrd(&mut r)?;
        }
        let vcl_hrd = r.flag()?;
        if vcl_hrd {
            avc_hrd(&mut r)?;
        }
        if nal_hrd || vcl_hrd {
            r.skip(1)?;
        }
        r.skip(1)?;
        if r.flag()? {
            r.skip(1)?;
            for _ in 0..4 {
                r.ue()?;
            }
            max_num_reorder = Some(r.ue()?);
            r.ue()?;
        }
    }
    Ok(SpsInfo {
        profile_idc,
        level_idc,
        high_tier: false,
        chroma_format_idc: chroma as u8,
        bit_depth_luma: bdl.min(255) as u8,
        bit_depth_chroma: bdc.min(255) as u8,
        width: (full_w - cw) as u32,
        height: (full_h - ch) as u32,
        frame_mbs_only,
        poc_type: poc_type as u8,
        log2_max_frame_num: log2_max_frame_num as u8,
        log2_max_poc_lsb: log2_max_poc_lsb as u8,
        max_num_reorder,
        full_range,
        separate_colour_plane: sep,
    })
}

/// `profile_tier_level(1, max_sub_layers_minus1)` (ISO/IEC 23008-2 §7.3.3): the general
/// profile, tier and level.
fn hevc_ptl(r: &mut Bits<'_>, max_sub_layers_minus1: u64) -> Result<(u8, bool, u8), Reason> {
    r.skip(2)?;
    let tier = r.flag()?;
    let profile = r.u(5)? as u8;
    r.skip(32 + 4 + 43 + 1)?;
    let level = r.u(8)? as u8;
    let mut present = Vec::new();
    for _ in 0..max_sub_layers_minus1 {
        present.push((r.flag()?, r.flag()?));
    }
    if max_sub_layers_minus1 > 0 {
        for _ in max_sub_layers_minus1..8 {
            r.skip(2)?;
        }
    }
    for (p, l) in present {
        if p {
            r.skip(88)?;
        }
        if l {
            r.skip(8)?;
        }
    }
    Ok((profile, tier, level))
}

/// Read an HEVC sequence parameter set's RBSP (ISO/IEC 23008-2 §7.3.2.2.1), up to its
/// sub-layer ordering information.
pub fn hevc_sps(rbsp: &[u8]) -> Result<SpsInfo, Reason> {
    let mut r = Bits::new(rbsp);
    r.skip(4)?;
    let max_sub = r.u(3)?;
    if max_sub > 6 {
        return refuse(MALFORMED, "sps_max_sub_layers_minus1 above 6");
    }
    r.skip(1)?;
    let (profile_idc, high_tier, level_idc) = hevc_ptl(&mut r, max_sub)?;
    r.ue()?;
    let chroma = r.ue()?;
    if chroma > 3 {
        return refuse(MALFORMED, "chroma_format_idc above 3");
    }
    let sep = chroma == 3 && r.flag()?;
    let pw = r.ue()?;
    let ph = r.ue()?;
    let (mut cl, mut cr, mut ct, mut cb) = (0u64, 0u64, 0u64, 0u64);
    if r.flag()? {
        cl = r.ue()?;
        cr = r.ue()?;
        ct = r.ue()?;
        cb = r.ue()?;
    }
    let (sub_w, sub_h) = match (chroma, sep) {
        (1, _) => (2, 2),
        (2, false) => (2, 1),
        _ => (1, 1),
    };
    let (cw, ch) = (sub_w * (cl + cr), sub_h * (ct + cb));
    if pw == 0 || ph == 0 || cw >= pw || ch >= ph || pw > 1 << 16 || ph > 1 << 16 {
        return refuse(MALFORMED, "a picture size of 0, or a conformance window larger than the picture");
    }
    let bdl = r.ue()? + 8;
    let bdc = r.ue()? + 8;
    let log2_max_poc_lsb = r.ue()? + 4;
    if log2_max_poc_lsb > 16 {
        return refuse(MALFORMED, "log2_max_pic_order_cnt_lsb above 16");
    }
    let ordering_all = r.flag()?;
    let mut reorder = 0;
    for _ in (if ordering_all { 0 } else { max_sub })..=max_sub {
        r.ue()?;
        reorder = r.ue()?;
        r.ue()?;
    }
    Ok(SpsInfo {
        profile_idc,
        level_idc,
        high_tier,
        chroma_format_idc: chroma as u8,
        bit_depth_luma: bdl.min(255) as u8,
        bit_depth_chroma: bdc.min(255) as u8,
        width: (pw - cw) as u32,
        height: (ph - ch) as u32,
        frame_mbs_only: true,
        poc_type: 0,
        log2_max_frame_num: 0,
        log2_max_poc_lsb: log2_max_poc_lsb as u8,
        max_num_reorder: Some(reorder),
        full_range: None,
        separate_colour_plane: sep,
    })
}

fn is_vcl(codec: Codec, kind: u8) -> bool {
    match codec {
        Codec::Avc => matches!(kind, 1..=5),
        Codec::Hevc => kind < 32,
    }
}

/// Whether a NAL unit of this type, met after the last VCL NAL unit of an access unit,
/// begins the next one (ISO/IEC 14496-10 §7.4.1.2.3; ISO/IEC 23008-2 §7.4.2.4.4).
fn starts_next(codec: Codec, kind: u8) -> bool {
    match codec {
        Codec::Avc => matches!(kind, 6..=9 | 14..=18),
        Codec::Hevc => matches!(kind, 32..=35 | 39 | 41..=44 | 48..=55),
    }
}

/// Whether a VCL NAL unit is the first slice of its picture: AVC `first_mb_in_slice` 0
/// (Main and High have no arbitrary slice order); HEVC `first_slice_segment_in_pic_flag`.
fn first_slice(codec: Codec, s: &[u8], nal: &Nal) -> Result<bool, Reason> {
    let rb = rbsp(codec, &s[nal.range.clone()]);
    let mut r = Bits::new(&rb);
    Ok(match codec {
        Codec::Avc => r.ue()? == 0,
        Codec::Hevc => r.flag()?,
    })
}

/// The access units of a stream, as ranges of NAL unit indices.
pub fn access_units(codec: Codec, s: &[u8], nals: &[Nal]) -> Result<Vec<Range<usize>>, Reason> {
    let mut units = Vec::new();
    let mut start = 0usize;
    let mut seen_vcl = false;
    for (i, n) in nals.iter().enumerate() {
        let vcl = is_vcl(codec, n.kind);
        let begins = if vcl { seen_vcl && first_slice(codec, s, n)? } else { seen_vcl && starts_next(codec, n.kind) };
        if begins {
            units.push(start..i);
            start = i;
            seen_vcl = false;
        }
        if vcl {
            seen_vcl = true;
        }
    }
    if !seen_vcl {
        return refuse(MALFORMED, if nals.is_empty() { "no NAL unit".to_string() } else { "NAL units after the last picture's slices begin no picture".to_string() });
    }
    units.push(start..nals.len());
    Ok(units)
}

/// A stream held to §4.3.
#[derive(Debug, Clone)]
pub struct Checked {
    /// Each access unit's byte range in the stream, from its first NAL unit's start code
    /// prefix to the byte before the next access unit's.
    pub units: Vec<Range<usize>>,
    pub sps: SpsInfo,
    /// Which access units are IDR.
    pub idr: Vec<bool>,
    /// Each access unit's NAL units, as index ranges into [`nal_units`]'s list.
    pub nal_units: Vec<Range<usize>>,
}

fn kinds(codec: Codec) -> (&'static [u8], u8) {
    match codec {
        // SPS, PPS; the IDR slice type.
        Codec::Avc => (&[7, 8], 5),
        // VPS, SPS, PPS.
        Codec::Hevc => (&[32, 33, 34], 19),
    }
}

/// Hold `s` to [platform-encode] §4.3 for `frames` pictures with IDR access units at
/// exactly `idr_at`.
///
/// [platform-encode]: ../../../../extensions/platform-encode.md
pub fn check_stream(codec: Codec, s: &[u8], frames: u64, idr_at: &[u64]) -> Result<Checked, Reason> {
    let nals = nal_units(codec, s)?;
    let aus = access_units(codec, s, &nals)?;
    if aus.len() as u64 != frames {
        return refuse("encode-frame-count", format!("{} access units for {frames} pictures", aus.len()));
    }
    let (ps_kinds, _) = kinds(codec);
    let sps_kind = ps_kinds[ps_kinds.len() - 2];
    let mut first_ps: Option<Vec<&[u8]>> = None;
    let mut sps: Option<SpsInfo> = None;
    let mut idr = Vec::with_capacity(aus.len());
    // AVC picture order: the previous reference picture's PicOrderCntMsb and lsb (§8.2.1.1).
    let (mut prev_msb, mut prev_lsb, mut last_poc) = (0i64, 0i64, i64::MIN);
    for (a, au) in aus.iter().enumerate() {
        let au_nals = &nals[au.clone()];
        // One pass over the access unit's NAL units: what it holds, and whether each
        // parameter set and slice header reads (§4.3 item 1).
        let mut is_idr = None;
        let mut mixed = false;
        let mut other_irap = None;
        let mut ps: Vec<&[u8]> = Vec::new();
        let mut new_sps: Option<SpsInfo> = None;
        let mut first_vcl: Option<&Nal> = None;
        let mut layer = None;
        for n in au_nals {
            if codec == Codec::Hevc && n.aux != 0 {
                layer.get_or_insert(n.aux);
            }
            let bytes = &s[n.range.clone()];
            if ps_kinds.contains(&n.kind) {
                ps.push(bytes);
                if n.kind == sps_kind {
                    let info = match codec {
                        Codec::Avc => avc_sps(&rbsp(codec, bytes))?,
                        Codec::Hevc => hevc_sps(&rbsp(codec, bytes))?,
                    };
                    if new_sps.as_ref().is_some_and(|p| *p != info) {
                        return refuse("encode-parameter-sets", format!("access unit {a} carries two sequence parameter sets"));
                    }
                    new_sps = Some(info);
                }
            }
            if is_vcl(codec, n.kind) {
                let this_idr = match codec {
                    Codec::Avc => n.kind == 5,
                    Codec::Hevc => matches!(n.kind, 19 | 20),
                };
                if codec == Codec::Hevc && (16..=23).contains(&n.kind) && !this_idr {
                    other_irap.get_or_insert(n.kind);
                }
                mixed |= is_idr.is_some_and(|k| k != this_idr);
                is_idr = Some(this_idr);
                first_vcl.get_or_insert(n);
            }
        }
        let is_idr = is_idr.expect("every access unit has a slice");
        let fv = first_vcl.expect("every access unit has a slice");
        // §4.3 item 3: IDR access units at exactly the plan's pictures.
        if mixed {
            return refuse(MALFORMED, format!("access unit {a} mixes IDR and non-IDR slices"));
        }
        if let Some(k) = other_irap {
            return refuse("encode-keyframe-mismatch", format!("access unit {a} is a random access point that is not an IDR (nal_unit_type {k})"));
        }
        let planned = idr_at.binary_search(&(a as u64)).is_ok();
        if planned != is_idr {
            return refuse("encode-keyframe-mismatch", if planned { format!("picture {a} is planned as an IDR and is not one") } else { format!("picture {a} is an IDR the plan does not name") });
        }
        // §4.3 item 4: the parameter sets.
        if is_idr {
            if !ps_kinds.iter().all(|k| au_nals.iter().any(|n| n.kind == *k)) {
                return refuse("encode-parameter-sets", format!("IDR access unit {a} does not carry every parameter set"));
            }
            match &first_ps {
                None => first_ps = Some(ps.clone()),
                Some(f) if *f != ps => return refuse("encode-parameter-sets", format!("IDR access unit {a} carries other parameter set bytes than the first")),
                Some(_) => {}
            }
        } else if !ps.is_empty() && first_ps.as_ref().is_none_or(|f| ps.iter().any(|p| !f.contains(p))) {
            return refuse("encode-parameter-sets", format!("access unit {a} carries other parameter set bytes than the first IDR's"));
        }
        if let Some(n) = new_sps {
            if sps.as_ref().is_some_and(|p| *p != n) {
                return refuse("encode-parameter-sets", format!("access unit {a} carries another sequence parameter set"));
            }
            sps = Some(n);
        }
        let Some(info) = sps.as_ref() else {
            return refuse("encode-parameter-sets", format!("access unit {a} comes before any sequence parameter set"));
        };
        // §4.3 item 5: what this revision reads. AVC's scalable, multiview and 3-D NAL
        // unit types (14, 15, 20, 21) are other layers.
        if codec == Codec::Avc
            && let Some(k) = au_nals.iter().map(|n| n.kind).find(|k| matches!(k, 14 | 15 | 20 | 21))
        {
            return refuse("encode-stream-unsupported", format!("access unit {a} holds a NAL unit of type {k}, of another layer"));
        }
        if let Some(l) = layer {
            return refuse("encode-stream-unsupported", format!("access unit {a} holds a NAL unit of layer {l}"));
        }
        if info.chroma_format_idc != 1 || info.bit_depth_luma != 8 || info.bit_depth_chroma != 8 || info.separate_colour_plane {
            return refuse("encode-stream-unsupported", "the stream is not 8-bit 4:2:0");
        }
        if !info.frame_mbs_only {
            return refuse("encode-stream-unsupported", "the stream may code fields");
        }
        if codec == Codec::Avc && info.poc_type == 1 {
            return refuse("encode-stream-unsupported", "pic_order_cnt_type 1 is not read by this revision");
        }
        // §4.3 item 6: no reordering. HEVC by the SPS; AVC by the VUI where it says, and by
        // the picture order count, which must increase in decode order from each IDR.
        match codec {
            Codec::Hevc => {
                if info.max_num_reorder != Some(0) {
                    return refuse("encode-stream-reordered", "sps_max_num_reorder_pics is not 0");
                }
            }
            Codec::Avc => {
                if info.max_num_reorder.is_some_and(|m| m != 0) {
                    return refuse("encode-stream-reordered", "the VUI's max_num_reorder_frames is not 0");
                }
                if info.poc_type == 0 {
                    let rb = rbsp(codec, &s[fv.range.clone()]);
                    let mut r = Bits::new(&rb);
                    r.ue()?;
                    r.ue()?;
                    r.ue()?;
                    r.skip(u32::from(info.log2_max_frame_num))?;
                    if is_idr {
                        r.ue()?;
                        prev_msb = 0;
                        prev_lsb = 0;
                        last_poc = i64::MIN;
                    }
                    let lsb = r.u(u32::from(info.log2_max_poc_lsb))? as i64;
                    let max = 1i64 << info.log2_max_poc_lsb;
                    let msb = if lsb < prev_lsb && prev_lsb - lsb >= max / 2 {
                        prev_msb + max
                    } else if lsb > prev_lsb && lsb - prev_lsb > max / 2 {
                        prev_msb - max
                    } else {
                        prev_msb
                    };
                    let poc = msb + lsb;
                    if poc <= last_poc {
                        return refuse("encode-stream-reordered", format!("access unit {a}'s picture order count {poc} does not follow {last_poc}"));
                    }
                    last_poc = poc;
                    if fv.aux != 0 {
                        prev_msb = msb;
                        prev_lsb = lsb;
                    }
                }
            }
        }
        idr.push(is_idr);
    }
    // Byte ranges: from each access unit's first start code prefix.
    let starts: Vec<usize> = aus.iter().map(|au| prefix_start(s, nals[au.start].range.start)).collect();
    let units = starts.iter().enumerate().map(|(i, &b)| b..starts.get(i + 1).copied().unwrap_or(s.len())).collect();
    Ok(Checked { units, sps: sps.expect("checked above"), idr, nal_units: aus })
}

/// A piece's payload ([platform-encode] §9.1): the NAL units of access units
/// `units` of a checked stream, each after `00 00 00 01`.
///
/// [platform-encode]: ../../../../extensions/platform-encode.md
pub fn piece_payload(codec: Codec, s: &[u8], units: std::ops::Range<usize>) -> Result<Vec<u8>, Reason> {
    let nals = nal_units(codec, s)?;
    let aus = access_units(codec, s, &nals)?;
    let mut out = Vec::new();
    for au in aus.get(units).ok_or_else(|| Reason { code: "encode-frame-count", detail: "a piece past the stream's access units".into() })? {
        for n in &nals[au.clone()] {
            out.extend_from_slice(&[0, 0, 0, 1]);
            out.extend_from_slice(&s[n.range.clone()]);
        }
    }
    Ok(out)
}

/// The first byte of the start code prefix (with any zero bytes before it that follow
/// the previous NAL unit's trailing zeros) of the NAL unit beginning at `nal_start`.
fn prefix_start(s: &[u8], nal_start: usize) -> usize {
    let mut b = nal_start - 3;
    while b > 0 && s[b - 1] == 0 {
        b -= 1;
    }
    b
}

/// Rewrite NAL units each preceded by a big-endian length of `len_size` bytes (the form
/// ISO/IEC 14496-15 stores) as Annex B NAL units, each after `00 00 00 01`.
pub fn length_prefixed_to_annexb(data: &[u8], len_size: usize) -> Result<Vec<u8>, Reason> {
    if !(1..=4).contains(&len_size) {
        return refuse(MALFORMED, format!("a NAL length field of {len_size} bytes"));
    }
    let mut out = Vec::with_capacity(data.len() + 16);
    let mut at = 0usize;
    while at < data.len() {
        let head = data.get(at..at + len_size).ok_or_else(|| Reason { code: MALFORMED, detail: "a NAL length field runs past the sample".into() })?;
        let n = head.iter().fold(0usize, |v, &b| (v << 8) | b as usize);
        at += len_size;
        let nal = data.get(at..at + n).filter(|n| !n.is_empty()).ok_or_else(|| Reason { code: MALFORMED, detail: "a NAL unit runs past the sample".into() })?;
        out.extend_from_slice(&[0, 0, 0, 1]);
        out.extend_from_slice(nal);
        at += n;
    }
    Ok(out)
}

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

    #[test]
    fn exp_golomb_reads_the_table_of_9_1() {
        // 1, 010, 011, 00100, 00101 → 0, 1, 2, 3, 4; then se: 010 → 1, 011 → −1.
        let bits = [0b1010_0110, 0b0100_0010, 0b1010_0110];
        let mut r = Bits::new(&bits);
        assert_eq!([r.ue().unwrap(), r.ue().unwrap(), r.ue().unwrap(), r.ue().unwrap(), r.ue().unwrap()], [0, 1, 2, 3, 4]);
        assert_eq!([r.se().unwrap(), r.se().unwrap()], [1, -1]);
    }

    #[test]
    fn emulation_prevention_is_removed() {
        assert_eq!(rbsp(Codec::Avc, &[0x67, 0, 0, 3, 1, 0, 0, 3, 0, 3]), vec![0, 0, 1, 0, 0, 0, 3]);
    }

    #[test]
    fn nal_units_split_on_three_and_four_byte_prefixes() {
        let s = [0, 0, 0, 1, 0x67, 1, 2, 0, 0, 1, 0x68, 3, 0, 0, 0, 0, 1, 0x65, 9];
        let n = nal_units(Codec::Avc, &s).unwrap();
        assert_eq!(n.iter().map(|n| (n.range.clone(), n.kind)).collect::<Vec<_>>(), vec![(4..7, 7), (10..12, 8), (17..19, 5)]);
        assert_eq!(nal_units(Codec::Avc, &[1, 0x67]).unwrap_err().code, MALFORMED);
        assert_eq!(nal_units(Codec::Avc, &[0, 0, 1, 0x67, 0, 0, 2, 1]).unwrap_err().code, MALFORMED);
        assert_eq!(nal_units(Codec::Avc, &[0, 0, 1, 0xe7, 1]).unwrap_err().code, MALFORMED);
    }

    #[test]
    fn length_prefixed_samples_become_annex_b() {
        assert_eq!(length_prefixed_to_annexb(&[0, 0, 0, 2, 0x65, 1, 0, 0, 0, 1, 0x06], 4).unwrap(), vec![0, 0, 0, 1, 0x65, 1, 0, 0, 0, 1, 0x06]);
        assert!(length_prefixed_to_annexb(&[0, 0, 0, 9, 0x65], 4).is_err());
    }
}
