//! The queries against real AV1 floors and constructed sequence headers and SPSs, the
//! query vectors of `wai/platform-decode-conformance/queries.json`, and the answer a page
//! enters in browser-sink's report. The order the answers feed is `energy_order`'s and
//! browser-sink §4.7's, tested there.

use super::*;
use serde_json::Value;

const QUERIES: &str = include_str!("../../../platform-decode-conformance/queries.json");

// ---------------------------------------------------------------- a bit writer

/// Bits, most significant first, as the AV1 and H.265 headers are written.
#[derive(Default)]
struct W {
    bits: Vec<bool>,
}

impl W {
    fn f(&mut self, n: u32, v: u64) -> &mut Self {
        for k in (0..n).rev() {
            self.bits.push((v >> k) & 1 == 1);
        }
        self
    }
    fn ue(&mut self, v: u32) -> &mut Self {
        let x = u64::from(v) + 1;
        let len = 64 - x.leading_zeros();
        self.f(len - 1, 0).f(len, x)
    }
    /// The bytes, padded with `trailing_bits` (a 1, then zeros).
    fn bytes(&self) -> Vec<u8> {
        let mut b = self.bits.clone();
        b.push(true);
        while !b.len().is_multiple_of(8) {
            b.push(false);
        }
        b.chunks(8).map(|c| c.iter().fold(0u8, |a, &x| (a << 1) | u8::from(x))).collect()
    }
}

/// An RBSP with emulation-prevention bytes inserted (H.265 §7.4.2).
fn escape(rbsp: &[u8]) -> Vec<u8> {
    let mut out = Vec::new();
    let mut zeros = 0;
    for &b in rbsp {
        if zeros >= 2 && b <= 3 {
            out.push(3);
            zeros = 0;
        }
        zeros = if b == 0 { zeros + 1 } else { 0 };
        out.push(b);
    }
    out
}

// ---------------------------------------------------------------- AV1

struct Seq {
    profile: u64,
    level: u64,
    tier: u64,
    high: bool,
    twelve: bool,
    mono: bool,
    color: Option<(u64, u64, u64)>,
    range: bool,
    ss: Option<(bool, bool)>,
    csp: u64,
}

impl Seq {
    fn new(profile: u64, level: u64) -> Seq {
        Seq { profile, level, tier: 0, high: false, twelve: false, mono: false, color: None, range: false, ss: None, csp: 0 }
    }

    /// A sequence header OBU's payload (AV1 §5.5): one operating point, frames of
    /// 640×360, order hints, screen content tools chosen per frame.
    fn payload(&self) -> Vec<u8> {
        let mut w = W::default();
        w.f(3, self.profile).f(1, 0).f(1, 0); // still_picture, reduced_still_picture_header
        w.f(1, 0).f(1, 0).f(5, 0); // timing_info, initial_display_delay, one operating point
        w.f(12, 0).f(5, self.level);
        if self.level > 7 {
            w.f(1, self.tier);
        }
        w.f(4, 15).f(4, 15).f(16, 639).f(16, 359);
        w.f(1, 0); // frame_id_numbers_present_flag
        w.f(3, 0b011); // 128x128 superblock, filter intra, intra edge
        w.f(4, 0b1111).f(1, 1).f(2, 0b11); // interintra..dual filter; order hint; jnt, ref mvs
        w.f(1, 1).f(1, 1); // seq_choose_screen_content_tools, seq_choose_integer_mv
        w.f(3, 6); // order_hint_bits_minus_1
        w.f(3, 0b011); // superres, cdef, restoration
        w.f(1, u64::from(self.high));
        if self.profile == 2 && self.high {
            w.f(1, u64::from(self.twelve));
        }
        if self.profile != 1 {
            w.f(1, u64::from(self.mono));
        }
        match self.color {
            Some((cp, tc, mc)) => {
                w.f(1, 1).f(8, cp).f(8, tc).f(8, mc);
            }
            None => {
                w.f(1, 0);
            }
        }
        let srgb = self.color == Some((1, 13, 0));
        if self.mono {
            w.f(1, u64::from(self.range));
        } else if !srgb {
            w.f(1, u64::from(self.range));
            if self.profile == 2 && self.twelve {
                let (x, y) = self.ss.unwrap_or((true, true));
                w.f(1, u64::from(x));
                if x {
                    w.f(1, u64::from(y));
                }
            }
            let (x, y) = match self.profile {
                0 => (true, true),
                1 => (false, false),
                _ => self.ss.unwrap_or((true, false)),
            };
            if x && y {
                w.f(2, self.csp);
            }
        }
        if !self.mono {
            w.f(1, 0); // separate_uv_delta_q
        }
        w.f(1, 0); // film_grain_params_present
        w.bytes()
    }
}

