//! IVF, YUV4MPEG2, FLV, AV1 OBU streams and Annex B streams (media-identify §5.14 to §5.17, §5.20).
//! Frames are counted by their headers; no sample is read.

use super::{be24, be32, clean, get, le16, le32, Candidate, Source, Strength, MAX_STEPS};
use crate::codecs::*;
use crate::format_probe::bits::BitReader;
use serde_json::{json, Value};
use std::io;

pub fn ivf(src: &dyn Source, head: &[u8]) -> io::Result<Option<Candidate>> {
    if head.get(..4) != Some(b"DKIF") {
        return Ok(None);
    }
    let size = src.size();
    let mut c = Candidate::new("ivf", "container");
    c.see(0, 4, "dkif");
    let (Some(hdr), Some(fourcc)) = (le16(head, 6), head.get(8..12)) else {
        c.walked_to = Some(0);
        c.stop(0, "overrun");
        return Ok(Some(c));
    };
    if hdr < 32 || le16(head, 4) != Some(0) {
        c.walked_to = Some(0);
        c.stop(0, "invalid");
        return Ok(Some(c));
    }
    let codec = match fourcc {
        b"AV01" => Some("av1"),
        b"VP90" => Some("vp9"),
        b"VP80" => Some("vp8"),
        _ => None,
    };
    c.track(super::Track {
        media: "video",
        stated: clean(fourcc),
        codec,
        width: le16(head, 12),
        height: le16(head, 14),
        ..super::Track::default()
    });
    let (mut pos, mut frames) = (hdr, 0u64);
    loop {
        if pos == size {
            break;
        }
        if frames == MAX_STEPS {
            c.stop(pos, "limit");
            break;
        }
        let h = get(src, pos, 12)?;
        let Some(n) = le32(&h, 0).filter(|_| h.len() == 12) else {
            c.stop(pos, "overrun");
            break;
        };
        if pos + 12 + n > size {
            c.stop(pos, "overrun");
            break;
        }
        pos += 12 + n;
        frames += 1;
    }
    c.walked_to = Some(pos);
    c.fact("frames_read", frames);
    if frames >= 1 {
        c.strength = Strength::Structure;
    }
    Ok(Some(c))
}

/// The bytes of a frame of `w` by `h` under YUV4MPEG2's colour tag, where the tag is known.
fn y4m_frame(tag: &[u8], w: u64, h: u64) -> Option<u64> {
    let (cw, ch) = (w / 2 + w % 2, h / 2 + h % 2);
    let tag = std::str::from_utf8(tag).ok()?;
    let (base, depth) = match tag.split_once('p') {
        Some((b, d)) if !d.is_empty() && d.bytes().all(|c| c.is_ascii_digit()) => (b, d.parse::<u64>().ok()?),
        _ => (tag, 8),
    };
    let samples = match base {
        "420" | "420jpeg" | "420mpeg2" | "420paldv" => w.checked_mul(h)?.checked_add(cw.checked_mul(ch)?.checked_mul(2)?)?,
        "422" => w.checked_mul(h)?.checked_add(cw.checked_mul(h)?.checked_mul(2)?)?,
        "444" => w.checked_mul(h)?.checked_mul(3)?,
        "mono" => w.checked_mul(h)?,
        _ => return None,
    };
    samples.checked_mul(if depth > 8 { 2 } else { 1 })
}

