//! Matroska and WebM (media-identify §5.4): the EBML header's document type, the Segment's elements
//! in order, and each track as `Tracks` states it. Clusters are not read.

use super::{clean, get, Candidate, Source, Strength, Track, HEAD, MAX_STEPS};
use serde_json::json;
use std::io;

const EBML: u32 = 0x1a45_dfa3;
const SEGMENT: u32 = 0x1853_8067;
const TRACKS: u32 = 0x1654_ae6b;
const CLUSTER: u32 = 0x1f43_b675;

/// An element's identifier and size at the start of `b`, and the bytes they take: `None` for a
/// size is an unknown size.
fn element(b: &[u8]) -> Option<(u32, Option<u64>, usize)> {
    let first = *b.first()?;
    let idn = first.leading_zeros() as usize + 1;
    if idn > 4 {
        return None;
    }
    let mut id = 0u32;
    for &c in b.get(..idn)? {
        id = (id << 8) | c as u32;
    }
    let s = *b.get(idn)?;
    let sn = s.leading_zeros() as usize + 1;
    if sn > 8 {
        return None;
    }
    let mut size = (s as u64) & (0xffu64 >> sn);
    let mut unknown = size == (0xffu64 >> sn);
    for &c in b.get(idn + 1..idn + sn)? {
        size = (size << 8) | c as u64;
        unknown &= c == 0xff;
    }
    Some((id, (!unknown).then_some(size), idn + sn))
}

/// The elements inside `b`, each with a known size that fits.
fn children(mut b: &[u8]) -> Vec<(u32, &[u8])> {
    let mut out = Vec::new();
    while let Some((id, Some(size), hdr)) = element(b) {
        let Some(end) = usize::try_from(size).ok().and_then(|s| hdr.checked_add(s)) else { break };
        let Some(content) = b.get(hdr..end) else { break };
        out.push((id, content));
        b = &b[end..];
        if out.len() as u64 == MAX_STEPS {
            break;
        }
    }
    out
}

fn uint(b: &[u8]) -> Option<u64> {
    if b.is_empty() || b.len() > 8 {
        return None;
    }
    Some(b.iter().fold(0u64, |a, &c| (a << 8) | c as u64))
}

/// A sampling frequency: an IEEE float of four or eight bytes, taken toward zero to whole hertz.
fn rate(b: &[u8]) -> Option<u64> {
    match b.len() {
        4 => Some(f32::from_be_bytes(b.try_into().ok()?) as u64),
        8 => Some(f64::from_be_bytes(b.try_into().ok()?) as u64),
        _ => None,
    }
}

pub(super) fn codec(id: &[u8]) -> Option<&'static str> {
    if id.starts_with(b"A_AAC") {
        return Some("aac");
    }
    Some(match id {
        b"V_MPEG4/ISO/AVC" => "avc",
        b"V_MPEGH/ISO/HEVC" => "hevc",
        b"V_MPEGI/ISO/VVC" => "vvc",
        b"V_AV1" => "av1",
        b"V_VP8" => "vp8",
        b"V_VP9" => "vp9",
        b"A_OPUS" => "opus",
        b"A_FLAC" => "flac",
        b"A_VORBIS" => "vorbis",
        b"A_AC3" => "ac3",
        b"A_EAC3" => "eac3",
        b"A_MPEG/L3" => "mp3",
        b"S_TEXT/WEBVTT" => "webvtt",
        _ => return None,
    })
}

