//! The AVS3 probes (`wai.video.avs3`, `wai.audio.avs3`; format-probes §8 and §9). T/AI
//! 109.2 and T/AI 109.3 are not openly readable. The video sequence header and the
//! pseudo-start-code prevention are read as an open decoder reads them, within the values
//! ETSI TS 101 154 V2.10.1 §5.16.2 constrains; the audio frame header as an open
//! demultiplexer reads it, within the values that document's §6.10.2 constrains.

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

pub const SEQUENCE_HEADER: u8 = 0xb0;
pub const SEQUENCE_END: u8 = 0xb1;
pub const INTRA_PICTURE: u8 = 0xb3;
pub const INTER_PICTURE: u8 = 0xb6;

/// The units of an AVS3 video stream: the byte after each start code prefix `0x000001`
/// and the bytes up to the next prefix. The stream must begin, after any zero bytes, with
/// a prefix.
pub fn units(b: &[u8]) -> Result<Vec<(u8, &[u8])>, Refusal> {
    let mut i = 0;
    while i < b.len() && b[i] == 0 {
        i += 1;
    }
    if i < 2 || i >= b.len() || b[i] != 1 {
        return Err(Refusal::NoSignature);
    }
    let mut starts = Vec::new();
    let mut j = 0;
    while j + 3 < b.len() {
        if b[j] == 0 && b[j + 1] == 0 && b[j + 2] == 1 {
            starts.push(j + 3);
            j += 3;
        } else {
            j += 1;
        }
    }
    let mut out = Vec::new();
    for (k, &s) in starts.iter().enumerate() {
        let end = if k + 1 < starts.len() { starts[k + 1] - 3 } else { b.len() };
        // A prefix directly after a prefix: that unit's start code is the next prefix's first
        // byte, and it holds nothing.
        out.push((b[s], b.get(s + 1..end).unwrap_or(&[])));
    }
    Ok(out)
}

/// Removes the pseudo-start-code prevention bits: in each byte `0x02` that follows two
/// zero bytes of the unit, its two low bits `10` (format-probes §8.2).
pub fn demulate(b: &[u8]) -> Vec<u8> {
    let mut bits: Vec<u8> = Vec::with_capacity(b.len() * 8);
    for (i, &x) in b.iter().enumerate() {
        let n = if i >= 2 && b[i - 2] == 0 && b[i - 1] == 0 && x == 2 { 6 } else { 8 };
        for k in 0..n {
            bits.push((x >> (7 - k)) & 1);
        }
    }
    bits.chunks(8).map(|c| c.iter().enumerate().fold(0u8, |a, (k, &bit)| a | (bit << (7 - k)))).collect()
}

fn marker(r: &mut BitReader) -> Result<(), Refusal> {
    if r.u(1)? != 1 {
        return Err(Refusal::ReservedValue);
    }
    Ok(())
}

pub fn sequence_header(rbsp: &[u8]) -> Result<Value, Refusal> {
    let mut r = BitReader::new(rbsp);
    let profile_id = r.u(8)?;
    let level_id = r.u(8)?;
    let progressive = r.flag()?;
    let field_coded = r.flag()?;
    if r.flag()? {
        // A library stream: what follows is not read by this probe.
        return Err(Refusal::OutsideProbe);
    }
    if r.flag()? {
        // Library pictures enabled: a further flag follows whose place the open texts do
        // not establish.
        return Err(Refusal::OutsideProbe);
    }
    marker(&mut r)?;
    let width = r.u(14)?;
    marker(&mut r)?;
    let height = r.u(14)?;
    let chroma_format = r.u(2)?;
    let sample_precision = r.u(3)?;
    let encoding_precision = if profile_id == 0x22 || profile_id == 0x32 { Some(r.u(3)?) } else { None };
    marker(&mut r)?;
    let aspect_ratio = r.u(4)?;
    let frame_rate_code = r.u(4)?;
    marker(&mut r)?;
    let bit_rate_lower = r.u(18)?;
    marker(&mut r)?;
    let bit_rate_upper = r.u(12)?;
    let low_delay = r.flag()?;
    let temporal_id = r.flag()?;
    marker(&mut r)?;
    let bbv = r.u(18)?;
    marker(&mut r)?;
    let max_dpb_size = r.u(4)? + 1;
    if width == 0 || height == 0 || chroma_format == 0 || sample_precision == 0 || frame_rate_code == 0 || frame_rate_code > 13 {
        return Err(Refusal::ReservedValue);
    }
    Ok(json!({
        "profile_id": profile_id,
        "level_id": level_id,
        "progressive": progressive,
        "field_coded": field_coded,
        "size": [width, height],
        "chroma_format": chroma_format,
        "sample_precision": sample_precision,
        "encoding_precision": encoding_precision,
        "aspect_ratio": aspect_ratio,
        "frame_rate_code": frame_rate_code,
        "bit_rate": (bit_rate_upper << 18) | bit_rate_lower,
        "low_delay": low_delay,
        "temporal_id_enabled": temporal_id,
        "bbv_buffer_size": bbv,
        "max_dpb_size": max_dpb_size,
    }))
}