/// A WAI AV1 payload (SPEC §5): `n`, `h`, `w`, `fps_num`, `fps_den`, then each temporal
/// unit with its length.
fn framing(units: &[Vec<u8>], h: u32, w: u32, fps: (u32, u32)) -> Vec<u8> {
    let mut p = Vec::new();
    for v in [units.len() as u32, h, w, fps.0, fps.1] {
        p.extend_from_slice(&v.to_le_bytes());
    }
    for u in units {
        p.extend_from_slice(&(u.len() as u32).to_le_bytes());
        p.extend_from_slice(u);
    }
    p
}

/// One temporal unit: a temporal delimiter, then a sequence header OBU around `seq`.
fn unit_with(seq: &[u8]) -> Vec<u8> {
    let mut u = vec![0x12, 0x00, 0x0a];
    u.push(seq.len() as u8);
    u.extend_from_slice(seq);
    u
}

#[test]
fn av1_codec_strings_follow_the_isobmff_binding() {
    let cs = |s: &Seq| Av1Sequence::parse(&s.payload()).unwrap().codec_string();
    // The binding's own examples (AV1-ISOBMFF §5).
    let mut hdr = Seq::new(0, 4);
    hdr.high = true;
    hdr.color = Some((9, 16, 9));
    hdr.csp = 2;
    assert_eq!(cs(&hdr), "av01.0.04M.10.0.112.09.16.09.0", "the binding's example: Main, level 3.0, 10-bit, BT.2100 PQ");
    assert_eq!(cs(&Seq::new(0, 1)), "av01.0.01M.08", "the binding's example: Main, level 2.1, 8-bit, the defaults omitted");
    let mut high_tier = Seq::new(0, 8);
    high_tier.tier = 1;
    assert_eq!(cs(&high_tier), "av01.0.08H.08");
    let mut full_range_444 = Seq::new(1, 5);
    full_range_444.range = true;
    assert_eq!(cs(&full_range_444), "av01.1.05M.08.0.000.01.01.01.1", "4:4:4 and full range are not the defaults");
    let mut srgb = Seq::new(1, 5);
    srgb.color = Some((1, 13, 0));
    assert_eq!(cs(&srgb), "av01.1.05M.08.0.000.01.13.00.1", "sRGB implies full range and no subsampling");
    let mut mono = Seq::new(0, 4);
    mono.mono = true;
    assert_eq!(cs(&mono), "av01.0.04M.08.1.110.01.01.01.0");
    let mut p2 = Seq::new(2, 10);
    p2.high = true;
    p2.twelve = true;
    p2.ss = Some((true, false));
    assert_eq!(cs(&p2), "av01.2.10M.12.0.100.01.01.01.0", "12-bit 4:2:2");
    let mut csp = Seq::new(0, 4);
    csp.csp = 1;
    assert_eq!(cs(&csp), "av01.0.04M.08.0.111.01.01.01.0", "a chroma sample position other than unknown is written");
}

#[test]
fn an_av1_query_reads_the_framing_and_the_first_sequence_header() {
    let mut s = Seq::new(0, 8);
    s.high = true;
    let seq = s.payload();
    let units = vec![unit_with(&seq), vec![0x12, 0x00, 0x32, 0x01, 0x00]];
    let p = framing(&units, 360, 640, (30000, 1001));
    let q = query(codecs::CAP_VIDEO_AV1, &p, None).unwrap();
    assert_eq!(q.codec, "av01.0.08M.10");
    assert_eq!((q.width, q.height, q.coded_width, q.coded_height, q.bit_depth), (640, 360, 640, 360, 10));
    assert_eq!(q.framerate, Some(Rate { num: 30000, den: 1001 }));
    let want = ((p.len() as u128) * 8 * 30000).div_ceil(2 * 1001) as u64;
    assert_eq!(q.bitrate, Some(want));
    let mc = q.decoding_configuration("media-source").unwrap();
    assert_eq!(mc["video"]["contentType"], "video/mp4; codecs=\"av01.0.08M.10\"");
    assert_eq!(mc["type"], "media-source");
    let wc = q.decoder_config(HardwareAcceleration::PreferHardware);
    assert_eq!(wc["hardwareAcceleration"], "prefer-hardware");
    assert_eq!(wc["codedWidth"], 640);
    // A payload with a zero frame rate has a query with no rate, and so no
    // decodingInfo() configuration.
    let q0 = query(codecs::CAP_VIDEO_AV1, &framing(&units, 360, 640, (0, 1)), None).unwrap();
    assert_eq!((q0.framerate, q0.bitrate), (None, None));
    assert!(q0.decoding_configuration("file").is_none());
    // Truncated, and a unit with no sequence header.
    assert!(matches!(query(codecs::CAP_VIDEO_AV1, &p[..30], None), Err(QueryError::Malformed(_))));
    let no_seq = framing(&[vec![0x12, 0x00]], 360, 640, (30, 1));
    assert!(matches!(query(codecs::CAP_VIDEO_AV1, &no_seq, None), Err(QueryError::Malformed(_))));
    assert_eq!(query(codecs::CAP_NEURAL_INT_SYNTH, &p, None), Err(QueryError::NotMapped(codecs::CAP_NEURAL_INT_SYNTH.into())));
}

