//! The H.265 probe (`wai.video.hevc`, `wai.video.mvhevc`; format-probes §4): the Annex B
//! byte stream's NAL units, the video parameter set with its multilayer extension
//! (Annex F), and the base layer's sequence parameter set. Syntax from ITU-T H.265 (V11,
//! 01/2026) §7.3.1.2, §7.3.2.1, §7.3.2.2.1, §7.3.3, §E.2.2–§E.2.3, §F.7.3.2.1 and
//! §F.7.3.2.1.1; derivations from §F.7.4.3.1.1 (F-3)–(F-6).

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

pub const NAL_VPS: u64 = 32;
pub const NAL_SPS: u64 = 33;
pub const NAL_PPS: u64 = 34;

/// One `profile_tier_level( )`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Ptl {
    /// `general_profile_idc`, or `None` when the structure carries no profile.
    pub profile_idc: Option<u64>,
    pub tier: Option<u64>,
    /// `general_profile_compatibility_flag[ j ]`, bit `j` = flag `j`.
    pub compat: u32,
    pub level_idc: u64,
}

impl Ptl {
    /// Whether the structure indicates profile `idc`, by its idc or its compatibility flag.
    pub fn indicates(&self, idc: u64) -> bool {
        self.profile_idc == Some(idc) || (idc < 32 && self.compat & (1 << idc) != 0)
    }
}

/// `profile_tier_level( profilePresentFlag, maxNumSubLayersMinus1 )` (§7.3.3). The general
/// profile part is 88 bits whatever its branches (the syntax's own note), so it is read as
/// fields and a fixed skip.
pub fn ptl(r: &mut BitReader, profile_present: bool, max_sub_layers_minus1: u64) -> Result<Ptl, Refusal> {
    let (mut profile_idc, mut tier, mut compat) = (None, None, 0u32);
    if profile_present {
        let _space = r.u(2)?;
        tier = Some(r.u(1)?);
        profile_idc = Some(r.u(5)?);
        compat = r.u(32)? as u32;
        // progressive, interlaced, non-packed, frame-only (4), then 43 constraint bits,
        // then general_inbld_flag or general_reserved_zero_bit.
        r.skip(4 + 43 + 1)?;
    }
    let level_idc = r.u(8)?;
    let n = max_sub_layers_minus1 as usize;
    let mut present = Vec::with_capacity(n);
    for _ in 0..n {
        present.push((r.flag()?, r.flag()?));
    }
    if n > 0 {
        r.skip((8 - n) * 2)?;
    }
    for &(profile, level) in &present {
        if profile {
            r.skip(88)?;
        }
        if level {
            r.skip(8)?;
        }
    }
    Ok(Ptl { profile_idc, tier, compat, level_idc })
}

/// The common information of the previous `hrd_parameters( )` a VPS carries, which one
/// with `cprms_present_flag` equal to 0 reuses (§E.2.2).
#[derive(Default, Clone, Copy)]
struct HrdCommon {
    nal: bool,
    vcl: bool,
    sub_pic: bool,
}

/// `hrd_parameters( commonInfPresentFlag, maxNumSubLayersMinus1 )` (§E.2.2), read and
/// discarded.
fn hrd(r: &mut BitReader, common: bool, prev: HrdCommon, max_sub: u64) -> Result<HrdCommon, Refusal> {
    let mut c = prev;
    if common {
        c.nal = r.flag()?;
        c.vcl = r.flag()?;
        c.sub_pic = false;
        if c.nal || c.vcl {
            c.sub_pic = r.flag()?;
            if c.sub_pic {
                r.skip(8 + 5 + 1 + 5)?;
            }
            r.skip(4 + 4)?;
            if c.sub_pic {
                r.skip(4)?;
            }
            r.skip(5 + 5 + 5)?;
        }
    }
    for _ in 0..=max_sub {
        let general = r.flag()?;
        let within = if general { true } else { r.flag()? };
        let mut low_delay = false;
        if within {
            r.ue()?;
        } else {
            low_delay = r.flag()?;
        }
        let cpb_cnt = if low_delay { 1 } else { r.ue()? + 1 };
        if cpb_cnt > 32 {
            return Err(Refusal::ReservedValue);
        }
        for present in [c.nal, c.vcl] {
            if present {
                for _ in 0..cpb_cnt {
                    r.ue()?;
                    r.ue()?;
                    if c.sub_pic {
                        r.ue()?;
                        r.ue()?;
                    }
                    r.skip(1)?;
                }
            }
        }
    }
    Ok(c)
}