pub fn probe_video(b: &[u8]) -> Result<Value, Refusal> {
    let us = units(b)?;
    let Some(&(first, _)) = us.first() else { return Err(Refusal::NoSignature) };
    let Some(seq) = us.iter().find(|(c, _)| *c == SEQUENCE_HEADER) else { return Err(Refusal::MissingStructure) };
    if first != SEQUENCE_HEADER {
        return Err(Refusal::OutOfOrder);
    }
    let mut out = sequence_header(&demulate(seq.1))?;
    let intra = us.iter().filter(|(c, _)| *c == INTRA_PICTURE).count();
    let inter = us.iter().filter(|(c, _)| *c == INTER_PICTURE).count();
    if intra == 0 {
        return Err(Refusal::MissingStructure);
    }
    out["pictures"] = json!([intra, inter]);
    out["sequence_end"] = json!(us.iter().any(|(c, _)| *c == SEQUENCE_END));
    Ok(out)
}

/// The sampling rates of `sampling_frequency_index` 0 to 8.
pub const AUDIO_RATES: [u32; 9] = [192000, 96000, 48000, 44100, 32000, 24000, 22050, 16000, 8000];

/// `wai.audio.avs3`: the first frame header of the stream (format-probes §9).
pub fn probe_audio(b: &[u8]) -> Result<Value, Refusal> {
    let mut r = BitReader::new(b);
    if r.u(12).map_err(|_| Refusal::NoSignature)? != 0xfff {
        return Err(Refusal::NoSignature);
    }
    // The lossless codec (1) and frames carrying ancillary data are outside what the open
    // reading covers; any other codec id is not one it defines.
    let codec_id = r.u(4)?;
    if codec_id == 1 {
        return Err(Refusal::OutsideProbe);
    }
    if codec_id != 2 {
        return Err(Refusal::ReservedValue);
    }
    if r.flag()? {
        return Err(Refusal::OutsideProbe);
    }
    let nn_type = r.u(3)?;
    if nn_type > 1 {
        return Err(Refusal::ReservedValue);
    }
    let coding_profile = r.u(3)?;
    if coding_profile > 2 {
        return Err(Refusal::ReservedValue);
    }
    let sfi = r.u(4)?;
    if sfi > 8 {
        return Err(Refusal::ReservedValue);
    }
    r.skip(8)?;
    let mut out = json!({
        "codec_id": codec_id,
        "nn_type": nn_type,
        "coding_profile": coding_profile,
        "sampling_rate": AUDIO_RATES[sfi as usize],
    });
    match coding_profile {
        0 => out["channel_number_index"] = json!(r.u(7)?),
        1 => {
            let soundbed = r.u(2)?;
            out["soundbed_type"] = json!(soundbed);
            match soundbed {
                0 => {
                    out["objects"] = json!(r.u(7)? + 1);
                    out["bitrate_index_per_object"] = json!(r.u(4)?);
                }
                1 => {
                    out["channel_number_index"] = json!(r.u(7)?);
                    out["bitrate_index"] = json!(r.u(4)?);
                    out["objects"] = json!(r.u(7)? + 1);
                    out["bitrate_index_per_object"] = json!(r.u(4)?);
                }
                _ => return Err(Refusal::ReservedValue),
            }
        }
        _ => {
            let order = r.u(4)? + 1;
            if !(1..=3).contains(&order) {
                return Err(Refusal::ReservedValue);
            }
            out["ambisonic_order"] = json!(order);
        }
    }
    let resolution = r.u(2)?;
    if resolution > 2 {
        return Err(Refusal::ReservedValue);
    }
    out["resolution_bits"] = json!(8 * (resolution + 1));
    if coding_profile != 1 {
        out["bitrate_index"] = json!(r.u(4)?);
    }
    r.skip(8)?;
    Ok(out)
}