pub fn y4m(src: &dyn Source, head: &[u8]) -> io::Result<Option<Candidate>> {
    if head.get(..10) != Some(b"YUV4MPEG2 ") {
        return Ok(None);
    }
    let size = src.size();
    let mut c = Candidate::new("y4m", "elementary");
    c.see(0, 10, "signature");
    let Some(nl) = head.iter().take(256).position(|&b| b == b'\n') else {
        c.walked_to = Some(0);
        c.stop(0, if (head.len() as u64) == size && head.len() < 256 { "overrun" } else { "invalid" });
        return Ok(Some(c));
    };
    let (mut w, mut h, mut colour) = (None, None, b"420".as_slice());
    for tok in head[10..nl].split(|&b| b == b' ') {
        let parse = |t: &[u8]| std::str::from_utf8(t).ok().and_then(|s| s.parse::<u64>().ok());
        match tok.split_first() {
            Some((b'W', v)) => w = parse(v),
            Some((b'H', v)) => h = parse(v),
            Some((b'C', v)) => colour = v,
            Some((b'F', v)) => {
                c.fact("frame_rate", clean(v));
            }
            _ => {}
        }
    }
    c.fact("width", w);
    c.fact("height", h);
    c.fact("colour", clean(colour));
    let mut pos = nl as u64 + 1;
    let Some(frame) = w.zip(h).and_then(|(w, h)| y4m_frame(colour, w, h)) else {
        c.walked_to = Some(pos);
        c.stop(pos, "invalid");
        return Ok(Some(c));
    };
    let mut frames = 0u64;
    loop {
        if pos == size {
            break;
        }
        if frames == MAX_STEPS {
            c.stop(pos, "limit");
            break;
        }
        let l = get(src, pos, 256)?;
        if !l.starts_with(b"FRAME") {
            c.stop(pos, if l.len() < 5 { "overrun" } else { "invalid" });
            break;
        }
        let Some(e) = l.iter().position(|&b| b == b'\n') else {
            c.stop(pos, "invalid");
            break;
        };
        let Some(next) = (pos + e as u64 + 1).checked_add(frame).filter(|&n| n <= size) else {
            c.stop(pos, "overrun");
            break;
        };
        pos = next;
        frames += 1;
    }
    c.walked_to = Some(pos);
    c.fact("frames_read", frames);
    if frames >= 1 {
        c.strength = Strength::Structure;
    }
    Ok(Some(c))
}

