//! The LCEVC probe (`wai.video.lcevc`; format-probes §6): the byte stream's NAL units, each
//! one's process blocks, and the first sequence and global configurations. ISO/IEC 23094-2's
//! text is not openly readable; the layout is the one two independent open implementations
//! agree on, each citing its clauses (format-probes §6.1).

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

/// The sizes of `resolution_type` 1 to 50 (ISO/IEC 23094-2 Table 20, as the open decoder
/// reproduces it); 0 and 51 to 62 are not sizes, and 63 is a custom size.
pub const RESOLUTIONS: [(u16, u16); 50] = [
    (360, 200), (400, 240), (480, 320), (640, 360), (640, 480), (768, 480), (800, 600),
    (852, 480), (854, 480), (856, 480), (960, 540), (960, 640), (1024, 576), (1024, 600),
    (1024, 768), (1152, 864), (1280, 720), (1280, 800), (1280, 1024), (1360, 768),
    (1366, 768), (1400, 1050), (1440, 900), (1600, 1200), (1680, 1050), (1920, 1080),
    (1920, 1200), (2048, 1080), (2048, 1152), (2048, 1536), (2160, 1440), (2560, 1440),
    (2560, 1600), (2560, 2048), (3200, 1800), (3200, 2048), (3200, 2400), (3440, 1440),
    (3840, 1600), (3840, 2160), (3840, 2400), (4096, 2160), (4096, 3072), (5120, 2880),
    (5120, 3200), (5120, 4096), (6400, 4096), (6400, 4800), (7680, 4320), (7680, 4800),
];

pub const NUT_NON_IDR: u8 = 28;
pub const NUT_IDR: u8 = 29;

pub const BLOCK_SEQUENCE_CONFIG: u8 = 0;
pub const BLOCK_GLOBAL_CONFIG: u8 = 1;
pub const BLOCK_PICTURE_CONFIG: u8 = 2;

/// A multi-byte value: 7 bits a byte, most significant first, while the byte's top bit is
/// set; at most ten bytes.
fn multibyte(b: &[u8], at: &mut usize) -> Result<u64, Refusal> {
    let mut v: u64 = 0;
    for _ in 0..10 {
        let x = *b.get(*at).ok_or(Refusal::Truncated)?;
        *at += 1;
        v = (v << 7) | u64::from(x & 0x7f);
        if x & 0x80 == 0 {
            return Ok(v);
        }
    }
    Err(Refusal::ReservedValue)
}

/// The blocks of one NAL unit's payload (after its two header bytes, emulation prevention
/// removed, the final `0x80` stop byte removed): `(type, body)`.
pub fn blocks(payload: &[u8]) -> Result<Vec<(u8, &[u8])>, Refusal> {
    let mut out = Vec::new();
    let mut at = 0;
    while at < payload.len() {
        let h = payload[at];
        at += 1;
        let (size_type, kind) = (h >> 5, h & 0x1f);
        let size = match size_type {
            0..=5 => u64::from(size_type),
            6 => return Err(Refusal::ReservedValue),
            _ => multibyte(payload, &mut at)?,
        } as usize;
        let body = payload.get(at..at.checked_add(size).ok_or(Refusal::LengthMismatch)?).ok_or(Refusal::LengthMismatch)?;
        out.push((kind, body));
        at += size;
    }
    Ok(out)
}

#[derive(Debug, Clone)]
pub struct SequenceConfig {
    pub profile_idc: u64,
    pub level_idc: u64,
    pub sublevel_idc: u64,
    pub conformance_window: Option<[u64; 4]>,
}

pub fn sequence_config(b: &[u8]) -> Result<SequenceConfig, Refusal> {
    let mut r = BitReader::new(b);
    let profile_idc = r.u(4)?;
    let level_idc = r.u(4)?;
    let sublevel_idc = r.u(2)?;
    let window = r.flag()?;
    if r.u(5)? != 0 {
        return Err(Refusal::ReservedValue);
    }
    if profile_idc == 15 || level_idc == 15 {
        // The extended profile and level: the two open readings of their byte disagree.
        return Err(Refusal::OutsideProbe);
    }
    let mut at = r.position() / 8;
    let conformance_window = if window {
        let mut w = [0u64; 4];
        for x in &mut w {
            *x = multibyte(b, &mut at)?;
        }
        Some(w)
    } else {
        None
    };
    if at != b.len() {
        return Err(Refusal::LengthMismatch);
    }
    Ok(SequenceConfig { profile_idc, level_idc, sublevel_idc, conformance_window })
}

#[derive(Debug, Clone)]
pub struct GlobalConfig {
    pub resolution_type: u64,
    pub transform_type: u64,
    pub chroma_sampling_type: u64,
    pub base_depth_type: u64,
    pub enhancement_depth_type: u64,
    pub temporal_enabled: bool,
    pub upsample_type: u64,
    pub tile_dimensions_type: u64,
    pub custom_size: Option<[u64; 2]>,
}

