#!/usr/bin/env python3
"""The grain-free AV1 refinement corpus's second implementation (platform-decode §6).

    python3 -I verify.py [corpus dir] [--dav1d PATH]

Written from the text, sharing no code with wai-rs: it reads the AV1 framing of SPEC §5,
every OBU of every temporal unit and every sequence header (AV1 §5.3, §5.5) to the
film_grain_params_present flag; forms the codec string of platform-decode §2; checks each
frame of base.i420 against probe.json's digests (browser-sink §4.4's sample layout);
converts each frame to RGB8 by §6.3's JFIF integer conversion; and replays every case of
cases.json through the staged-delivery sink of ../../staged-conformance/staged_verify.py
(itself independent of wai-rs), a base made from the frames at the tier its decoder
gives (§6.2), and requires every report, unit, stage hash and clip hash. Then it requires
the mutations to fail: the grain flag of base.wai set alone (refused av1_film_grain), one
sample of base.i420 changed (other stage hashes, and a probe that no longer matches), and
a layer's base hash changed (the layer held for a base that never arrives).

With --dav1d PATH it also decodes base.wai's stream with the dav1d command-line tool
(its section 5 input) and requires base.i420 byte for byte. Standard library only.
"""

import hashlib
import json
import os
import struct
import subprocess
import sys
import tempfile

HERE = os.path.dirname(os.path.abspath(__file__))
WAI = os.path.abspath(os.path.join(HERE, "..", ".."))
sys.path.insert(0, os.path.join(WAI, "staged-conformance"))
import staged_verify as SV  # noqa: E402

AV1 = "wai.video.av1"


class Refused(Exception):
    def __init__(self, code, why):
        super().__init__(f"{code}: {why}")
        self.code = code


# ---------------------------------------------------------------- AV1 (SPEC §5, AV1 §5)


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

    def f(self, n):
        v = 0
        for _ in range(n):
            i = self.at >> 3
            if i >= len(self.b):
                raise Refused("av1_malformed", "a header runs past its OBU")
            v = (v << 1) | ((self.b[i] >> (7 - (self.at & 7))) & 1)
            self.at += 1
        return v

    def uvlc(self):
        lz = 0
        while self.f(1) == 0:
            lz += 1
            if lz > 32:
                raise Refused("av1_malformed", "uvlc past 32 leading zeros")
        return (1 << 32) - 1 if lz >= 32 else self.f(lz) + (1 << lz) - 1


def sequence_header(b):
    """sequence_header_obu() (AV1 §5.5) to film_grain_params_present."""
    r = Bits(b)
    profile = r.f(3)
    if profile > 2:
        raise Refused("av1_malformed", f"seq_profile {profile}")
    r.f(1)  # still_picture
    reduced = r.f(1)
    if reduced:
        level, tier = r.f(5), 0
    else:
        timing = r.f(1)
        model, delay_len = 0, 0
        if timing:
            r.f(32), r.f(32)
            if r.f(1):  # equal_picture_interval
                r.uvlc()
            model = r.f(1)
            if model:
                delay_len = r.f(5) + 1
                r.f(32), r.f(5), r.f(5)
        idd = r.f(1)
        level = tier = None
        for i in range(r.f(5) + 1):
            r.f(12)
            lv = r.f(5)
            t = r.f(1) if lv > 7 else 0
            if model and r.f(1):
                r.f(delay_len), r.f(delay_len), r.f(1)
            if idd and r.f(1):
                r.f(4)
            if i == 0:
                level, tier = lv, t
    wb, hb = r.f(4) + 1, r.f(4) + 1
    r.f(wb), r.f(hb)
    if not reduced and r.f(1):  # frame_id_numbers_present_flag
        r.f(4), r.f(3)
    r.f(3)
    if not reduced:
        r.f(4)
        oh = r.f(1)
        if oh:
            r.f(2)
        sct = 2 if r.f(1) else r.f(1)
        if sct > 0 and not r.f(1):
            r.f(1)
        if oh:
            r.f(3)
    r.f(3)
    # color_config() (AV1 §5.5.2)
    hbd = r.f(1)
    twelve = r.f(1) if profile == 2 and hbd else 0
    depth = 12 if twelve else (10 if hbd else 8)
    mono = 0 if profile == 1 else r.f(1)
    desc = (r.f(8), r.f(8), r.f(8)) if r.f(1) else None
    if mono:
        rng, sx, sy, csp = r.f(1), 1, 1, 0
    elif desc == (1, 13, 0):
        rng, sx, sy, csp = 1, 0, 0, 0
    else:
        rng = r.f(1)
        if profile == 0:
            sx, sy = 1, 1
        elif profile == 1:
            sx, sy = 0, 0
        elif depth == 12:
            sx = r.f(1)
            sy = r.f(1) if sx else 0
        else:
            sx, sy = 1, 0
        csp = r.f(2) if sx and sy else 0
    if not mono:
        r.f(1)  # separate_uv_delta_q
    grain = r.f(1)
    return {"profile": profile, "level": level, "tier": tier, "depth": depth, "mono": mono, "desc": desc,
            "range": rng, "sx": sx, "sy": sy, "csp": csp, "grain": grain}