/// What a tag names its codec, as it states it, and the codec where this extension names one.
type Named = Option<(String, Option<&'static str>)>;

pub fn flv(src: &dyn Source, head: &[u8]) -> io::Result<Option<Candidate>> {
    if head.get(..4) != Some(b"FLV\x01") {
        return Ok(None);
    }
    let size = src.size();
    let mut c = Candidate::new("flv", "container");
    c.see(0, 4, "signature");
    let Some(data) = be32(head, 5).filter(|&d| d >= 9) else {
        c.walked_to = Some(0);
        c.stop(0, if head.len() < 9 { "overrun" } else { "invalid" });
        return Ok(Some(c));
    };
    let mut pos = data + 4;
    let (mut tags, mut audio, mut video): (u64, Named, Named) = (0, None, None);
    loop {
        if pos == size {
            break;
        }
        if tags == MAX_STEPS {
            c.stop(pos, "limit");
            break;
        }
        let h = get(src, pos, 16)?;
        if h.len() < 11 {
            c.stop(pos, "overrun");
            break;
        }
        let ty = h[0] & 0x1f;
        if !matches!(ty, 8 | 9 | 18) {
            c.stop(pos, "invalid");
            break;
        }
        let n = be24(&h, 1).unwrap_or(0);
        if let (8, None, Some(&b)) = (ty, &audio, h.get(11).filter(|_| n > 0)) {
            let format = b >> 4;
            let codec = match format {
                2 => Some("mp3"),
                10 => Some("aac"),
                _ => None,
            };
            audio = Some((format!("sound_format:{format}"), codec));
        }
        if let (9, None, Some(&b)) = (ty, &video, h.get(11).filter(|_| n > 0)) {
            // The enhanced form names its codec by four bytes after the first; the plain one by a number.
            video = if b & 0x80 != 0 {
                h.get(12..16).map(|f| {
                    let codec = match f {
                        b"avc1" => Some("avc"),
                        b"hvc1" => Some("hevc"),
                        b"av01" => Some("av1"),
                        b"vp09" => Some("vp9"),
                        _ => None,
                    };
                    (clean(f), codec)
                })
            } else {
                let id = b & 0x0f;
                Some((format!("codec_id:{id}"), (id == 7).then_some("avc")))
            };
        }
        let prev = get(src, pos + 11 + n, 4)?;
        let Some(p) = be32(&prev, 0) else {
            c.stop(pos, "overrun");
            break;
        };
        if p != 11 + n {
            c.stop(pos, "invalid");
            break;
        }
        pos += 15 + n;
        tags += 1;
    }
    c.walked_to = Some(pos);
    c.fact("tags_read", tags);
    if let Some((stated, codec)) = video {
        c.track(super::Track { media: "video", stated, codec, ..super::Track::default() });
    }
    if let Some((stated, codec)) = audio {
        c.track(super::Track { media: "audio", stated, codec, ..super::Track::default() });
    }
    if tags >= 1 {
        c.strength = Strength::Structure;
    }
    Ok(Some(c))
}

/// An AV1 sequence header's profile, level and largest frame (AV1 §5.5).
fn av1_sequence(b: &[u8]) -> Option<(u64, u64, u64, u64)> {
    let mut r = BitReader::new(b);
    let mut u = |n: u32| r.u(n).ok();
    let profile = u(3)?;
    let _still = u(1)?;
    let reduced = u(1)?;
    let level = if reduced == 1 {
        u(5)?
    } else {
        let (mut decoder_model, mut buffer_delay) = (false, 0);
        if u(1)? == 1 {
            u(32)?;
            u(32)?;
            if u(1)? == 1 {
                // uvlc: leading zeros, a one, then as many bits.
                let mut zeros = 0;
                while u(1)? == 0 {
                    zeros += 1;
                    if zeros >= 32 {
                        return None;
                    }
                }
                u(zeros)?;
            }
            if u(1)? == 1 {
                decoder_model = true;
                buffer_delay = u(5)? + 1;
                u(32)?;
                u(5)?;
                u(5)?;
            }
        }
        let delay = u(1)? == 1;
        let count = u(5)? + 1;
        let mut first = None;
        for _ in 0..count {
            u(12)?;
            let l = u(5)?;
            first.get_or_insert(l);
            if l > 7 {
                u(1)?;
            }
            if decoder_model && u(1)? == 1 {
                u(buffer_delay as u32)?;
                u(buffer_delay as u32)?;
                u(1)?;
            }
            if delay && u(1)? == 1 {
                u(4)?;
            }
        }
        first?
    };
    let wbits = u(4)? as u32 + 1;
    let hbits = u(4)? as u32 + 1;
    let w = u(wbits)? + 1;
    let h = u(hbits)? + 1;
    Some((profile, level, w, h))
}

pub fn av1(head: &[u8], size: u64) -> Option<Candidate> {
    let obu = |b: &[u8]| -> Option<(u8, usize, u64)> {
        let h = *b.first()?;
        if h & 0x81 != 0 || h & 0x02 == 0 {
            return None;
        }
        let ty = (h >> 3) & 15;
        if !matches!(ty, 1..=8 | 15) {
            return None;
        }
        let mut i = 1 + ((h >> 2) & 1) as usize;
        let (mut n, mut shift) = (0u64, 0u32);
        loop {
            let c = *b.get(i)?;
            i += 1;
            n |= ((c & 0x7f) as u64) << shift;
            if c & 0x80 == 0 {
                break;
            }
            shift += 7;
            if shift > 49 {
                return None;
            }
        }
        Some((ty, i, n))
    };
    let (first, ..) = obu(head)?;
    if first != 1 && first != 2 {
        return None;
    }
    let mut c = Candidate::new("av1-obu", "elementary");
    c.cap(CAP_VIDEO_AV1);
    let (mut pos, mut count, mut seen) = (0usize, 0u64, false);
    loop {
        if count == MAX_STEPS {
            c.stop(pos as u64, "limit");
            break;
        }
        let rest = &head[pos..];
        if rest.is_empty() {
            if (pos as u64) < size {
                c.stop(pos as u64, "limit");
            }
            break;
        }
        let Some((ty, hdr, n)) = obu(rest) else {
            c.stop(pos as u64, if rest.len() < 8 && (head.len() as u64) < size { "limit" } else { "invalid" });
            break;
        };
        if (pos as u64) + hdr as u64 + n > size {
            c.stop(pos as u64, "overrun");
            break;
        }
        if hdr as u64 + n > rest.len() as u64 {
            c.stop(pos as u64, "limit");
            break;
        }
        if ty == 1 && !seen {
            seen = true;
            c.see(pos as u64, hdr as u64 + n, "sequence_header");
            if let Some((p, l, w, h)) = av1_sequence(&rest[hdr..hdr + n as usize]) {
                c.fact("profile", p);
                c.fact("level", l);
                c.fact("width", w);
                c.fact("height", h);
            }
        }
        pos += hdr + n as usize;
        count += 1;
    }
    c.walked_to = Some(pos as u64);
    c.fact("units_read", count);
    if count == 0 {
        return None;
    }
    if seen && count >= 2 && c.facts.contains_key("width") {
        c.strength = Strength::Structure;
        Some(c)
    } else {
        // Two bytes that merely fit are not a claim of AV1.
        None
    }
}

/// The NAL units of an Annex B stream at the start of `head`, each with its offset, at most
/// [`MAX_STEPS`] of them; `None` unless it begins as §5.20 says.
fn nal_units(head: &[u8]) -> Option<Vec<(usize, &[u8])>> {
    let lead = head.iter().position(|&b| b != 0)?;
    if lead < 2 || head[lead] != 1 {
        return None;
    }
    let mut starts = Vec::new();
    let mut i = 0;
    while i + 2 < head.len() {
        if head[i] == 0 && head[i + 1] == 0 && head[i + 2] == 1 {
            starts.push(i + 3);
            if starts.len() as u64 > MAX_STEPS {
                break;
            }
            i += 3;
        } else {
            i += 1;
        }
    }
    let mut units = Vec::with_capacity(starts.len());
    for (k, &s) in starts.iter().enumerate() {
        if k as u64 == MAX_STEPS {
            break;
        }
        let mut end = if k + 1 < starts.len() { starts[k + 1] - 3 } else { head.len() };
        while end > s && head[end - 1] == 0 {
            end -= 1;
        }
        units.push((s, &head[s..end]));
    }
    Some(units)
}

/// An H.264 sequence parameter set (ITU-T H.264 §7.3.2.1.1): its profile, level and picture size.
/// The profiles §5.20 reads, in its order: the first three hold no chroma format.
pub(super) const AVC_PROFILES: [u64; 16] = [66, 77, 88, 100, 110, 122, 244, 44, 83, 86, 118, 128, 138, 139, 134, 135];

fn avc_sps(rbsp: &[u8]) -> Option<Value> {
    let high = &AVC_PROFILES[3..];
    let mut r = BitReader::new(rbsp);
    let profile = r.u(8).ok()?;
    if !AVC_PROFILES.contains(&profile) {
        return None;
    }
    let constraints = r.u(8).ok()?;
    let level = r.u(8).ok()?;
    r.ue().ok()?;
    let mut chroma = 1u64;
    let mut separate = false;
    if high.contains(&profile) {
        chroma = r.ue().ok()?;
        if chroma > 3 {
            return None;
        }
        if chroma == 3 {
            separate = r.flag().ok()?;
        }
        r.ue().ok()?;
        r.ue().ok()?;
        r.flag().ok()?;
        if r.flag().ok()? {
            for i in 0..if chroma != 3 { 8 } else { 12 } {
                if r.flag().ok()? {
                    let size = if i < 6 { 16 } else { 64 };
                    let (mut last, mut next) = (8i64, 8i64);
                    for _ in 0..size {
                        if next != 0 {
                            let code = r.ue().ok()? as i64;
                            let delta = if code % 2 == 1 { (code + 1) / 2 } else { -(code / 2) };
                            next = (last + delta + 256) % 256;
                        }
                        last = if next == 0 { last } else { next };
                    }
                }
            }
        }
    }
    r.ue().ok()?;
    match r.ue().ok()? {
        0 => {
            r.ue().ok()?;
        }
        1 => {
            r.flag().ok()?;
            r.ue().ok()?;
            r.ue().ok()?;
            let n = r.ue().ok()?;
            if n > 255 {
                return None;
            }
            for _ in 0..n {
                r.ue().ok()?;
            }
        }
        2 => {}
        _ => return None,
    }
    r.ue().ok()?;
    r.flag().ok()?;
    let w_mbs = r.ue().ok()?.checked_add(1)?;
    let h_units = r.ue().ok()?.checked_add(1)?;
    let frame_mbs_only = r.flag().ok()?;
    if !frame_mbs_only {
        r.flag().ok()?;
    }
    r.flag().ok()?;
    let (mut l, mut rt, mut t, mut b) = (0, 0, 0, 0);
    if r.flag().ok()? {
        l = r.ue().ok()?;
        rt = r.ue().ok()?;
        t = r.ue().ok()?;
        b = r.ue().ok()?;
    }
    let array = if separate { 0 } else { chroma };
    let (sub_w, sub_h) = match array {
        1 => (2, 2),
        2 => (2, 1),
        _ => (1, 1),
    };
    let (crop_x, crop_y) = (if array == 0 { 1 } else { sub_w }, (if array == 0 { 1 } else { sub_h }) * if frame_mbs_only { 1 } else { 2 });
    let width = w_mbs.checked_mul(16)?.checked_sub(crop_x * l.checked_add(rt)?)?;
    let height = h_units.checked_mul(if frame_mbs_only { 16 } else { 32 })?.checked_sub(crop_y * t.checked_add(b)?)?;
    Some(json!({
        "profile_idc": profile,
        "constraint_flags": constraints,
        "level_idc": level,
        "chroma_format_idc": chroma,
        "width": width,
        "height": height,
    }))
}

/// Whether the parameter-set types came in `order`, each before the first coded picture.
fn in_order(types: &[u8], order: &[u8], is_vcl: impl Fn(u8) -> bool) -> bool {
    let first_vcl = types.iter().position(|&t| is_vcl(t)).unwrap_or(types.len());
    let mut last = 0;
    for &want in order {
        let Some(at) = types.iter().position(|&t| t == want) else { return false };
        if at >= first_vcl || at < last {
            return false;
        }
        last = at;
    }
    true
}

fn elementary(id: &'static str, cap: &'static str, units: &[(usize, &[u8])]) -> Candidate {
    let mut c = Candidate::new(id, "elementary");
    c.cap(cap);
    c.strength = Strength::Structure;
    c.fact("units_read", units.len() as u64);
    c
}

/// The H.264, H.265 and H.266 readings the stream's NAL unit headers allow (§5.20), each kept only
/// where every header is legal and the parameter sets come in order before the first picture. A
/// stream two readings fit is reported as both.
pub fn annex_b(head: &[u8], size: u64) -> Vec<Candidate> {
    let Some(units) = nal_units(head) else { return Vec::new() };
    if units.is_empty() || units[0].1.is_empty() {
        return Vec::new();
    }
    let first = units[0].1[0];
    // An MPEG program stream or MPEG video sequence begins the same way (§6).
    if matches!(first, 0xba | 0xb3) {
        return Vec::new();
    }
    let mut out = Vec::new();
    let cut = (head.len() as u64) < size;
    let finish = |c: &mut Candidate| {
        let last = units.len();
        if cut || last as u64 == MAX_STEPS {
            c.stop(units[last - 1].0 as u64, "limit");
        }
    };

    // H.264: one byte, forbidden_zero_bit, nal_ref_idc nonzero on what a stream cannot lose.
    let avc: Option<Vec<u8>> = units
        .iter()
        .map(|(_, u)| {
            let b = *u.first()?;
            let ty = b & 31;
            let legal = b & 0x80 == 0 && (1..=21).contains(&ty) && !(matches!(ty, 5 | 7 | 8) && b >> 5 == 0);
            legal.then_some(ty)
        })
        .collect();
    if let Some(types) = avc
        && in_order(&types, &[7, 8], |t| (1..=5).contains(&t))
    {
        let (at, sps) = units.iter().find(|(_, u)| u[0] & 31 == 7).copied().unwrap_or((0, &[]));
        let rbsp = crate::format_probe::bits::unescape(sps);
        if let Some(v) = avc_sps(rbsp.get(1..).unwrap_or(&[])) {
            let mut c = elementary("avc-annexb", CAP_VIDEO_AVC, &units);
            c.see(at as u64, sps.len() as u64, "nal:sps");
            for (k, f) in v.as_object().into_iter().flatten() {
                c.facts.insert(k.clone(), f.clone());
            }
            finish(&mut c);
            out.push(c);
        }
    }

    // H.265: two bytes; a temporal identifier plus one that is not zero.
    let hevc: Option<Vec<u8>> = units
        .iter()
        .map(|(_, u)| {
            let (a, b) = (*u.first()?, *u.get(1)?);
            let ty = (a >> 1) & 63;
            let legal = a & 0x80 == 0 && b & 7 != 0 && (ty <= 9 || (16..=21).contains(&ty) || (32..=40).contains(&ty));
            legal.then_some(ty)
        })
        .collect();
    if let Some(types) = hevc
        && in_order(&types, &[32, 33, 34], |t| t <= 9 || (16..=21).contains(&t))
    {
        let mut c = elementary("hevc-annexb", CAP_VIDEO_HEVC, &units);
        if let Some((at, u)) = units.iter().find(|(_, u)| (u[0] >> 1) & 63 == 33) {
            c.see(*at as u64, u.len() as u64, "nal:sps");
        }
        finish(&mut c);
        out.push(c);
    }

    // H.266: two bytes; forbidden and reserved bits zero, a temporal identifier plus one not zero.
    let vvc: Option<Vec<u8>> = units
        .iter()
        .map(|(_, u)| {
            let (a, b) = (*u.first()?, *u.get(1)?);
            let ty = b >> 3;
            let legal = a & 0xc0 == 0 && b & 7 != 0 && ty <= 25;
            legal.then_some(ty)
        })
        .collect();
    if let Some(types) = vvc
        && in_order(&types, &[15, 16], |t| t <= 10)
    {
        let mut c = elementary("vvc-annexb", CAP_VIDEO_VVC, &units);
        if let Some((at, u)) = units.iter().find(|(_, u)| u[1] >> 3 == 15) {
            c.see(*at as u64, u.len() as u64, "nal:sps");
        }
        finish(&mut c);
        out.push(c);
    }
    out
}
