//! The H.266 probe (`wai.video.vvc`; format-probes §5): the Annex B byte stream's NAL
//! units, the start of the video parameter set, and the sequence parameter set through its
//! conformance window. Syntax from ITU-T H.266 (V4, 01/2026) §7.3.1.2, §7.3.2.3, §7.3.2.4,
//! §7.3.3.1 and §7.3.3.2.

use super::bits::{annex_b_units, unescape, BitReader};
use super::hevc::cropped;
use super::Refusal;
use serde_json::{json, Value};

pub const NAL_VPS: u64 = 14;
pub const NAL_SPS: u64 = 15;

/// `general_constraints_info( )` (§7.3.3.2): its 71 fixed bits, the additional bits and
/// the zero alignment, read and discarded.
fn gci(r: &mut BitReader) -> Result<(), Refusal> {
    if r.flag()? {
        r.skip(71)?;
        // gci_num_additional_bits: six named flags when it exceeds 5, then reserved bits,
        // so as many bits as it states either way.
        let n = r.u(8)? as usize;
        r.skip(n)?;
    }
    r.align(Some(0))
}

#[derive(Debug, Clone)]
pub struct Ptl {
    pub profile_idc: Option<u64>,
    pub tier: Option<u64>,
    pub level_idc: u64,
    pub frame_only: bool,
    pub multilayer: bool,
    pub sub_profiles: Vec<u64>,
}

/// `profile_tier_level( profileTierPresentFlag, MaxNumSubLayersMinus1 )` (§7.3.3.1).
pub fn ptl(r: &mut BitReader, present: bool, max_sub: u64) -> Result<Ptl, Refusal> {
    let (mut profile_idc, mut tier) = (None, None);
    if present {
        profile_idc = Some(r.u(7)?);
        tier = Some(r.u(1)?);
    }
    let level_idc = r.u(8)?;
    let frame_only = r.flag()?;
    let multilayer = r.flag()?;
    if present {
        gci(r)?;
    }
    let mut level_present = Vec::new();
    for _ in 0..max_sub {
        level_present.push(r.flag()?);
    }
    r.align(None)?;
    for &p in &level_present {
        if p {
            r.skip(8)?;
        }
    }
    let mut sub_profiles = Vec::new();
    if present {
        let n = r.u(8)?;
        for _ in 0..n {
            sub_profiles.push(r.u(32)?);
        }
    }
    Ok(Ptl { profile_idc, tier, level_idc, frame_only, multilayer, sub_profiles })
}

#[derive(Debug, Clone)]
pub struct Sps {
    pub ptl: Option<Ptl>,
    pub chroma_format_idc: u64,
    pub ctu_size: u64,
    pub width: u64,
    pub height: u64,
    pub crop: [u64; 4],
}

/// `seq_parameter_set_rbsp( )` (§7.3.2.4) through the conformance window.
pub fn sps(rbsp: &[u8]) -> Result<Sps, Refusal> {
    let mut r = BitReader::new(rbsp);
    let _id = r.u(4)?;
    let _vps_id = r.u(4)?;
    let max_sub = r.u(3)?;
    if max_sub > 6 {
        return Err(Refusal::ReservedValue);
    }
    let chroma = r.u(2)?;
    let ctu_minus5 = r.u(2)?;
    if ctu_minus5 > 2 {
        return Err(Refusal::ReservedValue);
    }
    let ptl_present = r.flag()?;
    let p = if ptl_present { Some(ptl(&mut r, true, max_sub)?) } else { None };
    let _gdr = r.flag()?;
    if r.flag()? {
        r.skip(1)?;
    }
    let width = r.ue()?;
    let height = r.ue()?;
    if width == 0 || height == 0 {
        return Err(Refusal::ReservedValue);
    }
    let mut crop = [0u64; 4];
    if r.flag()? {
        for c in &mut crop {
            *c = r.ue()?;
        }
    }
    Ok(Sps { ptl: p, chroma_format_idc: chroma, ctu_size: 1 << (ctu_minus5 + 5), width, height, crop })
}

pub fn probe(stream: &[u8]) -> Result<Value, Refusal> {
    let units = annex_b_units(stream)?;
    let has_sps = units.iter().any(|u| u.bytes.len() >= 2 && u64::from(u.bytes[1] >> 3) == NAL_SPS);
    let mut s: Option<Sps> = None;
    let mut vps_layers: Option<u64> = None;
    let mut layers_seen = std::collections::BTreeSet::new();
    let mut vcl = 0u64;
    for u in &units {
        if u.bytes.len() < 2 {
            return Err(Refusal::Truncated);
        }
        let h = u16::from_be_bytes([u.bytes[0], u.bytes[1]]);
        // forbidden_zero_bit and nuh_reserved_zero_bit.
        if h & 0xc000 != 0 {
            return Err(Refusal::ReservedValue);
        }
        let layer = u64::from((h >> 8) & 0x3f);
        let nut = u64::from((h >> 3) & 0x1f);
        if h & 7 == 0 {
            return Err(Refusal::ReservedValue);
        }
        if layer > 55 {
            return Err(Refusal::ReservedValue);
        }
        let rbsp = unescape(&u.bytes[2..]);
        if nut <= 11 {
            if s.is_none() {
                return Err(if has_sps { Refusal::OutOfOrder } else { Refusal::MissingStructure });
            }
            vcl += 1;
            layers_seen.insert(layer);
        } else if nut == NAL_VPS && vps_layers.is_none() {
            let mut r = BitReader::new(&rbsp);
            let _id = r.u(4)?;
            vps_layers = Some(r.u(6)? + 1);
        } else if nut == NAL_SPS && s.is_none() {
            s = Some(sps(&rbsp)?);
        }
    }
    let Some(s) = s else { return Err(Refusal::MissingStructure) };
    if vcl == 0 {
        return Err(Refusal::MissingStructure);
    }
    let (cw, ch) = cropped(s.chroma_format_idc, s.width, s.height, s.crop)?;
    let p = s.ptl.as_ref();
    Ok(json!({
        "profile_idc": p.and_then(|p| p.profile_idc),
        "tier": p.and_then(|p| p.tier),
        "level_idc": p.map(|p| p.level_idc),
        "frame_only": p.map(|p| p.frame_only),
        "multilayer_enabled": p.map(|p| p.multilayer),
        "sub_profiles": p.map(|p| p.sub_profiles.clone()).unwrap_or_default(),
        "chroma_format_idc": s.chroma_format_idc,
        "ctu_size": s.ctu_size,
        "coded": [s.width, s.height],
        "cropped": [cw, ch],
        "layers": vps_layers.unwrap_or(1),
        "layers_present": layers_seen.iter().copied().collect::<Vec<_>>(),
    }))
}