/// The low-overhead OBU stream split into temporal units at each temporal delimiter.
fn temporal_units(obus: &[u8]) -> Vec<Vec<u8>> {
    let mut units: Vec<Vec<u8>> = Vec::new();
    let mut at = 0;
    while at < obus.len() {
        let header = obus[at];
        let hl = if header & 0x04 != 0 { 2 } else { 1 };
        let (size, sl) = leb128(&obus[at + hl..]).unwrap();
        let end = at + hl + sl + size as usize;
        if (header >> 3) & 0x0f == 2 || units.is_empty() {
            units.push(Vec::new());
        }
        units.last_mut().unwrap().extend_from_slice(&obus[at..end]);
        at = end;
    }
    units
}

/// Every AV1 floor the packaging corpus supplies, wrapped in SPEC §5's framing at the
/// pack's frame rate: its codec string is one the packager wrote into the pack's HLS
/// master (`CODECS`), and its size one of the master's `RESOLUTION`s. The packager reads
/// the sequence header with its own reader (`wai-pack` `av1.rs`).
#[test]
fn the_packaging_corpus_floors_give_their_catalogs_codec_strings() {
    let packs = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../pack-conformance/packs");
    let mut checked = 0;
    let mut names: Vec<_> = std::fs::read_dir(&packs).unwrap().filter_map(|e| e.ok()).map(|e| e.path()).collect();
    names.sort();
    for pack in names {
        let master = pack.join("package/hls/classic/master.m3u8");
        let (Ok(floors), Ok(m)) = (std::fs::read_dir(pack.join("floor")), std::fs::read_to_string(&master)) else { continue };
        let codecs_of: Vec<String> = m.split("CODECS=\"").skip(1).filter_map(|s| s.split('"').next()).flat_map(|s| s.split(',').map(str::to_string).collect::<Vec<_>>()).filter(|c| c.starts_with("av01")).collect();
        let sizes: Vec<(u32, u32)> = m
            .split("RESOLUTION=")
            .skip(1)
            .filter_map(|s| {
                let r: String = s.chars().take_while(|c| c.is_ascii_alphanumeric()).collect();
                let (w, h) = r.split_once('x')?;
                Some((w.parse().ok()?, h.parse().ok()?))
            })
            .collect();
        for f in floors.filter_map(|e| e.ok()).map(|e| e.path()).filter(|p| p.extension().is_some_and(|x| x == "obu")) {
            let obus = std::fs::read(&f).unwrap();
            let units = temporal_units(&obus);
            let seq = obus_sequence(&units[0]);
            let (w, h) = sizes.iter().copied().find(|&(w, h)| (w, h) == seq).unwrap_or(seq);
            let p = framing(&units, h, w, (25, 1));
            let q = query(codecs::CAP_VIDEO_AV1, &p, None).unwrap_or_else(|e| panic!("{}: {e}", f.display()));
            assert!(codecs_of.contains(&q.codec), "{}: {} is not among the master's {codecs_of:?}", f.display(), q.codec);
            assert!(sizes.contains(&(q.width, q.height)), "{}: {}x{} is not among the master's {sizes:?}", f.display(), q.width, q.height);
            checked += 1;
        }
    }
    assert!(checked >= 4, "only {checked} floors were checked");
}