impl GlobalConfig {
    /// The enhanced picture's size: Table 20's for a listed type, the custom size for 63.
    pub fn size(&self) -> [u64; 2] {
        match self.custom_size {
            Some(s) => s,
            None => {
                let (w, h) = RESOLUTIONS[self.resolution_type as usize - 1];
                [u64::from(w), u64::from(h)]
            }
        }
    }
}

/// The four fixed bytes of the global configuration and the conditional fields up to the
/// custom resolution (format-probes §6.4).
pub fn global_config(b: &[u8]) -> Result<GlobalConfig, Refusal> {
    let mut r = BitReader::new(b);
    let planes_flag = r.flag()?;
    let resolution_type = r.u(6)?;
    let transform_type = r.u(1)?;
    let chroma_sampling_type = r.u(2)?;
    let base_depth_type = r.u(2)?;
    let enhancement_depth_type = r.u(2)?;
    let tsw_modifier = r.flag()?;
    let _predicted_residual = r.flag()?;
    let _temporal_tile_intra = r.flag()?;
    let temporal_enabled = r.flag()?;
    let upsample_type = r.u(3)?;
    let l1_filter = r.flag()?;
    let _scaling_l1 = r.u(2)?;
    let _scaling_l2 = r.u(2)?;
    let tile_dimensions_type = r.u(2)?;
    let _user_data = r.u(2)?;
    let _l1_depth = r.flag()?;
    let _chroma_step = r.flag()?;
    if planes_flag {
        r.skip(8)?;
    }
    if tsw_modifier {
        r.skip(8)?;
    }
    if upsample_type == 4 {
        r.skip(64)?;
    }
    if l1_filter {
        r.skip(8)?;
    }
    if tile_dimensions_type > 0 {
        if tile_dimensions_type == 3 {
            r.skip(32)?;
        }
        r.skip(8)?;
    }
    if resolution_type == 0 || (51..63).contains(&resolution_type) || upsample_type > 4 {
        return Err(Refusal::ReservedValue);
    }
    let custom_size = if resolution_type == 63 { Some([r.u(16)?, r.u(16)?]) } else { None };
    if custom_size.is_some_and(|[w, h]| w == 0 || h == 0) {
        return Err(Refusal::ReservedValue);
    }
    Ok(GlobalConfig {
        resolution_type, transform_type, chroma_sampling_type, base_depth_type, enhancement_depth_type,
        temporal_enabled, upsample_type, tile_dimensions_type, custom_size,
    })
}

pub fn probe(stream: &[u8]) -> Result<Value, Refusal> {
    let units = annex_b_units(stream)?;
    let (mut seq, mut glob) = (None, None);
    let (mut idr, mut non_idr, mut pictures) = (0u64, 0u64, 0u64);
    for (i, u) in units.iter().enumerate() {
        let b = u.bytes;
        if b.len() < 3 {
            return Err(Refusal::Truncated);
        }
        // forbidden_zero_bit 0, forbidden_one_bit 1, every reserved_flag bit 1.
        if b[0] & 0xc1 != 0x41 || b[1] != 0xff {
            return Err(Refusal::ReservedValue);
        }
        let nut = (b[0] >> 1) & 0x1f;
        match nut {
            NUT_IDR => idr += 1,
            NUT_NON_IDR => non_idr += 1,
            _ => return Err(Refusal::ReservedValue),
        }
        if i == 0 && nut != NUT_IDR {
            return Err(Refusal::OutOfOrder);
        }
        let rbsp = unescape(&b[2..]);
        let Some((&stop, payload)) = rbsp.split_last() else { return Err(Refusal::Truncated) };
        if stop != 0x80 {
            return Err(Refusal::ReservedValue);
        }
        for (kind, body) in blocks(payload)? {
            match kind {
                BLOCK_SEQUENCE_CONFIG if seq.is_none() => seq = Some(sequence_config(body)?),
                BLOCK_GLOBAL_CONFIG if glob.is_none() => glob = Some(global_config(body)?),
                BLOCK_PICTURE_CONFIG => {
                    if seq.is_none() || glob.is_none() {
                        return Err(Refusal::OutOfOrder);
                    }
                    pictures += 1;
                }
                _ => {}
            }
        }
    }
    let (Some(s), Some(g)) = (seq, glob) else { return Err(Refusal::MissingStructure) };
    Ok(json!({
        "nal_units": [idr, non_idr],
        "pictures": pictures,
        "profile_idc": s.profile_idc,
        "level_idc": s.level_idc,
        "sublevel_idc": s.sublevel_idc,
        "conformance_window": s.conformance_window,
        "resolution_type": g.resolution_type,
        "size": g.size(),
        "transform_type": g.transform_type,
        "chroma_sampling_type": g.chroma_sampling_type,
        "base_depth": 8 + 2 * g.base_depth_type,
        "enhancement_depth": 8 + 2 * g.enhancement_depth_type,
        "temporal_enabled": g.temporal_enabled,
        "upsample_type": g.upsample_type,
        "tile_dimensions_type": g.tile_dimensions_type,
    }))
}
