//! The JPEG XS probe (`wai.image.jxs`, `wai.video.jxs`; format-probes §3): the codestream's
//! markers from SOC to the first slice header, the picture header, the component table and
//! the trailing EOC, and the media-type parameters of RFC 9134 §7.1 they give.
//!
//! ISO/IEC 21122-1's text is not openly readable. The marker codes and segment layouts here
//! are the ones two independent open implementations agree on (format-probes §3.1); the
//! profile codes are ISO/IEC 21122-2:2022/Amd 1:2022 Table A.7 and the sublevel codes its
//! Table A.16, both read in the amendment's public preview.

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

pub const SOC: u16 = 0xff10;
pub const EOC: u16 = 0xff11;
pub const PIH: u16 = 0xff12;
pub const CDT: u16 = 0xff13;
pub const WGT: u16 = 0xff14;
pub const COM: u16 = 0xff15;
pub const NLT: u16 = 0xff16;
pub const CWD: u16 = 0xff17;
pub const CTS: u16 = 0xff18;
pub const CRG: u16 = 0xff19;
pub const SLH: u16 = 0xff20;
pub const CAP: u16 = 0xff50;

/// ISO/IEC 21122-2:2022/Amd 1:2022 Table A.7: `Ppih` and the profile's name as RFC 9134
/// §7.1 writes it (white space omitted).
pub const PROFILES: &[(u16, &str)] = &[
    (0x0000, "Unrestricted"),
    (0x1500, "Light422.10"),
    (0x1a00, "Light444.12"),
    (0x2500, "Light-Subline422.10"),
    (0x3240, "Main420.12"),
    (0x3540, "Main422.10"),
    (0x3a40, "Main444.12"),
    (0x3e40, "Main4444.12"),
    (0x4240, "High420.12"),
    (0x4a40, "High444.12"),
    (0x4e40, "High4444.12"),
    (0x6ec0, "MLS.12"),
    (0x9300, "LightBayer"),
    (0xb340, "MainBayer"),
    (0xc340, "HighBayer"),
];

/// The level, `Plev`'s high byte.
pub const LEVELS: &[(u8, &str)] = &[
    (0x00, "Unrestricted"),
    (0x04, "1k-1"),
    (0x10, "2k-1"),
    (0x20, "4k-1"),
    (0x24, "4k-2"),
    (0x28, "4k-3"),
    (0x2c, "5k-1"),
    (0x30, "8k-1"),
    (0x34, "8k-2"),
    (0x38, "8k-3"),
    (0x40, "10k-1"),
];

/// The sublevel, `Plev`'s low byte (Amd 1 Table A.16).
pub const SUBLEVELS: &[(u8, &str)] = &[
    (0x00, "Unrestricted"),
    (0x80, "Full"),
    (0x10, "Sublev12bpp"),
    (0x0c, "Sublev9bpp"),
    (0x08, "Sublev6bpp"),
    (0x06, "Sublev4bpp"),
    (0x04, "Sublev3bpp"),
    (0x03, "Sublev2bpp"),
];

fn be16(b: &[u8], at: usize) -> Result<u16, Refusal> {
    b.get(at..at + 2).map(|s| u16::from_be_bytes([s[0], s[1]])).ok_or(Refusal::Truncated)
}

/// One component of the component table.
#[derive(Debug, Clone, Copy)]
pub struct Component {
    pub depth: u8,
    pub sx: u8,
    pub sy: u8,
}

/// What the probe reads from one codestream.
#[derive(Debug, Clone)]
pub struct Codestream {
    pub ppih: u16,
    pub plev: u16,
    pub lcod: u32,
    pub width: u16,
    pub height: u16,
    pub cw: u16,
    pub hsl: u16,
    pub nc: u8,
    pub bw: u8,
    pub fq: u8,
    pub cpih: u8,
    pub nlx: u8,
    pub nly: u8,
    pub lh: bool,
    pub rl: bool,
    pub qpih: u8,
    pub cap_bits: Vec<u32>,
    pub components: Vec<Component>,
}