/// What the probe reads from a video parameter set.
#[derive(Debug, Clone, Default)]
pub struct Vps {
    pub base_layer_internal: bool,
    pub max_layers_minus1: u64,
    pub max_sub_layers_minus1: u64,
    pub base_ptl: Option<Ptl>,
    /// Present when `vps_extension_flag` is 1 and `vps_max_layers_minus1` > 0.
    pub ext: Option<VpsExt>,
}

#[derive(Debug, Clone, Default)]
pub struct VpsExt {
    pub scalability_mask: u16,
    /// `layer_id_in_nuh[ i ]`, for i from 0 to MaxLayersMinus1.
    pub layer_ids: Vec<u64>,
    /// `ViewOrderIdx[ layer_id_in_nuh[ i ] ]`.
    pub view_order: Vec<u64>,
    /// `AuxId[ layer_id_in_nuh[ i ] ]`.
    pub aux: Vec<u64>,
    pub num_views: u64,
    /// `view_id_val[ i ]`, for each view order index.
    pub view_ids: Vec<u64>,
    /// `direct_dependency_flag[ i ][ j ]`.
    pub deps: Vec<Vec<bool>>,
    /// The profile_tier_level( ) structures of the extension, in order (the first, when
    /// present, carries no profile).
    pub ptls: Vec<Ptl>,
}

/// `video_parameter_set_rbsp( )` (§F.7.3.2.1) through `vps_extension( )` (§F.7.3.2.1.1)
/// up to and including its profile_tier_level( ) structures.
pub fn vps(rbsp: &[u8]) -> Result<Vps, Refusal> {
    let mut r = BitReader::new(rbsp);
    let _id = r.u(4)?;
    let base_layer_internal = r.flag()?;
    let _base_layer_available = r.flag()?;
    let max_layers_minus1 = r.u(6)?;
    let max_sub = r.u(3)?;
    if max_sub > 6 {
        return Err(Refusal::ReservedValue);
    }
    let _nesting = r.flag()?;
    if r.u(16)? != 0xffff {
        return Err(Refusal::ReservedValue);
    }
    let base_ptl = ptl(&mut r, true, max_sub)?;
    let ordering = r.flag()?;
    for _ in (if ordering { 0 } else { max_sub })..=max_sub {
        r.ue()?;
        r.ue()?;
        r.ue()?;
    }
    let max_layer_id = r.u(6)?;
    let num_layer_sets_minus1 = r.ue()?;
    if num_layer_sets_minus1 > 1023 {
        return Err(Refusal::ReservedValue);
    }
    r.skip(num_layer_sets_minus1 as usize * (max_layer_id as usize + 1))?;
    if r.flag()? {
        r.skip(64)?;
        if r.flag()? {
            r.ue()?;
        }
        let num_hrd = r.ue()?;
        if num_hrd > num_layer_sets_minus1 + 1 {
            return Err(Refusal::ReservedValue);
        }
        let mut prev = HrdCommon::default();
        for i in 0..num_hrd {
            r.ue()?;
            let cprms = if i > 0 { r.flag()? } else { true };
            prev = hrd(&mut r, cprms, prev, max_sub)?;
        }
    }
    let mut out = Vps { base_layer_internal, max_layers_minus1, max_sub_layers_minus1: max_sub, base_ptl: Some(base_ptl), ext: None };
    if !r.flag()? {
        return Ok(out);
    }
    r.align(Some(1))?;
    if max_layers_minus1 == 0 {
        // An extension with one layer describes no further layer: nothing the probe reads.
        return Ok(out);
    }
    out.ext = Some(vps_extension(&mut r, base_layer_internal, max_layers_minus1, max_sub, base_ptl)?);
    Ok(out)
}