def codec_string(s):
    """platform-decode §2: av01, profile, level and tier, depth; the rest unless default."""
    head = f"av01.{s['profile']}.{s['level']:02d}{'H' if s['tier'] else 'M'}.{s['depth']:02d}"
    cp, tc, mc = s["desc"] or (1, 1, 1)
    third = s["csp"] if s["sx"] and s["sy"] else 0
    tail = f".{s['mono']}.{s['sx']}{s['sy']}{third}.{cp:02d}.{tc:02d}.{mc:02d}.{s['range']}"
    return head if tail == ".0.110.01.01.01.0" else head + tail


def leb128(b, at):
    v = 0
    for i in range(8):
        if at + i >= len(b):
            raise Refused("av1_malformed", "an OBU size runs past its unit")
        v |= (b[at + i] & 0x7F) << (7 * i)
        if not b[at + i] & 0x80:
            return v, i + 1
    raise Refused("av1_malformed", "an OBU size longer than 8 bytes")


def read_stream(p):
    """SPEC §5's framing read whole: (n, h, w, units, sequence headers)."""
    if len(p) < 20:
        raise Refused("av1_malformed", "short header")
    n, h, w, _, _ = struct.unpack_from("<5I", p, 0)
    if n == 0 or h == 0 or w == 0:
        raise Refused("av1_malformed", "no frames, or a zero size")
    units, seqs, at = [], [], 20
    while at < len(p):
        if at + 4 > len(p):
            raise Refused("av1_malformed", "a unit length past the payload")
        (ln,) = struct.unpack_from("<I", p, at)
        tu = p[at + 4:at + 4 + ln]
        if ln == 0 or len(tu) != ln:
            raise Refused("av1_malformed", "an empty unit, or one past the payload")
        if len(units) >= n:
            raise Refused("av1_malformed", "more units than the header declares")
        o, has_seq = 0, False
        while o < len(tu):
            hd = tu[o]
            if hd & 0x80 or not hd & 0x02:
                raise Refused("av1_malformed", "forbidden bit, or no size field")
            hl = 2 if hd & 0x04 else 1
            size, sl = leb128(tu, o + hl)
            start = o + hl + sl
            if start + size > len(tu):
                raise Refused("av1_malformed", "an OBU past its unit")
            if (hd >> 3) & 0x0F == 1:
                seqs.append(sequence_header(tu[start:start + size]))
                has_seq = True
            o = start + size
        if not units and not has_seq:
            raise Refused("av1_malformed", "no sequence header in the first unit")
        units.append(tu)
        at += 4 + ln
    if len(units) != n:
        raise Refused("av1_malformed", f"{len(units)} units where the header declares {n}")
    return n, h, w, units, seqs


def read_base(p):
    """platform-decode §6.1, in its order."""
    n, h, w, units, seqs = read_stream(p)
    if any(s["grain"] for s in seqs):
        raise Refused("av1_film_grain", "film_grain_params_present 1")
    if any(s["depth"] != 8 or s["mono"] or not (s["sx"] and s["sy"]) for s in seqs):
        raise Refused("av1_format", "not 8-bit 4:2:0 with chroma")
    if w * h * 3 > 1 << 26 or n * w * h * 3 > 1 << 30:
        raise Refused("av1_too_large", "past the residual form's caps")
    return n, h, w, units, seqs