/// Reads one codestream, which is the whole of `b`: SOC first, EOC last.
pub fn codestream(b: &[u8]) -> Result<Codestream, Refusal> {
    if b.len() < 2 || be16(b, 0)? != SOC {
        return Err(Refusal::NoSignature);
    }
    // CAP immediately after SOC.
    if be16(b, 2)? != CAP {
        return Err(Refusal::OutOfOrder);
    }
    let lcap = be16(b, 4)? as usize;
    if lcap < 2 {
        return Err(Refusal::LengthMismatch);
    }
    let cap = b.get(6..4 + lcap).ok_or(Refusal::Truncated)?;
    let mut cap_bits = Vec::new();
    for (i, byte) in cap.iter().enumerate() {
        for k in 0..8 {
            if byte & (0x80 >> k) != 0 {
                cap_bits.push((i * 8 + k) as u32);
            }
        }
    }
    let mut at = 4 + lcap;
    // PIH immediately after CAP.
    if be16(b, at)? != PIH {
        return Err(Refusal::OutOfOrder);
    }
    if be16(b, at + 2)? != 26 {
        return Err(Refusal::LengthMismatch);
    }
    let p = b.get(at + 4..at + 28).ok_or(Refusal::Truncated)?;
    let lcod = u32::from_be_bytes([p[0], p[1], p[2], p[3]]);
    let ppih = u16::from_be_bytes([p[4], p[5]]);
    let plev = u16::from_be_bytes([p[6], p[7]]);
    let width = u16::from_be_bytes([p[8], p[9]]);
    let height = u16::from_be_bytes([p[10], p[11]]);
    let cw = u16::from_be_bytes([p[12], p[13]]);
    let hsl = u16::from_be_bytes([p[14], p[15]]);
    let (nc, ng, ss, bw) = (p[16], p[17], p[18], p[19]);
    let (fq, br) = (p[20] >> 4, p[20] & 15);
    let (_fslc, _ppoc, cpih) = (p[21] >> 7, (p[21] >> 4) & 7, p[21] & 15);
    let (nlx, nly) = (p[22] >> 4, p[22] & 15);
    let (lh, rl, qpih) = (p[23] >> 7 == 1, (p[23] >> 6) & 1 == 1, (p[23] >> 4) & 3);
    if width == 0 || height == 0 || hsl == 0 || !(1..=8).contains(&nc) || ng != 4 || ss != 8 || br != 4 {
        return Err(Refusal::ReservedValue);
    }
    if ![0, 1, 3].contains(&cpih) || qpih > 1 || nlx == 0 {
        return Err(Refusal::ReservedValue);
    }
    at += 28;
    let (mut cdt, mut wgt, mut cts, mut crg) = (None, false, false, false);
    loop {
        let m = be16(b, at)?;
        if m == SLH {
            break;
        }
        let l = be16(b, at + 2)? as usize;
        if l < 2 {
            return Err(Refusal::LengthMismatch);
        }
        let body = b.get(at + 4..at + 2 + l).ok_or(Refusal::Truncated)?;
        match m {
            CDT => {
                if l != 2 + 2 * nc as usize {
                    return Err(Refusal::LengthMismatch);
                }
                let comps: Vec<Component> =
                    body.chunks(2).map(|c| Component { depth: c[0], sx: c[1] >> 4, sy: c[1] & 15 }).collect();
                if comps.iter().any(|c| c.depth == 0 || c.sx == 0 || c.sy == 0) {
                    return Err(Refusal::ReservedValue);
                }
                cdt = Some(comps);
            }
            WGT => wgt = true,
            CTS => cts = true,
            CRG => crg = true,
            COM | NLT | CWD => {}
            SOC | EOC | PIH | CAP => return Err(Refusal::OutOfOrder),
            _ => return Err(Refusal::ReservedValue),
        }
        at += 2 + l;
    }
    let Some(components) = cdt else { return Err(Refusal::MissingStructure) };
    if !wgt || (cpih == 3 && !(cts && crg)) {
        return Err(Refusal::MissingStructure);
    }
    if b.len() < at + 4 || be16(b, b.len() - 2)? != EOC {
        return Err(Refusal::MissingStructure);
    }
    if lcod != 0 && lcod as usize != b.len() {
        return Err(Refusal::LengthMismatch);
    }
    Ok(Codestream { ppih, plev, lcod, width, height, cw, hsl, nc, bw, fq, cpih, nlx, nly, lh, rl, qpih, cap_bits, components })
}