fn ceil_log2(n: u64) -> u32 {
    if n <= 1 { 0 } else { 64 - (n - 1).leading_zeros() }
}

/// A structure with no profile takes the previous structure's profile, tier and
/// compatibility flags (§F.7.4.4: VpsProfileTierLevel[ k ] from VpsProfileTierLevel[ k − 1 ]).
fn inherit(mut p: Ptl, prev: &Ptl) -> Ptl {
    if p.profile_idc.is_none() {
        p.profile_idc = prev.profile_idc;
        p.tier = prev.tier;
        p.compat = prev.compat;
    }
    p
}

fn vps_extension(r: &mut BitReader, base_internal: bool, max_layers_minus1: u64, max_sub: u64, base_ptl: Ptl) -> Result<VpsExt, Refusal> {
    let mut ptls = Vec::new();
    if base_internal {
        ptls.push(inherit(ptl(r, false, max_sub)?, &base_ptl));
    }
    let splitting = r.flag()?;
    let mut mask = 0u16;
    for i in 0..16 {
        if r.flag()? {
            mask |= 1 << i;
        }
    }
    let types = mask.count_ones() as usize;
    let mut dim_len = vec![0u32; types];
    for len in dim_len.iter_mut().take(types.saturating_sub(usize::from(splitting))) {
        *len = r.u(3)? as u32 + 1;
    }
    let max_l = max_layers_minus1.min(62) as usize;
    if splitting && types > 0 {
        // (F-2): the last length makes the offsets reach 6 bits.
        let used: u32 = dim_len[..types - 1].iter().sum();
        if used >= 6 {
            return Err(Refusal::ReservedValue);
        }
        dim_len[types - 1] = 6 - used;
    }
    let id_present = r.flag()?;
    let mut layer_ids = vec![0u64; max_l + 1];
    let mut dims = vec![vec![0u64; types]; max_l + 1];
    for i in 1..=max_l {
        layer_ids[i] = if id_present { r.u(6)? } else { i as u64 };
        if layer_ids[i] <= layer_ids[i - 1] {
            return Err(Refusal::ReservedValue);
        }
        if !splitting {
            for j in 0..types {
                dims[i][j] = r.u(dim_len[j])?;
            }
        }
    }
    if splitting {
        for i in 0..=max_l {
            let mut off = 0u32;
            for j in 0..types {
                dims[i][j] = (layer_ids[i] >> off) & ((1 << dim_len[j]) - 1);
                off += dim_len[j];
            }
        }
    }
    // (F-3): ScalabilityId, ViewOrderIdx, AuxId and NumViews.
    let mut view_order = vec![0u64; max_l + 1];
    let mut aux = vec![0u64; max_l + 1];
    let mut num_views = 1u64;
    for i in 0..=max_l {
        let mut j = 0;
        let mut sid = [0u64; 16];
        for (sm, s) in sid.iter_mut().enumerate() {
            if mask & (1 << sm) != 0 {
                *s = dims[i][j];
                j += 1;
            }
        }
        view_order[i] = sid[1];
        aux[i] = sid[3];
        if i > 0 && !(0..i).any(|k| view_order[k] == view_order[i]) {
            num_views += 1;
        }
    }
    let view_id_len = r.u(4)? as u32;
    let mut view_ids = Vec::new();
    if view_id_len > 0 {
        for _ in 0..num_views {
            view_ids.push(r.u(view_id_len)?);
        }
    } else {
        view_ids = vec![0; num_views as usize];
    }
    let mut deps = vec![vec![false; max_l + 1]; max_l + 1];
    for (i, row) in deps.iter_mut().enumerate().skip(1) {
        for cell in row.iter_mut().take(i) {
            *cell = r.flag()?;
        }
    }
    // (F-4)–(F-6): the independent layers and their tree partitions.
    let n = max_l + 1;
    let mut dep = deps.clone();
    for i in 0..n {
        for k in 0..i {
            if deps[i][k] {
                let row_k = dep[k].clone();
                for (cell, &v) in dep[i].iter_mut().zip(&row_k) {
                    *cell |= v;
                }
            }
        }
    }
    let mut tree_sizes = Vec::new();
    let mut in_list = [false; 64];
    for i in 0..n {
        if (0..n).any(|j| deps[i][j]) {
            continue;
        }
        let mut h = 1u64;
        for j in 0..n {
            if dep[j][i] {
                let lid = layer_ids[j] as usize;
                if !in_list[lid] {
                    in_list[lid] = true;
                    h += 1;
                }
            }
        }
        tree_sizes.push(h);
    }
    let independent = tree_sizes.len();
    let num_add_layer_sets = if independent > 1 { r.ue()? } else { 0 };
    if num_add_layer_sets > 1023 {
        return Err(Refusal::ReservedValue);
    }
    for _ in 0..num_add_layer_sets {
        for &size in tree_sizes.iter().skip(1) {
            r.skip(ceil_log2(size + 1) as usize)?;
        }
    }
    if r.flag()? {
        r.skip(3 * n)?;
    }
    if r.flag()? {
        // max_tid_il_ref_pics_plus1[ i ][ j ] for each j above i that depends on i.
        for i in 0..max_l {
            r.skip(3 * deps.iter().skip(i + 1).filter(|row| row[i]).count())?;
        }
    }
    let _default_ref_layers_active = r.flag()?;
    let num_ptl_minus1 = r.ue()?;
    if num_ptl_minus1 > 63 {
        return Err(Refusal::ReservedValue);
    }
    let first = if base_internal { 2 } else { 1 };
    for _ in first..=num_ptl_minus1 {
        let present = r.flag()?;
        let p = ptl(r, present, max_sub)?;
        let prev = ptls.last().copied().unwrap_or(base_ptl);
        ptls.push(inherit(p, &prev));
    }
    Ok(VpsExt { scalability_mask: mask, layer_ids, view_order, aux, num_views, view_ids, deps, ptls })
}

