open-standards/wai/format-probe-conformance/format_probe_verify.py

"""The second implementation of the format probes: written from extensions/format-probes.md
and extensions/projection.md, sharing no code with wai-rs.

    python3 -I format_probe_verify.py <corpus-dir> [--fetched <dir>]

It probes every vector of <corpus-dir>/vectors.json and requires the outcome stated there
(the parameters, or the refusal), every binding's outcome, every fact (what an encoder was
asked for), that every file under ext/ and gen/ is a vector's, that vectors.json's probe and
refusal lists are the extension's §2.4 and §12 tables, and that every refusal is given. With
--fetched it requires every fetched LCEVC stream (fetch.sh) and what its encoder was asked for.
Then it mutates copies of the corpus and requires each copy to fail.

Exit status 0 only if everything verifies and every mutated copy fails. Python 3.8 or later;
standard library only.
"""

import copy
import hashlib
import json
import os
import re
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
EXT = os.path.join(HERE, "..", "extensions")


class Refuse(Exception):
    def __init__(self, code):
        Exception.__init__(self, code)
        self.code = code


# ---- bits (§2.2) -------------------------------------------------------------------

class Bits:
    def __init__(self, data):
        self.d = bytes(data)
        self.p = 0

    def left(self):
        return len(self.d) * 8 - self.p

    def u(self, n):
        if self.left() < n:
            raise Refuse("truncated")
        v = 0
        for _ in range(n):
            v = (v << 1) | ((self.d[self.p >> 3] >> (7 - (self.p & 7))) & 1)
            self.p += 1
        return v

    def flag(self):
        return self.u(1) == 1

    def skip(self, n):
        if self.left() < n:
            raise Refuse("truncated")
        self.p += n

    def ue(self):
        zeros = 0
        while self.u(1) == 0:
            zeros += 1
            if zeros > 31:
                raise Refuse("reserved_value")
        return (1 << zeros) - 1 + self.u(zeros) if zeros else 0

    def aligned(self):
        return self.p % 8 == 0

    def align(self, bit):
        while not self.aligned():
            b = self.u(1)
            if bit is not None and b != bit:
                raise Refuse("reserved_value")


def be16(b, at):
    if at + 2 > len(b):
        raise Refuse("truncated")
    return (b[at] << 8) | b[at + 1]


# ---- Annex B (§2.5) ----------------------------------------------------------------

def nal_units(b):
    i = 0
    while i < len(b) and b[i] == 0:
        i += 1
    if i < 2 or i >= len(b) or b[i] != 1:
        raise Refuse("no_signature")
    starts, j = [], 0
    while j + 2 < len(b):
        if b[j] == 0 and b[j + 1] == 0 and b[j + 2] == 1:
            starts.append(j + 3)
            j += 3
        else:
            j += 1
    out = []
    for k, s in enumerate(starts):
        e = starts[k + 1] - 3 if k + 1 < len(starts) else len(b)
        while e > s and b[e - 1] == 0:
            e -= 1
        out.append(b[s:e])
    return out


def unescape(n):
    out, zeros = bytearray(), 0
    for x in n:
        if zeros >= 2 and x == 3:
            zeros = 0
            continue
        zeros = zeros + 1 if x == 0 else 0
        out.append(x)
    return bytes(out)


def crop(chroma, w, h, c):
    sw, sh = {1: (2, 2), 2: (2, 1)}.get(chroma, (1, 1))
    dw, dh = sw * (c[0] + c[1]), sh * (c[2] + c[3])
    if dw >= w or dh >= h:
        raise Refuse("reserved_value")
    return [w - dw, h - dh]


# ---- JPEG XS (§3) ------------------------------------------------------------------

def spec_text(name):
    return open(os.path.join(EXT, name), encoding="utf-8").read()


def section(text, heading):
    at = text.index("\n" + heading)
    nxt = text.find("\n## ", at + 1)
    return text[at:nxt if nxt >= 0 else len(text)]


FP = spec_text("format-probes.md")
PROFILES = {}
for row in re.findall(r"^\|(.*)\|$", section(FP, "## 3. JPEG XS"), re.M):
    cells = [c.strip() for c in row.split("|")]
    for k in range(0, len(cells) - 1, 2):
        m = re.fullmatch(r"`0x([0-9A-F]{4})`", cells[k])
        if m and cells[k + 1]:
            PROFILES[int(m.group(1), 16)] = cells[k + 1]