/// The max frame size the first sequence header states.
fn obus_sequence(unit: &[u8]) -> (u32, u32) {
    let mut at = 0;
    loop {
        let header = unit[at];
        let hl = if header & 0x04 != 0 { 2 } else { 1 };
        let (size, sl) = leb128(&unit[at + hl..]).unwrap();
        let body = &unit[at + hl + sl..at + hl + sl + size as usize];
        if (header >> 3) & 0x0f == 1 {
            let mut r = Bits { b: body, at: 0 };
            let _profile = r.f(3).unwrap();
            r.f(1).unwrap();
            let reduced = r.flag().unwrap();
            assert!(!reduced);
            assert!(!r.flag().unwrap(), "no timing info in the corpus floors");
            assert!(!r.flag().unwrap());
            let cnt = r.f(5).unwrap() + 1;
            for _ in 0..cnt {
                r.f(12).unwrap();
                if r.f(5).unwrap() > 7 {
                    r.f(1).unwrap();
                }
            }
            let (wb, hb) = (r.f(4).unwrap() + 1, r.f(4).unwrap() + 1);
            return (r.f(wb).unwrap() + 1, r.f(hb).unwrap() + 1);
        }
        at += hl + sl + size as usize;
    }
}

// ---------------------------------------------------------------- HEVC

struct Sps {
    space: u64,
    tier: u64,
    idc: u64,
    compat: u32,
    constraints: [u8; 6],
    level: u64,
    sub_layers: Vec<(bool, bool)>,
    chroma: u32,
    width: u32,
    height: u32,
    window: Option<(u32, u32, u32, u32)>,
    depth: u32,
}

impl Sps {
    fn main(level: u64) -> Sps {
        Sps { space: 0, tier: 0, idc: 1, compat: 0b110, constraints: [0xB0, 0, 0, 0, 0, 0], level, sub_layers: Vec::new(), chroma: 1, width: 1920, height: 1088, window: Some((0, 0, 0, 4)), depth: 8 }
    }

    /// The NAL unit, header and emulation prevention included.
    fn nal(&self) -> Vec<u8> {
        let mut w = W::default();
        w.f(4, 0).f(3, self.sub_layers.len() as u64).f(1, 1);
        w.f(2, self.space).f(1, self.tier).f(5, self.idc);
        for j in 0..32 {
            w.f(1, u64::from((self.compat >> j) & 1));
        }
        for c in self.constraints {
            w.f(8, u64::from(c));
        }
        w.f(8, self.level);
        for &(p, l) in &self.sub_layers {
            w.f(1, u64::from(p)).f(1, u64::from(l));
        }
        if !self.sub_layers.is_empty() {
            for _ in self.sub_layers.len()..8 {
                w.f(2, 0);
            }
        }
        for &(p, l) in &self.sub_layers {
            if p {
                w.f(32, 0x1234_5678).f(32, 0x9abc_def0).f(24, 0x00ab_cdef);
            }
            if l {
                w.f(8, 90);
            }
        }
        w.ue(0).ue(self.chroma);
        if self.chroma == 3 {
            w.f(1, 0);
        }
        w.ue(self.width).ue(self.height);
        match self.window {
            Some((l, r, t, b)) => {
                w.f(1, 1).ue(l).ue(r).ue(t).ue(b);
            }
            None => {
                w.f(1, 0);
            }
        }
        w.ue(self.depth - 8).ue(self.depth - 8);
        let mut nal = vec![33 << 1, 1];
        nal.extend(escape(&w.bytes()));
        nal
    }
}

/// An Annex B stream: an SPS, then `pictures` slices, each the first of its picture.
fn hevc_stream(sps: &Sps, pictures: usize) -> Vec<u8> {
    let mut s = vec![0, 0, 0, 1];
    s.extend(sps.nal());
    for k in 0..pictures {
        s.extend_from_slice(&[0, 0, 1, 1 << 1, 1, 0x80 | (k as u8 & 0x7f), 0xAA, 0x55]);
    }
    s
}

