"""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


# ---- 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,
}


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))
    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))