LEVELS = {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"}
SUBLEVELS = {0x00: "Unrestricted", 0x80: "Full", 0x10: "Sublev12bpp", 0x0C: "Sublev9bpp", 0x08: "Sublev6bpp",
             0x06: "Sublev4bpp", 0x04: "Sublev3bpp", 0x03: "Sublev2bpp"}


def jxs_codestream(b):
    if len(b) < 2 or be16(b, 0) != 0xFF10:
        raise Refuse("no_signature")
    if be16(b, 2) != 0xFF50:
        raise Refuse("out_of_order")
    lcap = be16(b, 4)
    if lcap < 2:
        raise Refuse("length_mismatch")
    if 4 + lcap > len(b):
        raise Refuse("truncated")
    cap_bits = [i * 8 + k for i, x in enumerate(b[6:4 + lcap]) for k in range(8) if x & (0x80 >> k)]
    at = 4 + lcap
    if be16(b, at) != 0xFF12:
        raise Refuse("out_of_order")
    if be16(b, at + 2) != 26:
        raise Refuse("length_mismatch")
    if at + 28 > len(b):
        raise Refuse("truncated")
    p = b[at + 4:at + 28]
    lcod, ppih, plev, wf, hf, cw, hsl = struct.unpack(">IHHHHHH", p[:16])
    nc, ng, ss, bw = p[16], p[17], p[18], p[19]
    fq, br = p[20] >> 4, p[20] & 15
    cpih = p[21] & 15
    nlx, nly = p[22] >> 4, p[22] & 15
    lh, rl, qpih = p[23] >> 7, (p[23] >> 6) & 1, (p[23] >> 4) & 3
    if wf == 0 or hf == 0 or hsl == 0 or not 1 <= nc <= 8 or ng != 4 or ss != 8 or br != 4:
        raise Refuse("reserved_value")
    if cpih not in (0, 1, 3) or qpih > 1:
        raise Refuse("reserved_value")
    at += 28
    cdt, wgt, cts, crg = None, False, False, False
    while True:
        m = be16(b, at)
        if m == 0xFF20:
            break
        ln = be16(b, at + 2)
        if ln < 2:
            raise Refuse("length_mismatch")
        if at + 2 + ln > len(b):
            raise Refuse("truncated")
        body = b[at + 4:at + 2 + ln]
        if m == 0xFF13:
            if ln != 2 + 2 * nc:
                raise Refuse("length_mismatch")
            comps = [[body[2 * i], body[2 * i + 1] >> 4, body[2 * i + 1] & 15] for i in range(nc)]
            if any(0 in c for c in comps):
                raise Refuse("reserved_value")
            cdt = comps
        elif m == 0xFF14:
            wgt = True
        elif m == 0xFF18:
            cts = True
        elif m == 0xFF19:
            crg = True
        elif m in (0xFF15, 0xFF16, 0xFF17):
            pass
        elif m in (0xFF10, 0xFF11, 0xFF12, 0xFF50):
            raise Refuse("out_of_order")
        else:
            raise Refuse("reserved_value")
        at += 2 + ln
    if cdt is None or not wgt or (cpih == 3 and not (cts and crg)):
        raise Refuse("missing_structure")
    if len(b) < at + 4 or be16(b, len(b) - 2) != 0xFF11:
        raise Refuse("missing_structure")
    if lcod != 0 and lcod != len(b):
        raise Refuse("length_mismatch")
    structure = None
    if len(cdt) >= 3 and cdt[0][1:] == [1, 1] and cdt[1][1:] == cdt[2][1:]:
        structure = {(1, 1): "4:4:4", (2, 1): "4:2:2", (2, 2): "4:2:0"}.get(tuple(cdt[1][1:]))
    return {
        "profile": "0x%04x" % ppih, "level": "0x%04x" % plev, "width": wf, "height": hf,
        "components": cdt, "structure": structure, "colour_transform": cpih,
        "decompositions": [nlx, nly], "precinct_width": cw, "slice_height": hsl,
        "wavelet_precision": bw, "fractional_bits": fq, "quantizer": qpih,
        "long_header": lh == 1, "raw_mode": rl == 1, "cbr_bytes": lcod, "cap_bits": cap_bits,
        "rtp": {
            "profile": PROFILES.get(ppih), "level": LEVELS.get(plev >> 8), "sublevel": SUBLEVELS.get(plev & 0xFF),
            "depth": max(c[0] for c in cdt), "width": wf, "height": hf,
            "sampling": "RGB" if cpih != 0 and structure == "4:4:4" else None,
        },
    }


def jxs_video(b):
    if len(b) < 20:
        raise Refuse("truncated")
    n, h, w, num, den = struct.unpack("<5I", b[:20])
    if 0 in (n, h, w, num, den):
        raise Refuse("framing_invalid")
    at, first, lcods = 20, None, []
    for _ in range(n):
        if at + 4 > len(b):
            raise Refuse("truncated")
        ln = struct.unpack("<I", b[at:at + 4])[0]
        if at + 4 + ln > len(b):
            raise Refuse("truncated")
        c = jxs_codestream(b[at + 4:at + 4 + ln])
        if c["width"] > w or c["height"] > h:
            raise Refuse("framing_invalid")
        if first is None:
            first = c
        lcods.append(c["cbr_bytes"])
        at += 4 + ln
    if at != len(b):
        raise Refuse("trailing_bytes")
    constant = lcods[0] if all(x != 0 and x == lcods[0] for x in lcods) else None
    return {"frames": n, "size": [w, h], "fps": [num, den], "constant_bytes": constant, "first": first}


# ---- H.265 (§4) --------------------------------------------------------------------

def h265_ptl(r, present, max_sub):
    out = {"idc": None, "tier": None, "compat": 0}
    if present:
        r.u(2)
        out["tier"] = r.u(1)
        out["idc"] = r.u(5)
        out["compat"] = r.u(32)
        r.skip(48)
    out["level"] = r.u(8)
    flags = [(r.flag(), r.flag()) for _ in range(max_sub)]
    if max_sub > 0:
        r.skip((8 - max_sub) * 2)
    for p, l in flags:
        if p:
            r.skip(88)
        if l:
            r.skip(8)
    return out


def h265_hrd(r, common, prev, max_sub):
    nal, vcl, sub = prev
    if common:
        nal, vcl, sub = r.flag(), r.flag(), False
        if nal or vcl:
            sub = r.flag()
            if sub:
                r.skip(19)
            r.skip(8)
            if sub:
                r.skip(4)
            r.skip(15)
    for _ in range(max_sub + 1):
        general = r.flag()
        within = True if general else r.flag()
        low_delay = False
        if within:
            r.ue()
        else:
            low_delay = r.flag()
        cnt = 1 if low_delay else r.ue() + 1
        if cnt > 32:
            raise Refuse("reserved_value")
        for present in (nal, vcl):
            if present:
                for _ in range(cnt):
                    r.ue()
                    r.ue()
                    if sub:
                        r.ue()
                        r.ue()
                    r.skip(1)
    return (nal, vcl, sub)


def h265_vps(rbsp):
    r = Bits(rbsp)
    r.u(4)
    internal = r.flag()
    r.flag()
    max_layers = r.u(6)
    max_sub = r.u(3)
    if max_sub > 6:
        raise Refuse("reserved_value")
    r.flag()
    if r.u(16) != 0xFFFF:
        raise Refuse("reserved_value")
    base = h265_ptl(r, True, max_sub)
    ordering = r.flag()
    for _ in range(0 if ordering else max_sub, max_sub + 1):
        r.ue()
        r.ue()
        r.ue()
    max_layer_id = r.u(6)
    sets = r.ue()
    if sets > 1023:
        raise Refuse("reserved_value")
    r.skip(sets * (max_layer_id + 1))
    if r.flag():
        r.skip(64)
        if r.flag():
            r.ue()
        num_hrd = r.ue()
        if num_hrd > sets + 1:
            raise Refuse("reserved_value")
        prev = (False, False, False)
        for i in range(num_hrd):
            r.ue()
            common = r.flag() if i > 0 else True
            prev = h265_hrd(r, common, prev, max_sub)
    out = {"max_layers": max_layers, "ext": None}
    if not r.flag():
        return out
    r.align(1)
    if max_layers == 0:
        return out
    out["ext"] = h265_vps_ext(r, internal, max_layers, max_sub, base)
    return out


def inherit(p, prev):
    if p["idc"] is None:
        p = dict(p, idc=prev["idc"], tier=prev["tier"], compat=prev["compat"])
    return p


def h265_vps_ext(r, internal, max_layers, max_sub, base):
    ptls = []
    if internal:
        ptls.append(inherit(h265_ptl(r, False, max_sub), base))
    splitting = r.flag()
    mask = 0
    for i in range(16):
        if r.flag():
            mask |= 1 << i
    types = bin(mask).count("1")
    lens = [0] * types
    for j in range(types - (1 if splitting else 0)):
        lens[j] = r.u(3) + 1
    ml = min(62, max_layers)
    if splitting and types > 0:
        used = sum(lens[:types - 1])
        if used >= 6:
            raise Refuse("reserved_value")
        lens[types - 1] = 6 - used
    id_present = r.flag()
    ids = [0] * (ml + 1)
    dims = [[0] * types for _ in range(ml + 1)]
    for i in range(1, ml + 1):
        ids[i] = r.u(6) if id_present else i
        if ids[i] <= ids[i - 1]:
            raise Refuse("reserved_value")
        if not splitting:
            for j in range(types):
                dims[i][j] = r.u(lens[j])
    if splitting:
        for i in range(ml + 1):
            off = 0
            for j in range(types):
                dims[i][j] = (ids[i] >> off) & ((1 << lens[j]) - 1)
                off += lens[j]
    view_order, aux, views = [], [], 1
    for i in range(ml + 1):
        sid, j = [0] * 16, 0
        for s in range(16):
            if mask & (1 << s):
                sid[s] = dims[i][j]
                j += 1
        if i > 0 and sid[1] not in view_order:
            views += 1
        view_order.append(sid[1])
        aux.append(sid[3])
    vlen = r.u(4)
    view_ids = [r.u(vlen) for _ in range(views)] if vlen else [0] * views
    n = ml + 1
    deps = [[False] * n for _ in range(n)]
    for i in range(1, n):
        for j in range(i):
            deps[i][j] = r.flag()
    dep = [row[:] for row in deps]
    for i in range(n):
        for j in range(n):
            for k in range(i):
                if deps[i][k] and dep[k][j]:
                    dep[i][j] = True
    sizes, listed = [], set()
    for i in range(n):
        if any(deps[i]):
            continue
        h = 1
        for j in range(n):
            if dep[j][i] and ids[j] not in listed:
                listed.add(ids[j])
                h += 1
        sizes.append(h)
    add = r.ue() if len(sizes) > 1 else 0
    if add > 1023:
        raise Refuse("reserved_value")
    for _ in range(add):
        for s in sizes[1:]:
            r.skip(s.bit_length())
    if r.flag():
        r.skip(3 * n)
    if r.flag():
        for i in range(ml):
            for j in range(i + 1, ml + 1):
                if deps[j][i]:
                    r.skip(3)
    r.flag()
    num = r.ue()
    if num > 63:
        raise Refuse("reserved_value")
    for _ in range(2 if internal else 1, num + 1):
        present = r.flag()
        p = h265_ptl(r, present, max_sub)
        ptls.append(inherit(p, ptls[-1] if ptls else base))
    return {"mask": mask, "ids": ids, "view_order": view_order, "aux": aux, "views": views, "view_ids": view_ids, "ptls": ptls}


def h265_sps(rbsp):
    r = Bits(rbsp)
    r.u(4)
    max_sub = r.u(3)
    if max_sub > 6:
        raise Refuse("reserved_value")
    r.flag()
    p = h265_ptl(r, True, max_sub)
    if r.ue() > 15:
        raise Refuse("reserved_value")
    chroma = r.ue()
    if chroma > 3:
        raise Refuse("reserved_value")
    if chroma == 3:
        r.skip(1)
    w, h = r.ue(), r.ue()
    if w == 0 or h == 0:
        raise Refuse("reserved_value")
    c = [r.ue() for _ in range(4)] if r.flag() else [0, 0, 0, 0]
    bl, bc = r.ue(), r.ue()
    if bl > 8 or bc > 8:
        raise Refuse("reserved_value")
    return p, chroma, w, h, c, bl + 8, bc + 8


def h265(b, multiview):
    units = nal_units(b)

    def has(t):
        return any(len(u) >= 2 and (u[0] >> 1) & 0x3F == t and ((u[0] << 8 | u[1]) >> 3) & 0x3F == 0 for u in units)
    complete = has(32) and has(33)
    vps, sps, layers, vcl = None, None, set(), 0
    for u in units:
        if len(u) < 2:
            raise Refuse("truncated")
        h = u[0] << 8 | u[1]
        if h & 0x8000:
            raise Refuse("reserved_value")
        nut, layer = (h >> 9) & 0x3F, (h >> 3) & 0x3F
        if h & 7 == 0:
            raise Refuse("reserved_value")
        rbsp = unescape(u[2:])
        if nut < 32:
            if sps is None or vps is None:
                raise Refuse("out_of_order" if complete else "missing_structure")
            vcl += 1
            layers.add(layer)
        elif nut == 32 and vps is None:
            vps = h265_vps(rbsp)
        elif nut == 33 and layer == 0 and sps is None:
            sps = h265_sps(rbsp)
    if vps is None or sps is None or vcl == 0:
        raise Refuse("missing_structure")
    p, chroma, w, h, c, bl, bc = sps
    cropped = crop(chroma, w, h, c)
    e = vps["ext"]
    views = (e["views"] if e["mask"] & 2 else 1) if e else 1
    mvp = bool(e) and any(q["idc"] == 6 or q["compat"] & (1 << 6) for q in e["ptls"])
    if multiview and (views < 2 or not any(x > 0 for x in layers)):
        raise Refuse("not_multiview")
    out = {
        "profile_idc": p["idc"], "tier": p["tier"], "level_idc": p["level"], "chroma_format_idc": chroma,
        "bit_depth": [bl, bc], "coded": [w, h], "cropped": cropped,
        "layers": min(62, vps["max_layers"]) + 1, "layers_present": sorted(layers),
    }
    if e:
        out["multilayer"] = {
            "scalability_mask": e["mask"], "layer_ids": e["ids"], "view_order": e["view_order"], "aux": e["aux"],
            "views": views, "view_ids": e["view_ids"], "profiles": [q["idc"] for q in e["ptls"]], "multiview_profile": mvp,
        }
    return out


# ---- H.266 (§5) --------------------------------------------------------------------

def h266_ptl(r, present, max_sub):
    idc = tier = None
    if present:
        idc, tier = r.u(7), r.u(1)
    level = r.u(8)
    frame_only, multilayer = r.flag(), r.flag()
    if present:
        if r.flag():
            r.skip(71)
            r.skip(r.u(8))
        r.align(0)
    flags = [r.flag() for _ in range(max_sub)]
    r.align(None)
    for f in flags:
        if f:
            r.skip(8)
    subs = []
    if present:
        subs = [r.u(32) for _ in range(r.u(8))]
    return {"idc": idc, "tier": tier, "level": level, "frame_only": frame_only, "multilayer": multilayer, "subs": subs}


def h266(b):
    units = nal_units(b)
    has_sps = any(len(u) >= 2 and u[1] >> 3 == 15 for u in units)
    sps, layers, vcl, vps_layers = None, set(), 0, None
    for u in units:
        if len(u) < 2:
            raise Refuse("truncated")
        h = u[0] << 8 | u[1]
        if h & 0xC000:
            raise Refuse("reserved_value")
        layer, nut = (h >> 8) & 0x3F, (h >> 3) & 0x1F
        if h & 7 == 0 or layer > 55:
            raise Refuse("reserved_value")
        rbsp = unescape(u[2:])
        if nut <= 11:
            if sps is None:
                raise Refuse("out_of_order" if has_sps else "missing_structure")
            vcl += 1
            layers.add(layer)
        elif nut == 14 and vps_layers is None:
            r = Bits(rbsp)
            r.u(4)
            vps_layers = r.u(6) + 1
        elif nut == 15 and sps is None:
            r = Bits(rbsp)
            r.u(4)
            r.u(4)
            max_sub = r.u(3)
            if max_sub > 6:
                raise Refuse("reserved_value")
            chroma = r.u(2)
            ctu = r.u(2)
            if ctu > 2:
                raise Refuse("reserved_value")
            p = h266_ptl(r, True, max_sub) if r.flag() else None
            r.flag()
            if r.flag():
                r.skip(1)
            w, hh = r.ue(), r.ue()
            if w == 0 or hh == 0:
                raise Refuse("reserved_value")
            c = [r.ue() for _ in range(4)] if r.flag() else [0, 0, 0, 0]
            sps = (p, chroma, 1 << (ctu + 5), w, hh, c)
    if sps is None or vcl == 0:
        raise Refuse("missing_structure")
    p, chroma, ctu, w, h, c = sps
    return {
        "profile_idc": p["idc"] if p else None, "tier": p["tier"] if p else None,
        "level_idc": p["level"] if p else None, "frame_only": p["frame_only"] if p else None,
        "multilayer_enabled": p["multilayer"] if p else None, "sub_profiles": p["subs"] if p else [],
        "chroma_format_idc": chroma, "ctu_size": ctu, "coded": [w, h], "cropped": crop(chroma, w, h, c),
        "layers": vps_layers if vps_layers is not None else 1, "layers_present": sorted(layers),
    }


# ---- LCEVC (§6) --------------------------------------------------------------------

TABLE20 = [tuple(int(x) for x in m) for m in re.findall(r"(\d+)×(\d+)", section(FP, "## 6. LCEVC").split("Table 20, types 1 to 50:")[1])]
assert len(TABLE20) == 50, "Table 20 in format-probes §6.5"


def multibyte(b, at):
    v = 0
    for _ in range(10):
        if at >= len(b):
            raise Refuse("truncated")
        x = b[at]
        at += 1
        v = (v << 7) | (x & 0x7F)
        if not x & 0x80:
            return v, at
    raise Refuse("reserved_value")


def lcevc(b):
    units = nal_units(b)
    seq = glob = None
    idr = non = pics = 0
    for i, u in enumerate(units):
        if len(u) < 3:
            raise Refuse("truncated")
        if u[0] & 0xC1 != 0x41 or u[1] != 0xFF:
            raise Refuse("reserved_value")
        t = (u[0] >> 1) & 0x1F
        if t == 29:
            idr += 1
        elif t == 28:
            non += 1
        else:
            raise Refuse("reserved_value")
        if i == 0 and t != 29:
            raise Refuse("out_of_order")
        rbsp = unescape(u[2:])
        if rbsp[-1] != 0x80:
            raise Refuse("reserved_value")
        pay, at = rbsp[:-1], 0
        while at < len(pay):
            st, kind = pay[at] >> 5, pay[at] & 0x1F
            at += 1
            if st <= 5:
                size = st
            elif st == 6:
                raise Refuse("reserved_value")
            else:
                size, at = multibyte(pay, at)
            if at + size > len(pay):
                raise Refuse("length_mismatch")
            body = pay[at:at + size]
            at += size
            if kind == 0 and seq is None:
                seq = lcevc_seq(body)
            elif kind == 1 and glob is None:
                glob = lcevc_glob(body)
            elif kind == 2:
                if seq is None or glob is None:
                    raise Refuse("out_of_order")
                pics += 1
    if seq is None or glob is None:
        raise Refuse("missing_structure")
    out = {"nal_units": [idr, non], "pictures": pics}
    out.update(seq)
    out.update(glob)
    return out


def lcevc_seq(b):
    r = Bits(b)
    prof, lev, sub, win = r.u(4), r.u(4), r.u(2), r.flag()
    if r.u(5) != 0:
        raise Refuse("reserved_value")
    if prof == 15 or lev == 15:
        raise Refuse("outside_probe")
    at, window = 2, None
    if win:
        window = []
        for _ in range(4):
            v, at = multibyte(b, at)
            window.append(v)
    if at != len(b):
        raise Refuse("length_mismatch")
    return {"profile_idc": prof, "level_idc": lev, "sublevel_idc": sub, "conformance_window": window}


def lcevc_glob(b):
    r = Bits(b)
    planes, rt, tt = r.flag(), r.u(6), r.u(1)
    chroma, bd, ed = r.u(2), r.u(2), r.u(2)
    tsw = r.flag()
    r.flag()
    r.flag()
    temporal, up, l1 = r.flag(), r.u(3), r.flag()
    r.u(2)
    r.u(2)
    tiles = r.u(2)
    r.u(2)
    r.flag()
    r.flag()
    if planes:
        r.skip(8)
    if tsw:
        r.skip(8)
    if up == 4:
        r.skip(64)
    if l1:
        r.skip(8)
    if tiles > 0:
        if tiles == 3:
            r.skip(32)
        r.skip(8)
    if rt == 0 or 51 <= rt < 63 or up > 4:
        raise Refuse("reserved_value")
    if rt == 63:
        size = [r.u(16), r.u(16)]
        if 0 in size:
            raise Refuse("reserved_value")
    else:
        size = list(TABLE20[rt - 1])
    return {"resolution_type": rt, "size": size, "transform_type": tt, "chroma_sampling_type": chroma,
            "base_depth": 8 + 2 * bd, "enhancement_depth": 8 + 2 * ed, "temporal_enabled": temporal,
            "upsample_type": up, "tile_dimensions_type": tiles}


def unwrap_lcevcbin(b):
    if b[:8] != b"lcevcbin":
        raise ValueError("not the container")
    at, out = 12, b""
    while at < len(b):
        size = struct.unpack(">I", b[at + 2:at + 6])[0]
        out += b[at + 6 + 16:at + 6 + size]
        at += 6 + size
    return out


# ---- MHAS (§7) ---------------------------------------------------------------------

RATES = [96000, 88200, 64000, 48000, 44100, 32000, 24000, 22050, 16000, 12000, 11025, 8000, 7350, None, None,
         57600, 51200, 40000, 38400, 34150, 28800, 25600, 20000, 19200, 17075, 14400, 12800, 9600, None, None, None]


def escaped(r, n1, n2, n3):
    v = r.u(n1)
    if v == (1 << n1) - 1:
        a = r.u(n2)
        v += a
        if a == (1 << n2) - 1:
            v += r.u(n3)
    return v


def mhas(b):
    try:
        first = escaped(Bits(b), 3, 8, 8)
    except Refuse:
        raise Refuse("no_signature")
    if first not in (1, 6):
        raise Refuse("no_signature")
    packets, at = [], 0
    while at < len(b):
        r = Bits(b[at:])
        t, lab, ln = escaped(r, 3, 8, 8), escaped(r, 2, 8, 32), escaped(r, 11, 24, 24)
        s = at + r.p // 8
        if s + ln > len(b):
            raise Refuse("truncated")
        packets.append((t, lab, b[s:s + ln]))
        at = s + ln
    cfg, frames, syncs, configs = None, 0, 0, 0
    for t, lab, pay in packets:
        if t == 6:
            if lab != 0 or pay != b"\xa5":
                raise Refuse("reserved_value")
            syncs += 1
        elif t == 1:
            configs += 1
            if cfg is None:
                cfg = mhas_cfg(pay)
        elif t == 2:
            if cfg is None:
                raise Refuse("out_of_order")
            frames += 1
    if cfg is None:
        raise Refuse("missing_structure")
    cfg.update(packets=len(packets), frames=frames, syncs=syncs, configs=configs)
    return cfg


def mhas_cfg(b):
    r = Bits(b)
    pl, sfi = r.u(8), r.u(5)
    if sfi == 31:
        rate = r.u(24)
    else:
        rate = RATES[sfi]
        if rate is None:
            raise Refuse("reserved_value")
    if rate == 0:
        raise Refuse("reserved_value")
    core = r.u(3)
    r.u(1)
    r.u(1)
    lt = r.u(2)
    cicp = speakers = None
    if lt == 0:
        cicp = r.u(6)
    elif lt < 3:
        speakers = escaped(r, 5, 8, 16) + 1
    profile = level = None
    if 0x0B <= pl <= 0x0F:
        profile, level = "low-complexity", pl - 0x0A
    elif 0x10 <= pl <= 0x14:
        profile, level = "baseline", pl - 0x0F
    return {"profile_level": pl, "profile": profile, "level": level, "sampling_rate": rate,
            "core_frame_length_index": core, "speaker_layout_type": lt, "cicp_layout": cicp, "speakers": speakers}


# ---- AVS3 (§8, §9) -----------------------------------------------------------------

def avs3_units(b):
    i = 0
    while i < len(b) and b[i] == 0:
        i += 1
    if i < 2 or i >= len(b) or b[i] != 1:
        raise Refuse("no_signature")
    starts, j = [], 0
    while j + 3 < len(b):
        if b[j] == 0 and b[j + 1] == 0 and b[j + 2] == 1:
            starts.append(j + 3)
            j += 3
        else:
            j += 1
    return [(b[s], b[s + 1:(starts[k + 1] - 3 if k + 1 < len(starts) else len(b))]) for k, s in enumerate(starts)]


def demulate(b):
    bits = []
    for i, x in enumerate(b):
        n = 6 if i >= 2 and b[i - 2] == 0 and b[i - 1] == 0 and x == 2 else 8
        bits.extend((x >> (7 - k)) & 1 for k in range(n))
    while len(bits) % 8:
        bits.append(0)
    return bytes(int("".join(map(str, bits[k:k + 8])), 2) for k in range(0, len(bits), 8))


def avs3_video(b):
    us = avs3_units(b)
    if not us:
        raise Refuse("no_signature")
    seq = next((p for c, p in us if c == 0xB0), None)
    if seq is None:
        raise Refuse("missing_structure")
    if us[0][0] != 0xB0:
        raise Refuse("out_of_order")
    r = Bits(demulate(seq))

    def marker():
        if r.u(1) != 1:
            raise Refuse("reserved_value")
    prof, lev, prog, field = r.u(8), r.u(8), r.flag(), r.flag()
    if r.flag() or r.flag():
        raise Refuse("outside_probe")
    marker()
    w = r.u(14)
    marker()
    h = r.u(14)
    chroma, sp = r.u(2), r.u(3)
    ep = r.u(3) if prof in (0x22, 0x32) else None
    marker()
    ar, frc = r.u(4), r.u(4)
    marker()
    lo = r.u(18)
    marker()
    hi = r.u(12)
    ld, tid = r.flag(), r.flag()
    marker()
    bbv = r.u(18)
    marker()
    dpb = r.u(4) + 1
    if w == 0 or h == 0 or chroma == 0 or sp == 0 or frc == 0 or frc > 13:
        raise Refuse("reserved_value")
    intra = sum(1 for c, _ in us if c == 0xB3)
    inter = sum(1 for c, _ in us if c == 0xB6)
    if intra == 0:
        raise Refuse("missing_structure")
    return {"profile_id": prof, "level_id": lev, "progressive": prog, "field_coded": field, "size": [w, h],
            "chroma_format": chroma, "sample_precision": sp, "encoding_precision": ep, "aspect_ratio": ar,
            "frame_rate_code": frc, "bit_rate": hi * (1 << 18) + lo, "low_delay": ld, "temporal_id_enabled": tid,
            "bbv_buffer_size": bbv, "max_dpb_size": dpb, "pictures": [intra, inter],
            "sequence_end": any(c == 0xB1 for c, _ in us)}


def avs3_audio(b):
    r = Bits(b)
    try:
        sync = r.u(12)
    except Refuse:
        raise Refuse("no_signature")
    if sync != 0xFFF:
        raise Refuse("no_signature")
    codec = r.u(4)
    if codec == 1:
        raise Refuse("outside_probe")
    if codec != 2:
        raise Refuse("reserved_value")
    if r.flag():
        raise Refuse("outside_probe")
    nn = r.u(3)
    if nn > 1:
        raise Refuse("reserved_value")
    prof = r.u(3)
    if prof > 2:
        raise Refuse("reserved_value")
    sfi = r.u(4)
    if sfi > 8:
        raise Refuse("reserved_value")
    r.skip(8)
    out = {"codec_id": codec, "nn_type": nn, "coding_profile": prof,
           "sampling_rate": [192000, 96000, 48000, 44100, 32000, 24000, 22050, 16000, 8000][sfi]}
    if prof == 0:
        out["channel_number_index"] = r.u(7)
    elif prof == 1:
        bed = r.u(2)
        out["soundbed_type"] = bed
        if bed == 0:
            out["objects"] = r.u(7) + 1
            out["bitrate_index_per_object"] = r.u(4)
        elif bed == 1:
            out["channel_number_index"] = r.u(7)
            out["bitrate_index"] = r.u(4)
            out["objects"] = r.u(7) + 1
            out["bitrate_index_per_object"] = r.u(4)
        else:
            raise Refuse("reserved_value")
    else:
        order = r.u(4) + 1
        if not 1 <= order <= 3:
            raise Refuse("reserved_value")
        out["ambisonic_order"] = order
    res = r.u(2)
    if res > 2:
        raise Refuse("reserved_value")
    out["resolution_bits"] = 8 * (res + 1)
    if prof != 1:
        out["bitrate_index"] = r.u(4)
    r.skip(8)
    return out


# ---- leb128 (§14.2, §16.2) ---------------------------------------------------------

def leb(b, at):
    """(value, bytes taken). At most eight bytes, the value at most 2^32 - 1."""
    v = 0
    for i in range(8):
        if at + i >= len(b):
            raise Refuse("truncated")
        x = b[at + i]
        v |= (x & 0x7F) << (7 * i)
        if not x & 0x80:
            if v > 0xFFFFFFFF:
                raise Refuse("reserved_value")
            return v, i + 1
    raise Refuse("reserved_value")


def inside(f, *args):
    """A structure read inside a unit whose declared size is too short: length_mismatch."""
    try:
        return f(*args)
    except Refuse as e:
        if e.code == "truncated":
            raise Refuse("length_mismatch")
        raise


# ---- AV2 (§14) ---------------------------------------------------------------------

AV2_FRAMES = {4, 5, 6, 7, 10, 11, 12, 13, 14, 19, 21}
AV2_LEVELS = ["2.0", "2.1", "3.0", "3.1", "4.0", "4.1", "5.0", "5.1", "5.2", "5.3", "6.0", "6.1", "6.2", "6.3",
              "7.0", "7.1", "7.2", "7.3", "8.0", "8.1", "8.2", "8.3"]


def uvlc(r):
    lz = 0
    while r.u(1) == 0:
        lz += 1
        if lz >= 32:
            raise Refuse("reserved_value")
    return r.u(lz) + (1 << lz) - 1


def av2(b):
    try:
        size, n = leb(b, 0)
        if size == 0 or (b[n] >> 2) & 31 != 2:
            raise Refuse("no_signature")
    except (Refuse, IndexError):
        raise Refuse("no_signature")
    obus, at = [], 0
    while at < len(b):
        size, n = leb(b, at)
        at += n
        if at + size > len(b):
            raise Refuse("truncated")
        o = b[at:at + size]
        if not o:
            raise Refuse("length_mismatch")
        typ, tl = (o[0] >> 2) & 31, o[0] & 3
        if o[0] & 0x80:
            if len(o) < 2:
                raise Refuse("length_mismatch")
            ml, xl, pay = o[1] >> 5, o[1] & 31, o[2:]
        else:
            ml, xl, pay = 0, (31 if typ in (2, 20) else 0), o[1:]
        obus.append((typ, tl, ml, xl, pay))
        at += size
    for typ, tl, ml, xl, _ in obus:
        if typ in (2, 20) and xl != 31:
            raise Refuse("reserved_value")
        if typ == 1 and (tl or ml):
            raise Refuse("reserved_value")
    types = [o[0] for o in obus]
    if 1 not in types or not any(t in AV2_FRAMES for t in types):
        raise Refuse("missing_structure")
    sh = types.index(1)
    if min(i for i, t in enumerate(types) if t in AV2_FRAMES) < sh:
        raise Refuse("out_of_order")
    out = inside(av2_header, obus[sh][4])
    out.update(temporal_units=types.count(2), frames=sum(t in AV2_FRAMES for t in types), obus=len(obus), codec=None)
    return out


def av2_header(p):
    r = Bits(p)

    def check(ok):
        if not ok:
            raise Refuse("reserved_value")
    hid = uvlc(r)
    check(hid < 16)
    prof = r.u(5)
    check(not 5 <= prof <= 30)
    single = r.flag()
    lev = r.u(5)
    check(not 22 <= lev <= 30)
    tier = r.u(1) if lev > 3 and not single else 0
    cf = uvlc(r)
    if cf > 3:
        raise Refuse("reserved_value")
    bdi = uvlc(r)
    if bdi > 1:
        raise Refuse("reserved_value")
    allowed = {0: {0, 1}, 1: {0, 1}, 2: {0, 1}, 3: {0, 1, 3}, 4: {0, 1, 2}}.get(prof, {0, 1, 2, 3})
    if cf not in allowed:
        raise Refuse("reserved_value")
    if single:
        still, mt, mm, cnt = True, 0, 0, 1
    else:
        r.u(3)
        still, mt, mm, cnt = r.flag(), r.u(2), r.u(3), 1
        if mm > 0:
            m1 = r.u(mm.bit_length())  # CeilLog2(max_mlayer_id + 1)
            if m1 > mm:
                raise Refuse("reserved_value")
            cnt = m1 + 1
        r.u(1)
    wb, hb = r.u(4) + 1, r.u(4) + 1
    w, h = r.u(wb) + 1, r.u(hb) + 1
    crop = [uvlc(r) for _ in range(4)] if r.flag() else [0, 0, 0, 0]
    l, rt, t, bt = crop
    if max(l, rt) > w - 1 or max(t, bt) > h - 1 or l + rt >= w or t + bt >= h:
        raise Refuse("reserved_value")
    sub = {0: [1, 1], 1: [1, 1], 2: [0, 0], 3: [1, 0]}[cf]
    return {"seq_header_id": hid, "profile": prof, "level_idx": lev, "level": "maximum" if lev == 31 else AV2_LEVELS[lev],
            "tier": tier, "single_picture_header": single, "still_picture": still, "chroma_format_idc": cf,
            "monochrome": cf == 1, "subsampling": sub, "bit_depth": 10 if bdi == 0 else 8, "max_size": [w, h],
            "cropped": [w - l - rt, h - t - bt], "max_tlayer_id": mt, "max_mlayer_id": mm, "mlayers": cnt}


# ---- AC-4 (§15) --------------------------------------------------------------------

AC4_RATES = [(24000, 1001, 1920), (24, 1, 1920), (25, 1, 2048), (30000, 1001, 1536), (30, 1, 1536), (48000, 1001, 960),
             (48, 1, 960), (50, 1, 1024), (60000, 1001, 768), (60, 1, 768), (100, 1, 512), (120000, 1001, 384),
             (120, 1, 384), (375, 16, 2048)]


def crc16_ac4(data):
    c = 0
    for x in data:
        c ^= x << 8
        for _ in range(8):
            c = ((c << 1) ^ 0x8005) & 0xFFFF if c & 0x8000 else (c << 1) & 0xFFFF
    return c


def var_bits(r, n):
    v = 0
    while True:
        v += r.u(n)
        if not r.flag():
            return v
        v = (v << n) + (1 << n)
        if v > 0xFFFFFFFF:
            raise Refuse("outside_probe")


def ac4(b):
    if b[:2] not in (b"\xac\x40", b"\xac\x41"):
        raise Refuse("no_signature")
    frames, at = [], 0
    while at < len(b):
        if at + 2 > len(b):
            raise Refuse("truncated")
        sync = b[at:at + 2]
        if sync not in (b"\xac\x40", b"\xac\x41"):
            raise Refuse("reserved_value")
        if at + 4 > len(b):
            raise Refuse("truncated")
        size, head = int.from_bytes(b[at + 2:at + 4], "big"), 2
        if size == 0xFFFF:
            if at + 7 > len(b):
                raise Refuse("truncated")
            size, head = int.from_bytes(b[at + 4:at + 7], "big"), 5
        start = at + 2 + head
        end = start + size + (2 if sync == b"\xac\x41" else 0)
        if end > len(b):
            raise Refuse("truncated")
        if sync == b"\xac\x41" and crc16_ac4(b[at + 2:end]) != 0:
            raise Refuse("reserved_value")
        frames.append((sync == b"\xac\x41", b[start:start + size]))
        at = end
    out = inside(ac4_toc, frames[0][1])
    out.update(frames=len(frames), crc_frames=sum(1 for c, _ in frames if c))
    return out


def ac4_toc(raw):
    r = Bits(raw)
    bv = r.u(2)
    if bv == 3:
        bv += var_bits(r, 2)
    if bv > 2:
        raise Refuse("outside_probe")
    counter = r.u(10)
    wait = None
    if r.flag():
        wait = r.u(3)
        if wait:
            r.u(2)
    fs, fri = r.u(1), r.u(4)
    if fri > 13 or (fs == 0 and fri != 13):
        raise Refuse("reserved_value")
    ifr = r.flag()
    if r.flag():
        n = 1
    elif r.flag():
        n = var_bits(r, 2) + 2
    else:
        n = 0
    if r.flag() and r.u(5) == 31:
        var_bits(r, 3)
    prog = None
    if bv >= 2 and r.flag():
        prog = r.u(16)
        if r.flag():
            r.skip(128)
    pres = codec = None
    if n:
        cfg = None
        if not r.flag():
            cfg = r.u(3)
            if cfg == 7:
                cfg += var_bits(r, 2)
        pv = 0
        while r.flag():
            pv += 1
            if pv > 2:
                raise Refuse("reserved_value")
        if (bv, pv) == (1, 1):
            good = cfg == 7
        elif (bv, pv) in ((0, 0), (1, 0), (2, 1), (2, 2)):
            good = cfg is None or cfg <= 6
        else:
            good = False
        if not good:
            raise Refuse("reserved_value")
        md = None if cfg == 6 else r.u(3)
        pres = {"config": cfg, "version": pv, "md_compat": md}
        if md is not None:
            codec = "ac-4.%02x.%02x.%02x" % (bv, pv, md)
    num, den, length = (11025, 512, 2048) if fs == 0 else AC4_RATES[fri]
    return {"bitstream_version": bv, "sequence_counter": counter, "wait_frames": wait,
            "sampling_rate": 48000 if fs else 44100, "frame_rate_index": fri, "frame_rate": [num, den],
            "frame_length": length, "iframe_global": ifr, "presentations": n, "program_id": prog,
            "presentation": pres, "codec": codec}


# ---- IAMF (§16) --------------------------------------------------------------------

class Body:
    def __init__(self, b):
        self.b, self.at = b, 0

    def take(self, n):
        if self.at + n > len(self.b):
            raise Refuse("length_mismatch")
        s = self.b[self.at:self.at + n]
        self.at += n
        return s

    def num(self, n):
        return int.from_bytes(self.take(n), "big")

    def leb(self):
        v, k = inside(leb, self.b, self.at)
        self.at += k
        return v


def iamf_obu(b, at):
    if at >= len(b):
        raise Refuse("truncated")
    h = b[at]
    typ, red, bit, ext = h >> 3, (h >> 2) & 1, (h >> 1) & 1, h & 1
    size, n = leb(b, at + 1)
    s = at + 1 + n
    if s + size > len(b):
        raise Refuse("truncated")
    frame = 5 <= typ <= 23
    if bit and not (typ in (2, 4) or frame):
        raise Refuse("reserved_value")
    if red and (typ == 4 or frame):
        raise Refuse("reserved_value")
    body = Body(b[s:s + size])
    if frame and bit:
        body.leb()
        body.leb()
    if ext:
        body.take(body.leb())
    return (typ, red, body.b[body.at:]), s + size


FLAC_RATES = {88200, 176400, 192000, 8000, 16000, 22050, 24000, 32000, 44100, 48000, 96000}


def iamf_codec(p):
    r = Body(p)
    cid = r.leb()
    fourcc = r.take(4)
    if fourcc not in (b"Opus", b"mp4a", b"fLaC", b"ipcm"):
        raise Refuse("outside_probe")
    n = r.leb()
    if n == 0:
        raise Refuse("reserved_value")
    roll = r.num(2)
    roll -= 0x10000 if roll >= 0x8000 else 0
    want = {b"Opus": -((3840 + n - 1) // n), b"mp4a": -1}.get(fourcc, 0)
    if roll != want:
        raise Refuse("reserved_value")
    rate = depth = None
    if fourcc == b"Opus":
        r.num(1)
        ch = r.num(1)
        r.num(2)
        r.num(4)
        gain, fam = r.num(2), r.num(1)
        if ch != 2 or gain or fam:
            raise Refuse("reserved_value")
        rate = 48000
    elif fourcc == b"fLaC":
        first = True
        while True:
            hdr = r.num(4)
            last, kind, ln = hdr >> 31, (hdr >> 24) & 0x7F, hdr & 0xFFFFFF
            blk = r.take(ln)
            if first:
                if kind != 0 or ln != 34:
                    raise Refuse("reserved_value")
                bits = Bits(blk)
                mn, mx = bits.u(16), bits.u(16)
                bits.u(24)
                bits.u(24)
                sr, ch, bps = bits.u(20), bits.u(3) + 1, bits.u(5) + 1
                if mn != n or mx != n or sr not in FLAC_RATES or ch != 2 or bps not in (16, 24, 32):
                    raise Refuse("reserved_value")
                rate, depth, first = sr, bps, False
            if last:
                break
    elif fourcc == b"ipcm":
        flags, size, sr = r.num(1), r.num(1), r.num(4)
        if flags > 1 or size not in (16, 24, 32) or sr not in (16000, 32000, 44100, 48000, 96000):
            raise Refuse("reserved_value")
        rate, depth = sr, size
    return {"id": cid, "codec_id": fourcc.decode(), "samples_per_frame": n, "roll_distance": roll,
            "sample_rate": rate, "bit_depth": depth}


def iamf_element(p):
    r = Body(p)
    eid = r.leb()
    kind = r.num(1) >> 5
    cfg = r.leb()
    n = r.leb()
    if n == 0:
        raise Refuse("reserved_value")
    subs = [r.leb() for _ in range(n)]
    if kind > 2:
        raise Refuse("outside_probe")
    return {"id": eid, "type": ["channel", "scene", "object"][kind], "codec_config_id": cfg, "substreams": subs}


def iamf(b):
    try:
        (typ, _, pay), _ = iamf_obu(b, 0)
        ok = typ == 31 and pay[:4] == b"iamf"
    except Refuse:
        ok = False
    if not ok:
        raise Refuse("no_signature")
    obus, at = [], 0
    while at < len(b):
        o, at = iamf_obu(b, at)
        obus.append(o)
    r = Body(obus[0][2])
    r.take(4)
    prim, add = r.num(1), r.num(1)
    if prim > 5:
        raise Refuse("reserved_value")
    data = next((i for i, o in enumerate(obus) if o[0] in (3, 4) or 5 <= o[0] <= 23), len(obus))
    desc = obus[1:data]
    if data < len(obus) and not any(o[0] == 2 for o in desc):
        raise Refuse("out_of_order" if any(o[0] == 2 for o in obus[data:]) else "missing_structure")
    stage, cfgs, els, mixes = 0, [], [], 0
    for typ, _, pay in desc:
        if typ not in (0, 1, 2):
            continue
        if typ + 1 < stage:
            raise Refuse("out_of_order")
        stage = typ + 1
        if typ == 0:
            cfgs.append(iamf_codec(pay))
        elif typ == 1:
            els.append(iamf_element(pay))
        else:
            mixes += 1
    if not cfgs or not els or not mixes:
        raise Refuse("missing_structure")
    if any(e["codec_config_id"] not in [c["id"] for c in cfgs] for e in els):
        raise Refuse("missing_structure")
    ids = [c["codec_id"] for c in cfgs]
    if len(ids) > 2 or (len(ids) == 2 and "ipcm" not in ids):
        raise Refuse("reserved_value")
    named = next((c for c in ids if c != "ipcm"), "ipcm")
    units = frames = blocks = 0
    seqs = 1
    for typ, red, pay in obus[data:]:
        if typ == 4:
            units += 1
        elif typ == 3:
            blocks += 1
        elif typ == 5:
            if Body(pay).leb() <= 17:
                raise Refuse("reserved_value")
            frames += 1
        elif 6 <= typ <= 23:
            frames += 1
        elif typ == 31 and not red:
            seqs += 1
    if not frames:
        raise Refuse("missing_structure")
    element = {"Opus": "Opus", "mp4a": "mp4a.40.2", "fLaC": "fLaC", "ipcm": "ipcm"}[named]
    return {"primary_profile": prim, "additional_profile": add, "codec_configs": cfgs, "audio_elements": els,
            "mix_presentations": mixes, "temporal_units": units, "audio_frames": frames, "parameter_blocks": blocks,
            "sequences": seqs, "codec": "iamf.%03d.%03d.%s" % (prim, add, element)}


# ---- IVAS (§17) --------------------------------------------------------------------

IVAS_RATES = [13200, 16400, 24400, 32000, 48000, 64000, 80000, 96000, 128000, 160000, 192000, 256000, 384000, 512000]
EVS_RATES = [2800, 7200, 8000, 9600, 13200, 16400, 24400, 32000, 48000, 64000, 96000, 128000, 2400]


def ivas_frame(ft):
    evs, ind, br = ft >> 5, (ft >> 4) & 1, ft & 15
    if not (evs == 0 and ind == 1):
        if br == 15:
            return {"mode": "no_data", "bitrate": None, "sid": False, "bytes": 0}
        if br == 14:
            return {"mode": "speech_lost", "bitrate": None, "sid": False, "bytes": 0}
        if evs == 1:
            raise Refuse("outside_probe" if br <= 9 else "reserved_value")
        if br == 13:
            raise Refuse("reserved_value")
        rate = EVS_RATES[br]
        return {"mode": "evs", "bitrate": rate, "sid": br == 12, "bytes": rate // 400}
    if br == 14:
        raise Refuse("outside_probe")
    rate = 5200 if br == 15 else IVAS_RATES[br]
    return {"mode": "ivas", "bitrate": rate, "sid": br == 15, "bytes": rate // 400}


def ivas_format(data, f):
    r = Bits(data)
    if f["sid"]:
        v = r.u(3)
        if v == 4:
            raise Refuse("reserved_value")
        return {0: "STEREO", 1: "STEREO", 2: "ISM", 3: "MASA", 7: "MASA", 5: "SBA", 6: "SBA"}[v]
    v = r.u(2)
    if v < 2:
        return ["STEREO", "MC"][v]
    if v == 3:
        return "MASA" if r.u(1) else "SBA"
    if f["bitrate"] < 24400:
        return "ISM"
    if r.u(1) == 0:
        return "ISM"
    return "OSBA" if r.u(1) else "OMASA"


def ivas(b):
    at, cmr, ebytes, pi, skip, tocs = 0, None, 0, False, False, []
    while True:
        if at >= len(b):
            raise Refuse("truncated")
        x = b[at]
        at += 1
        if x & 0x80:
            if at == 1:
                cmr = {"t": (x >> 4) & 7, "value": x & 15}
            elif not skip:
                ebytes += 1
                et = (x >> 4) & 7
                if et == 1 and x & 0x08:
                    if at >= len(b):
                        raise Refuse("truncated")
                    at += 1
                elif et == 2:
                    pi = True
                elif et >= 4:
                    skip = True
            continue
        skip = False
        tocs.append(x & 0x3F)
        if (x & 0x3F) == 0x1E:
            if at >= len(b):
                raise Refuse("truncated")
            at += 1
        if not x & 0x40:
            break
    frames = [ivas_frame(t) for t in tocs]
    fmt = None
    for f in frames:
        if at + f["bytes"] > len(b):
            raise Refuse("truncated")
        if fmt is None and f["mode"] == "ivas":
            fmt = ivas_format(b[at:at + f["bytes"]], f)
        at += f["bytes"]
    rest = b[at:]
    if pi and not rest:
        raise Refuse("missing_structure")
    if not pi and any(rest):
        raise Refuse("trailing_bytes")
    return {"cmr": cmr, "e_bytes": ebytes, "pi_data": pi, "frames": frames, "format": fmt,
            "pi_bytes": len(rest) if pi else 0, "padding_bytes": 0 if pi else len(rest), "codec": "sivs"}


# ---- IEEE 1857.11 (§10) and the projection descriptor (projection.md) ---------------

def nic(b):
    if len(b) < 13:
        raise Refuse("truncated")
    ids = int.from_bytes(b[0:3], "big")
    h, w = struct.unpack(">II", b[3:11])
    fmt = int.from_bytes(b[11:13], "big")
    prof, depth = ids >> 21, (fmt >> 11) & 7
    if prof > 2 or fmt >> 14 or depth not in (1, 2) or h == 0 or w == 0:
        raise Refuse("reserved_value")
    return {"profile_id": prof, "level_id": (ids >> 15) & 0x3F, "model_id": (ids >> 12) & 7,
            "parameter_set_id": (ids >> 5) & 0x7F, "quality_id": ids & 0x1F, "size": [w, h],
            "bit_depth": 8 if depth == 1 else 10, "picture_format": (fmt >> 8) & 7,
            "slice_size": [1 << ((fmt & 15) + 6), 1 << (((fmt >> 4) & 15) + 6)]}


DEG = 1 << 16
PROJ = ["rectilinear", "equirectangular", "half-equirectangular", "fisheye"]
STEREO = ["mono", "side-by-side", "top-bottom", "layers"]


def projection_read(b):
    if len(b) < 4 or b[:4] != b"WPJ1":
        raise Refuse("no_signature")
    if len(b) != 40:
        raise Refuse("length_mismatch")
    if b[4] != 1 or b[5] > 3 or b[6] > 3 or b[7] > 1 or any(b[34:]):
        raise Refuse("reserved_value")
    az, el = struct.unpack("<II", b[8:16])
    ca, ce, ct = struct.unpack("<iii", b[16:28])
    base = struct.unpack("<I", b[28:32])[0]
    pr, st, fv, left, right = b[5], b[6], b[7], b[32], b[33]
    ranges = {
        0: (az == 0 and el == 0) or (0 < az < 180 * DEG and 0 < el < 180 * DEG),
        1: az == 360 * DEG and el == 180 * DEG,
        2: az == 180 * DEG and el == 180 * DEG,
        3: 0 < az <= 360 * DEG and 0 < el <= 360 * DEG,
    }[pr]
    centre = -180 * DEG <= ca < 180 * DEG and -90 * DEG <= ce <= 90 * DEG and -180 * DEG <= ct < 180 * DEG
    if st == 0:
        views = fv == 0 and base == 0 and left == 0 and right == 0
    elif st in (1, 2):
        views = left == 0 and right == 0
    else:
        views = left != right and left <= 62 and right <= 62
    if not (ranges and centre and views):
        raise Refuse("descriptor_invalid")
    return {"projection": PROJ[pr], "stereo": STEREO[st], "first_view": "left" if fv == 0 else "right",
            "coverage": [az, el], "centre": [ca, ce, ct], "baseline_um": base,
            "layers": [left, right] if st == 3 else None}


def binding(desc, base_cap, base_params):
    if desc["stereo"] != "layers":
        return "accept"
    if base_cap != "wai.video.mvhevc" or not all(l in base_params["layers_present"] for l in desc["layers"]):
        return {"refuse": "layer_absent"}
    return "accept"


PROBES = {
    "wai.image.jxs": jxs_codestream, "wai.video.jxs": jxs_video,
    "wai.video.hevc": lambda b: h265(b, False), "wai.video.mvhevc": lambda b: h265(b, True),
    "wai.video.vvc": h266, "wai.video.lcevc": lcevc, "wai.audio.mpegh": mhas,
    "wai.video.avs3": avs3_video, "wai.audio.avs3": avs3_audio,
    "wai.image.ieee1857_11": nic, "wai.meta.projection": projection_read,
    "wai.video.av2": av2, "wai.audio.ac4": ac4, "wai.audio.iamf": iamf, "wai.audio.ivas": ivas,
}


def probe(cap, b):
    if cap not in PROBES:
        return {"refuse": "no_probe"}
    try:
        return {"accept": PROBES[cap](b)}
    except Refuse as e:
        return {"refuse": e.code}


def pointer(doc, ptr):
    for part in ptr.strip("/").split("/"):
        if isinstance(doc, list):
            doc = doc[int(part)]
        elif isinstance(doc, dict) and part in doc:
            doc = doc[part]
        else:
            return None
    return doc


# ---- The corpus --------------------------------------------------------------------

def spec_tables():
    probes = [(m.group(1), m.group(2)) for m in re.finditer(r"^\| `(wai\.[a-z0-9_.]+)` \| (\w+) \| §\d+(?:\.\d+)? \|$", section(FP, "## 2. Common rules"), re.M)]
    codes = re.findall(r"^\| `([a-z_]+)` \|", section(FP, "## 12. Refusals"), re.M)
    return probes, codes


def check(doc, files, fetched=None):
    fails = []
    probes, codes = spec_tables()
    if [(p["capability"], p["basis"]) for p in doc["probes"]] != probes:
        fails.append("vectors.json's probes are not §2.4's table")
    if doc["refusals"] != codes:
        fails.append("vectors.json's refusals are not §12's table")
    if sorted(c for c, _ in probes) != sorted(PROBES):
        fails.append("§2.4's table names a probe this verifier does not have, or the reverse")
    results, given, referenced = {}, set(), set()
    for v in doc["vectors"]:
        name, f = v["name"], v["file"]
        referenced.add(f)
        if f not in files:
            fails.append("%s: %s missing" % (name, f))
            continue
        if hashlib.sha256(files[f]).hexdigest() != v["sha256"]:
            fails.append("%s: %s's digest differs" % (name, f))
        got = probe(v["capability"], files[f])
        results[(f, v["capability"])] = got
        if got != v["outcome"]:
            fails.append("%s (%s): stated %s, probe gives %s" % (name, v["capability"], json.dumps(v["outcome"])[:160], json.dumps(got)[:160]))
        if "refuse" in got:
            given.add(got["refuse"])
    for bnd in doc["bindings"]:
        referenced.update([bnd["descriptor"], bnd["base"]])
        try:
            d = projection_read(files[bnd["descriptor"]])
            base = PROBES[bnd["base_capability"]](files[bnd["base"]])
            got = binding(d, bnd["base_capability"], base)
        except (Refuse, KeyError) as e:
            fails.append("binding %s: %s" % (bnd["name"], e))
            continue
        if got != bnd["outcome"]:
            fails.append("binding %s: stated %s, gives %s" % (bnd["name"], bnd["outcome"], got))
        if isinstance(got, dict):
            given.add(got["refuse"])
    for fact in doc["facts"]:
        caps = [c for (f, c) in results if f == fact["file"]]
        ok = any("accept" in results[(fact["file"], c)] and pointer(results[(fact["file"], c)]["accept"], fact["pointer"]) == fact["value"] for c in caps)
        if not ok:
            fails.append("fact %s%s = %s does not hold" % (fact["file"], fact["pointer"], json.dumps(fact["value"])))
    for c in codes:
        if c not in given:
            fails.append("no vector is refused with %s" % c)
    for f in files:
        if (f.startswith("ext/") or f.startswith("gen/")) and f not in referenced:
            fails.append("%s is no vector's" % f)
    if fetched is not None:
        for e in doc["fetched"]:
            b = fetched.get(e["file"])
            if b is None:
                fails.append("fetched %s missing" % e["file"])
                continue
            if hashlib.sha256(b).hexdigest() != e["sha256"]:
                fails.append("fetched %s: digest differs" % e["file"])
                continue
            got = probe("wai.video.lcevc", unwrap_lcevcbin(b))
            if "accept" not in got:
                fails.append("fetched %s: %s" % (e["file"], got))
                continue
            for ptr, want in e["expect"].items():
                if pointer(got["accept"], ptr) != want:
                    fails.append("fetched %s%s: %s, expected %s" % (e["file"], ptr, pointer(got["accept"], ptr), want))
    return fails


def load(corpus):
    files = {}
    for sub in ("ext", "gen"):
        for root, _, names in os.walk(os.path.join(corpus, sub)):
            for n in names:
                p = os.path.join(root, n)
                files[os.path.relpath(p, corpus).replace(os.sep, "/")] = open(p, "rb").read()
    doc = json.load(open(os.path.join(corpus, "vectors.json"), encoding="utf-8"))
    return doc, files


def mutations(doc, files):
    """(label, doc, files) copies, each of which must fail."""
    out = []

    def mutate(label, f):
        d, fs = copy.deepcopy(doc), dict(files)
        f(d, fs)
        out.append((label, d, fs))

    def vec(d, name, cap=None):
        return next(v for v in d["vectors"] if v["name"] == name and (cap is None or v["capability"] == cap))

    def rewrite(d, fs, name, cap, at, byte):
        v = vec(d, name, cap)
        b = bytearray(fs[v["file"]])
        b[at] = byte
        fs[v["file"]] = bytes(b)
        v["sha256"] = hashlib.sha256(fs[v["file"]]).hexdigest()

    mutate("the MHAS configuration's profile-level byte, its digest restated",
           lambda d, fs: rewrite(d, fs, "lc3_48k_5_1", None, 5, 0x0E))
    mutate("the H.266 vector's stated cropped width", lambda d, fs: vec(d, "vvc_main10_100x58.266")["outcome"]["accept"]["cropped"].__setitem__(0, 101))
    mutate("a refusal stated with another code", lambda d, fs: vec(d, "no_soc")["outcome"].__setitem__("refuse", "truncated"))
    mutate("a refused binding stated accepted", lambda d, fs: next(b for b in d["bindings"] if b["name"] == "layer_2_over_mvhevc").__setitem__("outcome", "accept"))
    mutate("a fact restated", lambda d, fs: d["facts"][0].__setitem__("value", [65, 64]))
    mutate("the JPEG XS codestream's Ppih, its digest restated",
           lambda d, fs: rewrite(d, fs, "jxs_main422_10_64x64.jxs", None, 16, 0x3A))
    mutate("a vector dropped from the list", lambda d, fs: d["vectors"].remove(vec(d, "erp_mono")))
    mutate("a probe's basis restated", lambda d, fs: d["probes"][2].__setitem__("basis", "implementation"))
    mutate("the LCEVC sequence configuration's level, its digest restated",
           lambda d, fs: rewrite(d, fs, "main_1080p", None, 7, 0x13))
    mutate("a projection descriptor's coverage stated otherwise",
           lambda d, fs: vec(d, "half_erp_side_by_side")["outcome"]["accept"]["coverage"].__setitem__(0, 360 << 16))
    mutate("the AVS3 audio header's rate index, its digest restated",
           lambda d, fs: rewrite(d, fs, "channels_5_1", None, 3, 0x10))

    # §14 to §17: a stated profile, level, version, bit rate or codec string changed, and a
    # header byte changed with its digest restated.
    def state(name, path, value):
        def f(d, fs):
            o = vec(d, name)["outcome"]["accept"]
            for k in path[:-1]:
                o = o[k]
            o[path[-1]] = value
        return f

    def flip(name, at, x):
        return lambda d, fs: rewrite(d, fs, name, None, at, fs[vec(d, name)["file"]][at] ^ x)

    for label, f in [
        ("the AV2 vector's stated profile", state("main420_10_1080p", ["profile"], 1)),
        ("the AV2 vector's stated level", state("main420_10_1080p", ["level"], "4.1")),
        ("an AV2 codec string, where AOMedia has released none", state("main420_10_1080p", ["codec"], "av02.0.08M.10")),
        ("the AV2 sequence header's last profile bit, its digest restated", flip("main420_10_1080p", 4, 0x04)),
        ("the AC-4 vector's stated codec string", state("v2_crc_24fps", ["codec"], "ac-4.02.01.04")),
        ("the AC-4 vector's stated presentation version", state("v2_crc_24fps", ["presentation", "version"], 2)),
        ("the AC-4 table of contents' frame rate index, its digest restated", flip("v0_single_44k", 5, 0x01)),
        ("the IAMF vector's stated codec string", state("opus_simple", ["codec"], "iamf.000.000.opus")),
        ("the IAMF vector's stated primary profile", state("opus_simple", ["primary_profile"], 1)),
        ("the IAMF sequence header's primary profile, its digest restated", flip("opus_simple", 6, 0x01)),
        ("the IVAS vector's stated codec string", state("stereo_32k_two_frames_cmr", ["codec"], "sevs")),
        ("the IVAS vector's stated bit rate", state("stereo_32k_two_frames_cmr", ["frames", 0, "bitrate"], 24400)),
        ("the IVAS ToC byte's bit rate index, its digest restated", flip("stereo_32k_two_frames_cmr", 1, 0x07)),
    ]:
        mutate(label, f)
    return out


def main(argv):
    corpus = argv[1]
    fetched = None
    if len(argv) == 4 and argv[2] == "--fetched":
        fetched = {n: open(os.path.join(argv[3], n), "rb").read() for n in os.listdir(argv[3])}
    elif len(argv) != 2:
        print(__doc__)
        return 2
    doc, files = load(corpus)
    fails = check(doc, files, fetched)
    for f in fails:
        print("FAIL", f)
    print("%d vectors, %d bindings, %d facts%s: %d failures" % (
        len(doc["vectors"]), len(doc["bindings"]), len(doc["facts"]),
        ", %d fetched" % len(doc["fetched"]) if fetched is not None else "", len(fails)))
    caught_all = True
    for label, d, fs in mutations(doc, files):
        try:
            caught = check(d, fs)
        except Exception as e:  # a crash on a mutated copy counts as a failure to read it
            caught = ["raised %s" % e]
        print(("PASS" if caught else "FAIL") + " mutated copy fails: %s%s" % (label, " (%s)" % caught[0][:90] if caught else ""))
        caught_all &= bool(caught)
    return 0 if not fails and caught_all else 1


if __name__ == "__main__":
    sys.exit(main(sys.argv))