fn track(entry: &[u8]) -> Track {
    let mut t = Track { media: "other", ..Track::default() };
    for (id, b) in children(entry) {
        match id {
            0xd7 => t.track = uint(b),
            0x83 => {
                t.media = match uint(b) {
                    Some(1) => "video",
                    Some(2) => "audio",
                    Some(17) => "text",
                    _ => "other",
                }
            }
            0x86 => {
                t.stated = clean(b);
                t.codec = codec(b);
            }
            0xe0 => {
                for (k, v) in children(b) {
                    match k {
                        0xb0 => t.width = uint(v),
                        0xba => t.height = uint(v),
                        _ => {}
                    }
                }
            }
            0xe1 => {
                for (k, v) in children(b) {
                    match k {
                        0xb5 => t.sample_rate = rate(v),
                        0x9f => t.channels = uint(v),
                        _ => {}
                    }
                }
            }
            0x6d80 => {
                let encrypted = children(b).iter().any(|(k, e)| *k == 0x6240 && children(e).iter().any(|(j, _)| *j == 0x5035));
                if encrypted {
                    t.protected = Some(json!({"scheme": null, "original": null}));
                }
            }
            _ => {}
        }
    }
    t
}

pub fn walk(src: &dyn Source, head: &[u8]) -> io::Result<Option<Candidate>> {
    if head.get(..4) != Some(&EBML.to_be_bytes()[..]) {
        return Ok(None);
    }
    let size = src.size();
    let mut c = Candidate::new("matroska", "container");
    c.see(0, 4, "ebml");
    let Some((_, Some(hsize), hdr)) = element(head) else {
        c.walked_to = Some(0);
        c.stop(0, "invalid");
        return Ok(Some(c));
    };
    let Some(body) = usize::try_from(hsize).ok().and_then(|s| head.get(hdr..hdr.checked_add(s)?)) else {
        c.walked_to = Some(0);
        c.stop(0, "overrun");
        return Ok(Some(c));
    };
    for (id, b) in children(body) {
        match id {
            0x4282 => {
                c.fact("doc_type", clean(b));
            }
            0x4287 => {
                c.fact("doc_type_version", uint(b));
            }
            _ => {}
        }
    }
    let mut pos = (hdr as u64) + hsize;
    let seg = get(src, pos, 12)?;
    let at_segment = element(&seg).filter(|(id, _, _)| *id == SEGMENT);
    let Some((_, ssize, shdr)) = at_segment else {
        c.walked_to = Some(pos);
        c.stop(pos, if seg.len() < 5 { "overrun" } else { "invalid" });
        return Ok(Some(c));
    };
    c.see(pos, 4, "segment");
    let start = pos + shdr as u64;
    let end = ssize.and_then(|s| start.checked_add(s));
    if end.is_some_and(|e| e > size) {
        c.fact("segment_overruns_by", end.unwrap() - size);
    }
    let end = end.unwrap_or(size).min(size);
    pos = start;
    let mut children_read = 0u64;
    let mut clusters = 0u64;
    loop {
        if pos >= end {
            break;
        }
        if children_read == MAX_STEPS {
            c.stop(pos, "limit");
            break;
        }
        let h = get(src, pos, 12)?;
        let Some((id, s, hl)) = element(&h) else {
            c.stop(pos, if h.len() < 12 { "overrun" } else { "invalid" });
            break;
        };
        let Some(s) = s else {
            c.stop(pos, "unknown_size");
            break;
        };
        let Some(next) = pos.checked_add(hl as u64).and_then(|p| p.checked_add(s)).filter(|&n| n <= size) else {
            if id == CLUSTER {
                clusters += 1;
            }
            c.stop(pos, "overrun");
            break;
        };
        if id == TRACKS && s <= HEAD as u64 {
            c.see(pos, 4, "tracks");
            let b = get(src, pos + hl as u64, s as usize)?;
            for (k, entry) in children(&b) {
                if k == 0xae {
                    c.track(track(entry));
                }
            }
        }
        if id == CLUSTER {
            clusters += 1;
        }
        children_read += 1;
        pos = next;
    }
    c.walked_to = Some(pos);
    if clusters > 0 {
        c.fact("clusters", clusters);
    }
    if children_read > 0 {
        c.strength = Strength::Structure;
    }
    Ok(Some(c))
}