#[test]
fn hevc_codec_strings_and_queries() {
    let s = Sps::main(93);
    let parsed = HevcSps::parse(&s.nal()).unwrap();
    assert_eq!(parsed.codec_string(), "hev1.1.6.L93.B0");
    assert_eq!((parsed.width, parsed.height, parsed.cropped_width, parsed.cropped_height, parsed.bit_depth), (1920, 1088, 1920, 1080, 8));
    // Main 10, high tier, level 5.1, a second constraint byte.
    let mut m10 = Sps::main(153);
    (m10.idc, m10.compat, m10.tier, m10.depth, m10.constraints) = (2, 0b100, 1, 10, [0x90, 0x00, 0x12, 0, 0, 0]);
    assert_eq!(HevcSps::parse(&m10.nal()).unwrap().codec_string(), "hev1.2.4.H153.90.0.12");
    // A profile space, and sub-layers with a profile and a level present.
    let mut spaced = Sps::main(120);
    spaced.space = 1;
    spaced.sub_layers = vec![(true, false), (false, true)];
    let p = HevcSps::parse(&spaced.nal()).unwrap();
    assert_eq!(p.codec_string(), "hev1.A1.6.L120.B0");
    assert_eq!((p.cropped_width, p.cropped_height), (1920, 1080), "the sub-layer fields are skipped exactly");
    // Every compatibility flag clear: the RBSP holds a run of zero bytes, which the NAL
    // carries with emulation-prevention bytes.
    let mut zeros = Sps::main(93);
    (zeros.compat, zeros.constraints, zeros.window) = (0, [0; 6], None);
    let nal = zeros.nal();
    assert!(nal.windows(3).any(|w| w == [0, 0, 3]), "the constructed NAL needs emulation prevention");
    let z = HevcSps::parse(&nal).unwrap();
    assert_eq!(z.codec_string(), "hev1.1.0.L93");
    assert_eq!((z.cropped_width, z.cropped_height), (1920, 1088));
    // The query: the cropped size, the coded size, the presentation rate the sink gives,
    // and the bit rate over the pictures the stream holds.
    let stream = hevc_stream(&s, 4);
    let q = query(codecs::CAP_VIDEO_HEVC, &stream, Some(Rate { num: 25, den: 1 })).unwrap();
    assert_eq!((q.codec.as_str(), q.width, q.height, q.coded_width, q.coded_height), ("hev1.1.6.L93.B0", 1920, 1080, 1920, 1088));
    assert_eq!(q.bitrate, Some(((stream.len() as u64) * 8 * 25).div_ceil(4)));
    let unrated = query(codecs::CAP_VIDEO_HEVC, &stream, None).unwrap();
    assert_eq!((unrated.framerate, unrated.bitrate), (None, None), "the payload states no frame rate");
    assert!(unrated.decoding_configuration("media-source").is_none());
    assert!(matches!(query(codecs::CAP_VIDEO_HEVC, &[0, 0, 1, 2, 1, 0x80], None), Err(QueryError::Malformed(_))), "no SPS");
}

// ---------------------------------------------------------------- the answers

/// browser-sink §3's answer from what the platform said: `supported` is WebCodecs' alone, a
/// Media Capabilities `supported: false` overrides nothing, and `smooth` and
/// `power_efficient` are read only once WebCodecs said supported.
#[test]
fn the_report_answer_takes_supported_from_webcodecs() {
    let a = report_answer(true, Some((true, true)));
    assert_eq!(a, serde_json::json!({ "supported": true, "smooth": true, "power_efficient": true, "probe": "not_run" }));
    // Media Capabilities' own `supported` is not an input at all: a `media-source` answer is
    // about the MSE pipeline, not the decoder WebCodecs found.
    assert_eq!(report_answer(true, Some((false, false)))["supported"], true);
    let no = report_answer(false, Some((true, true)));
    assert_eq!((&no["supported"], &no["smooth"], &no["power_efficient"]), (&Value::Bool(false), &Value::Null, &Value::Null));
    assert_eq!(report_answer(true, None)["power_efficient"], Value::Null, "nothing asked: not a power-efficient answer");
}

/// `queries.json`: two floor objects of the packaging corpus, and an SPS the platform's
/// encoder wrote, configured with `prefer-hardware`, whose level is not the one the encoder
/// was configured with.
#[test]
fn the_query_vectors() {
    let v: Value = serde_json::from_str(QUERIES).unwrap();
    let wai = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("..");
    let mut n = 0;
    for case in v["cases"].as_array().unwrap() {
        let id = case["id"].as_str().unwrap();
        let cap = case["capability"].as_str().unwrap();
        let want = &case["expect"];
        let got = if let Some(f) = case["payload_file"].as_str() {
            let env = crate::container::Wai::from_bytes(&std::fs::read(wai.join(f)).unwrap()).unwrap();
            assert_eq!(env.capability(), cap, "{id}");
            query_json(&query(cap, &env.payload, None).unwrap())
        } else {
            let hex = case["sps_hex"].as_str().unwrap();
            let nal: Vec<u8> = (0..hex.len()).step_by(2).map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap()).collect();
            let sps = HevcSps::parse(&nal).unwrap();
            serde_json::json!({ "codec": sps.codec_string(), "width": sps.cropped_width, "height": sps.cropped_height, "codedWidth": sps.width, "codedHeight": sps.height, "bitDepth": sps.bit_depth })
        };
        for (k, w) in want.as_object().unwrap() {
            assert_eq!(&got[k.as_str()], w, "{id}: {k}");
        }
        n += 1;
    }
    assert_eq!(n, 3);
}