/// What the probe reads from the base layer's sequence parameter set (§7.3.2.2.1).
#[derive(Debug, Clone, Copy)]
pub struct Sps {
    pub ptl: Ptl,
    pub chroma_format_idc: u64,
    pub width: u64,
    pub height: u64,
    pub crop: [u64; 4],
    pub bit_depth_luma: u64,
    pub bit_depth_chroma: u64,
}

pub fn sps(rbsp: &[u8]) -> Result<Sps, Refusal> {
    let mut r = BitReader::new(rbsp);
    let _vps_id = r.u(4)?;
    let max_sub = r.u(3)?;
    if max_sub > 6 {
        return Err(Refusal::ReservedValue);
    }
    let _nesting = r.flag()?;
    let p = ptl(&mut r, true, max_sub)?;
    let id = r.ue()?;
    if id > 15 {
        return Err(Refusal::ReservedValue);
    }
    let chroma = r.ue()?;
    if chroma > 3 {
        return Err(Refusal::ReservedValue);
    }
    if chroma == 3 {
        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()?;
        }
    }
    let bdl = r.ue()?;
    let bdc = r.ue()?;
    if bdl > 8 || bdc > 8 {
        return Err(Refusal::ReservedValue);
    }
    Ok(Sps { ptl: p, chroma_format_idc: chroma, width, height, crop, bit_depth_luma: bdl + 8, bit_depth_chroma: bdc + 8 })
}