fn name<T: PartialEq + Copy>(table: &[(T, &'static str)], v: T) -> Option<&'static str> {
    table.iter().find(|(k, _)| *k == v).map(|(_, n)| *n)
}

/// The chroma structure the component table states: `4:4:4`, `4:2:2` or `4:2:0` for three
/// or four components whose first is full resolution and whose next two share a factor;
/// otherwise `null`.
fn structure(c: &[Component]) -> Option<&'static str> {
    if c.len() < 3 || c[0].sx != 1 || c[0].sy != 1 || c[1].sx != c[2].sx || c[1].sy != c[2].sy {
        return None;
    }
    match (c[1].sx, c[1].sy) {
        (1, 1) => Some("4:4:4"),
        (2, 1) => Some("4:2:2"),
        (2, 2) => Some("4:2:0"),
        _ => None,
    }
}

/// The parameters of one codestream, with the RFC 9134 §7.1 media-type parameters it
/// determines (format-probes §3.4).
pub fn params(c: &Codestream) -> Value {
    let profile = name(PROFILES, c.ppih);
    let level = name(LEVELS, (c.plev >> 8) as u8);
    let sublevel = name(SUBLEVELS, (c.plev & 0xff) as u8);
    let depth = c.components.iter().map(|x| x.depth).max().unwrap_or(0);
    let st = structure(&c.components);
    // RGB samples only where a colour transform is signalled over three full components.
    let sampling = if c.cpih != 0 && st == Some("4:4:4") { Some("RGB") } else { None };
    json!({
        "profile": format!("0x{:04x}", c.ppih),
        "level": format!("0x{:04x}", c.plev),
        "width": c.width,
        "height": c.height,
        "components": c.components.iter().map(|x| json!([x.depth, x.sx, x.sy])).collect::<Vec<_>>(),
        "structure": st,
        "colour_transform": c.cpih,
        "decompositions": [c.nlx, c.nly],
        "precinct_width": c.cw,
        "slice_height": c.hsl,
        "wavelet_precision": c.bw,
        "fractional_bits": c.fq,
        "quantizer": c.qpih,
        "long_header": c.lh,
        "raw_mode": c.rl,
        "cbr_bytes": c.lcod,
        "cap_bits": c.cap_bits,
        "rtp": {
            "profile": profile,
            "level": level,
            "sublevel": sublevel,
            "depth": depth,
            "width": c.width,
            "height": c.height,
            "sampling": sampling,
        },
    })
}

pub fn probe_image(b: &[u8]) -> Result<Value, Refusal> {
    Ok(params(&codestream(b)?))
}

/// `wai.video.jxs`: five `u32` little-endian values `n`, `h`, `w`, `fps_num`, `fps_den`,
/// then `n` codestreams each preceded by its `u32` little-endian length (format-probes
/// §3.3).
pub fn probe_video(b: &[u8]) -> Result<Value, Refusal> {
    let word = |i: usize| -> Result<u32, Refusal> {
        b.get(4 * i..4 * i + 4).map(|s| u32::from_le_bytes([s[0], s[1], s[2], s[3]])).ok_or(Refusal::Truncated)
    };
    let (n, h, w, num, den) = (word(0)?, word(1)?, word(2)?, word(3)?, word(4)?);
    if n == 0 || h == 0 || w == 0 || num == 0 || den == 0 {
        return Err(Refusal::FramingInvalid);
    }
    let mut at = 20usize;
    let mut first = None;
    let mut lcods = Vec::with_capacity(n as usize);
    for _ in 0..n {
        let len = b.get(at..at + 4).map(|s| u32::from_le_bytes([s[0], s[1], s[2], s[3]])).ok_or(Refusal::Truncated)? as usize;
        let frame = b.get(at + 4..at + 4 + len).ok_or(Refusal::Truncated)?;
        let c = codestream(frame)?;
        if u32::from(c.width) > w || u32::from(c.height) > h {
            return Err(Refusal::FramingInvalid);
        }
        if first.is_none() {
            first = Some(params(&c));
        }
        lcods.push(c.lcod);
        at += 4 + len;
    }
    if at != b.len() {
        return Err(Refusal::TrailingBytes);
    }
    // The bytes every frame takes, where every frame states the same non-zero Lcod: the
    // constant number of bytes a frame SMPTE ST 2110-22 §4 asks of a stream (§3.5).
    let constant = lcods.iter().all(|&l| l != 0 && l == lcods[0]).then_some(lcods[0]);
    Ok(json!({ "frames": n, "size": [w, h], "fps": [num, den], "constant_bytes": constant, "first": first }))
}