def i420_len(w, h):
    return w * h + 2 * ((w + 1) // 2) * ((h + 1) // 2)


def to_rgb8(f, w, h):
    """§6.3: JFIF in 16-bit fixed point, nearest chroma."""
    cw, ch = (w + 1) // 2, (h + 1) // 2
    y0, u0, v0 = 0, w * h, w * h + cw * ch
    out = bytearray(w * h * 3)
    k = 0
    for i in range(h):
        for j in range(w):
            y = f[y0 + i * w + j] << 16
            c = (i >> 1) * cw + (j >> 1)
            cb, cr = f[u0 + c] - 128, f[v0 + c] - 128
            for v in (y + 91881 * cr + 32768, y - 22554 * cb - 46802 * cr + 32768, y + 116130 * cb + 32768):
                out[k] = min(max(v >> 16, 0), 255)
                k += 1
    return bytes(out)


# ---------------------------------------------------------------- the cases


def make_base(env, frames, decoder):
    """A stage set's base (§6.2), or the refusal code."""
    m, _, payload = SV.read_wai1(env)
    mr = m["model_requirement"]
    if mr["capability"] != AV1 and mr.get("fallback") != AV1:
        return "not-offered"
    try:
        n, h, w, _, _ = read_base(payload)
    except Refused as e:
        return e.code
    fl = i420_len(w, h)
    if len(frames) != n * fl:
        return "av1_frames"
    units = [((h, w, 3), "u8", list(to_rgb8(frames[k * fl:(k + 1) * fl], w, h))) for k in range(n)]
    return {"hash": SV.content_hash(env), "capability": AV1, "via_fallback": mr["capability"] != AV1, "attested": False,
            "tier": SV.DE if decoder in ("reference", "probe-matched") else SV.NONE, "first_unit": 0, "units": units}


CFG = {"bound": {"max_envelopes": 16, "max_age_ms": 1000}, "cost_limit": None, "budget": 1 << 30,
       "record_budget": 1 << 20, "live": None}


def replay(case, files, frames):
    names = {SV.content_hash(b): f for f, b in files.items()}
    sink = SV.Sink(CFG)
    b = make_base(files[case["base"]], frames, case["decoder"])
    refusal = b if isinstance(b, str) else None
    if refusal is None:
        assert sink.add_base(b, b["units"]) == []
    steps = []
    for i, s in enumerate(case["steps"]):
        rs = sink.expire(int(s[len("expire@"):])) if s.startswith("expire@") else sink.offer(files[s], 1 + i, "none")
        steps.append({"step": s, "reports": [{"envelope": names.get(r["object"], r["object"]), "reason": r["envelope"],
                                              "held_for_base": r["held_for_base"], "applied": r["applied"],
                                              "held": r["held"], "refused": r["refused"]} for r in rs]})
    c = sink.bases.get(SV.content_hash(files[case["base"]]))
    units = None if c is None else [{"stage": s["stage"], "producer": names.get(s["producer"], s["producer"]),
                                     "tier": s["tier"], "sha256": c.unit_sha(u)} for u, s in enumerate(c.units)]
    return {"base_refusal": refusal, "steps": steps, "units": units, "clip_sha256": None if c is None else c.clip_sha()}


FILES = ["base.wai", "base-grain.wai", "layer1.wai", "layer2.wai", "layer1-over-grain.wai"]


def main(argv):
    d = HERE
    dav1d = None
    rest = list(argv)
    if "--dav1d" in rest:
        i = rest.index("--dav1d")
        dav1d = rest[i + 1]
        del rest[i:i + 2]
    if rest:
        d = rest[0]
    fails = []

    def check(ok, what):
        print(("ok    " if ok else "FAIL  ") + what)
        if not ok:
            fails.append(what)

    files = {f: open(os.path.join(d, f), "rb").read() for f in FILES}
    frames = open(os.path.join(d, "base.i420"), "rb").read()
    probe = json.load(open(os.path.join(d, "probe.json")))
    cases = json.load(open(os.path.join(d, "cases.json")))

    _, _, payload = SV.read_wai1(files["base.wai"])
    n, h, w, units, seqs = read_base(payload)
    fl = i420_len(w, h)
    check(probe["codec"] == codec_string(seqs[0]), f"the probe's codec string is the stream's ({codec_string(seqs[0])})")
    check((probe["codedWidth"], probe["codedHeight"], probe["frames"]) == (w, h, n), f"{n} frames of {w}x{h}")
    check(probe["film_grain_params_present"] is False and not any(s["grain"] for s in seqs), "every sequence header is grain-free")
    digests = [hashlib.sha256(frames[k * fl:(k + 1) * fl]).hexdigest() for k in range(n)]
    check(len(frames) == n * fl and digests == probe["frame_sha256"], "every frame of base.i420 has its probe digest")
    clip = hashlib.sha256(b"".join(bytes.fromhex(x) for x in digests)).hexdigest()
    check(clip == probe["clip_sha256"], "the probe's clip digest")
    rgb = [to_rgb8(frames[k * fl:(k + 1) * fl], w, h) for k in range(n)]
    check([hashlib.sha256(x).hexdigest() for x in rgb] == cases["stage0_sha256"], "every frame's RGB8 unit has its stage-0 hash")

    _, _, gp = SV.read_wai1(files["base-grain.wai"])
    try:
        read_base(gp)
        check(False, "base-grain.wai is refused as a base")
    except Refused as e:
        check(e.code == "av1_film_grain", f"base-grain.wai is refused as a base ({e.code})")

    for case in cases["cases"]:
        got = replay(case, files, frames)
        check(got == case["expect"], f"case {case['name']}")

    # Mutations, each of which must fail.
    seq_at = None
    tu = units[0]
    o = 0
    while o < len(tu):
        hl = 2 if tu[o] & 0x04 else 1
        size, sl = leb128(tu, o + hl)
        if (tu[o] >> 3) & 0x0F == 1:
            seq_at = (24 + o + hl + sl, size)
            break
        o += hl + sl + size
    flipped = None
    for bit in range(seq_at[1] * 8):
        p = bytearray(payload)
        p[seq_at[0] + bit // 8] ^= 0x80 >> (bit % 8)
        try:
            s = sequence_header(bytes(p[seq_at[0]:seq_at[0] + seq_at[1]]))
        except Refused:
            continue
        if s["grain"] and dict(s, grain=0) == seqs[0]:
            flipped = bytes(p)
            break
    try:
        read_base(flipped)
        check(False, "base.wai with only its grain flag set is refused")
    except Refused as e:
        check(e.code == "av1_film_grain", "base.wai with only its grain flag set is refused av1_film_grain")
    bad = bytearray(frames)
    bad[fl + 77] ^= 1
    bad = bytes(bad)
    check(hashlib.sha256(bad[fl:2 * fl]).hexdigest() != probe["frame_sha256"][1], "a changed sample fails the probe")
    two = next(c for c in cases["cases"] if c["name"] == "two-layers")
    got = replay(two, files, bad)
    check(got["units"][1]["sha256"] != two["expect"]["units"][1]["sha256"] and got["units"][0] == two["expect"]["units"][0],
          "a changed sample changes that frame's stage hash alone")
    check(make_base(files["base.wai"], frames[:-1], "reference") == "av1_frames", "frames short by one byte are refused")

    if dav1d:
        with tempfile.TemporaryDirectory() as t:
            obu, yuv = os.path.join(t, "s.obu"), os.path.join(t, "s.yuv")
            with open(obu, "wb") as fh:
                fh.write(b"".join(units))
            r = subprocess.run([dav1d, "-q", "--demuxer", "section5", "-i", obu, "-o", yuv], capture_output=True, timeout=120)
            out = open(yuv, "rb").read() if os.path.exists(yuv) else b""
            check(r.returncode == 0 and out == frames, f"dav1d's decode of base.wai is base.i420 ({len(out)} bytes)")

    print(f"{len(cases['cases'])} cases; {'FAIL' if fails else 'ok'}")
    return 1 if fails else 0


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