/// The cropped size: the coded size less the conformance window, whose offsets count
/// chroma samples (§7.4.3.2.1, SubWidthC and SubHeightC).
pub fn cropped(chroma: u64, w: u64, h: u64, crop: [u64; 4]) -> Result<(u64, u64), Refusal> {
    let (sw, sh) = match chroma {
        1 => (2, 2),
        2 => (2, 1),
        _ => (1, 1),
    };
    let dw = sw * (crop[0] + crop[1]);
    let dh = sh * (crop[2] + crop[3]);
    if dw >= w || dh >= h {
        return Err(Refusal::ReservedValue);
    }
    Ok((w - dw, h - dh))
}

/// The probe. `multiview` is true for `wai.video.mvhevc`, which further requires a
/// multiview layer (§4.3).
pub fn probe(stream: &[u8], multiview: bool) -> Result<Value, Refusal> {
    let units = annex_b_units(stream)?;
    // Whether the stream holds a VPS and a base-layer SPS anywhere: a slice before them is out
    // of order if they come later, and the stream lacks them if they never do.
    let has = |t: u64| units.iter().any(|u| u.bytes.len() >= 2 && u64::from((u.bytes[0] >> 1) & 0x3f) == t && (u16::from_be_bytes([u.bytes[0], u.bytes[1]]) >> 3) & 0x3f == 0);
    let complete = has(NAL_VPS) && has(NAL_SPS);
    let mut vps_info: Option<Vps> = None;
    let mut base_sps: Option<Sps> = None;
    let mut layers_seen = std::collections::BTreeSet::new();
    let mut vcl_units = 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]]);
        if h & 0x8000 != 0 {
            return Err(Refusal::ReservedValue);
        }
        let nut = u64::from((h >> 9) & 0x3f);
        let layer = u64::from((h >> 3) & 0x3f);
        let tid_plus1 = h & 7;
        if tid_plus1 == 0 {
            return Err(Refusal::ReservedValue);
        }
        let rbsp = unescape(&u.bytes[2..]);
        if nut < 32 {
            if base_sps.is_none() || vps_info.is_none() {
                return Err(if complete { Refusal::OutOfOrder } else { Refusal::MissingStructure });
            }
            vcl_units += 1;
            layers_seen.insert(layer);
        } else if nut == NAL_VPS && vps_info.is_none() {
            vps_info = Some(vps(&rbsp)?);
        } else if nut == NAL_SPS && layer == 0 && base_sps.is_none() {
            base_sps = Some(sps(&rbsp)?);
        }
    }
    let (Some(v), Some(s)) = (vps_info, base_sps) else { return Err(Refusal::MissingStructure) };
    if vcl_units == 0 {
        return Err(Refusal::MissingStructure);
    }
    let (cw, ch) = cropped(s.chroma_format_idc, s.width, s.height, s.crop)?;
    let views = v.ext.as_ref().map_or(1, |e| if e.scalability_mask & 2 != 0 { e.num_views } else { 1 });
    let mv_profile = v.ext.as_ref().is_some_and(|e| e.ptls.iter().any(|p| p.indicates(6)));
    let enhancement_present = layers_seen.iter().any(|&l| l > 0);
    if multiview && (views < 2 || !enhancement_present) {
        return Err(Refusal::NotMultiview);
    }
    let mut out = json!({
        "profile_idc": s.ptl.profile_idc,
        "tier": s.ptl.tier,
        "level_idc": s.ptl.level_idc,
        "chroma_format_idc": s.chroma_format_idc,
        "bit_depth": [s.bit_depth_luma, s.bit_depth_chroma],
        "coded": [s.width, s.height],
        "cropped": [cw, ch],
        "layers": v.max_layers_minus1.min(62) + 1,
        "layers_present": layers_seen.iter().copied().collect::<Vec<_>>(),
    });
    if let Some(e) = &v.ext {
        out["multilayer"] = json!({
            "scalability_mask": e.scalability_mask,
            "layer_ids": e.layer_ids,
            "view_order": e.view_order,
            "aux": e.aux,
            "views": views,
            "view_ids": e.view_ids,
            "profiles": e.ptls.iter().map(|p| p.profile_idc).collect::<Vec<_>>(),
            "multiview_profile": mv_profile,
        });
    }
    Ok(out)
}
