#!/usr/bin/env python3
"""Independent verifier for the integer payload corpus (extensions/integer-payloads.md).

Python standard library only; it shares no code with wai-rs. Written from the
extension's text, it:

  - parses every committed payload (WIV1, WIH1, WIA1, WIS1) with the rules and in the
    check order of the extension's §2-§6, and requires each to parse;
  - converts each payload's JSON source independently (§10) and requires the result to
    equal the committed payload byte for byte, and the payload's SHA-256 to be the one
    the index states;
  - decodes the WIV1 payloads with keyframe 0 in full - the rANS decoder with its
    bypass escape, its stop at the first exhausted symbol and its canonical end,
    whole-pel and half-pel motion compensation - and requires every frame's SHA-256
    and the clip's to be the ones the vector's sha256.json states;
  - parses every refusal fixture and requires a parse-layer fixture to be refused with
    its code, and a decode-layer fixture to parse; a decode-layer fixture of WIV1
    keyframe 0 is decoded and must be refused with its code (stream_exhausted,
    stream_noncanonical).

Not decoded independently: WIH1, WIV1 with keyframe 1 (its keyframe is a WIH1), WIA1
and WIS1, whose decodes run integer networks this script does not implement. Their
decoded outputs are checked by wai-rs against the same sidecars, on every leg of the
cross-architecture matrix for WIV1 and WIH1. The decode-layer fixtures of those
formats are checked here to parse, and by wai-rs for their codes.

    python3 int-payload-conformance/int_payload_verify.py int-payload-conformance

Exit status 1 on any failure.
"""
import base64
import hashlib
import json
import math
import os
import struct
import sys

MAGICS = (b"WIV1", b"WIH1", b"WIA1", b"WIS1")
MAX_TENSOR = 1 << 26
MAX_CLIP = 1 << 30
MAX_FRAMES = 1 << 20
MAX_MEDIAN_Q = 1 << 47
CAP_MAGIC = {
    "wai.video.int_motion": b"WIV1",
    "wai.video.int_hyper": b"WIV1",
    "wai.neural.int_hyper": b"WIH1",
    "wai.audio.int_codec": b"WIA1",
    "wai.splat.int_codec": b"WIS1",
}


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


class Reader:
    def __init__(self, b, p):
        self.b, self.p = b, p

    def left(self):
        return len(self.b) - self.p

    def take(self, n):
        if n > self.left():
            raise Refused("stream_length", "the payload ends inside a section")
        s = self.b[self.p:self.p + n]
        self.p += n
        return s

    def u8(self):
        return self.take(1)[0]

    def u16(self):
        return struct.unpack("<H", self.take(2))[0]

    def u32(self):
        return struct.unpack("<I", self.take(4))[0]

    def i32(self):
        return struct.unpack("<i", self.take(4))[0]

    def cdf(self, n):
        t = list(struct.unpack(f"<{n}I", self.take(4 * n)))
        check_cdf(t)
        return t

    def stream(self):
        n = self.u32()
        if n % 4 or n < 8:
            raise Refused("stream_length", "a rANS stream's length is not a multiple of 4, or is under 8")
        return self.take(n)

    def end(self):
        if self.left():
            raise Refused("trailing", "bytes follow the last section")


def check_cdf(t):
    if not 3 <= len(t) <= 4097:
        raise Refused("cdf_malformed", "length outside 3..=4097")
    if t[0] != 0 or t[-1] != 65536:
        raise Refused("cdf_malformed", "not from 0 to 65536")
    if any(a >= b for a, b in zip(t, t[1:])):
        raise Refused("cdf_malformed", "not strictly increasing")


def check_magic(p, want):
    m = p[:4]
    if len(m) == 4 and m == want:
        return
    if len(m) == 4 and m in MAGICS:
        raise Refused("format_capability_mismatch", "another integer format")
    raise Refused("bad_magic", "not the capability's magic")


def opening(p, magic, header):
    check_magic(p, magic)
    if len(p) < 5:
        raise Refused("stream_length", "no version byte")
    if p[4] != 1:
        raise Refused("version", "a version this revision does not define")
    if len(p) < header:
        raise Refused("stream_length", "the payload ends inside its fixed header")
    return Reader(p, 5)


def too_large(dims, cap):
    n = 1
    for d in dims:
        n *= d
    if n > cap:
        raise Refused("too_large", "over the cap")


def parse_wiv1(p, want_keyframe=None):
    r = opening(p, b"WIV1", 36)
    mvp, kf, res0 = r.u8(), r.u8(), r.u8()
    n_frames, h, w, block, q = r.u32(), r.u32(), r.u32(), r.u32(), r.u32()
    offset, cdf_len, res1 = r.i32(), r.u16(), r.u16()
    if res0 or res1:
        raise Refused("reserved_nonzero", "reserved")
    if mvp not in (1, 2):
        raise Refused("field_range", "mv_precision")
    if kf > 1:
        raise Refused("field_range", "keyframe")
    if 0 in (n_frames, h, w):
        raise Refused("field_range", "n_frames, h or w")
    if block == 0:
        raise Refused("field_range", "block")
    if not 1 <= q <= 65535:
        raise Refused("field_range", "q")
    if not 3 <= cdf_len <= 4097:
        raise Refused("field_range", "cdf_len")
    if n_frames == 1 and (mvp != 1 or block != 1):
        raise Refused("field_range", "a one-frame clip has mv_precision 1 and block 1")
    if want_keyframe is not None and want_keyframe != kf:
        raise Refused("keyframe_capability_mismatch", "keyframe")
    too_large([h, w, 3], MAX_TENSOR)
    too_large([n_frames, h, w, 3], MAX_CLIP)
    too_large([n_frames], MAX_FRAMES)
    nb = -(-h // block) * -(-w // block)
    too_large([2, n_frames - 1, nb], MAX_TENSOR)
    cdf = r.cdf(cdf_len)
    raw = r.take((n_frames - 1) * nb * 4)
    mvs = [struct.unpack_from("<hh", raw, 4 * i) for i in range((n_frames - 1) * nb)]
    rans = r.stream()
    wih1 = None
    if kf == 1:
        k = r.u32()
        if k != r.left():
            raise Refused("nested_length", "K is not the rest")
        wih1 = r.take(k)
        parse_wih1(wih1)
    else:
        r.end()
    return dict(mv_precision=mvp, keyframe=kf, n_frames=n_frames, h=h, w=w, block=block, q=q,
                offset=offset, cdf=cdf, mvs=mvs, rans=rans, wih1=wih1)


def parse_wih1(p):
    r = opening(p, b"WIH1", 32)
    res = r.take(3)
    m, sh, sw, n_z, zh, zw = (r.u32() for _ in range(6))
    if any(res):
        raise Refused("reserved_nonzero", "reserved")
    if 0 in (m, sh, sw, n_z, zh, zw):
        raise Refused("field_range", "a dimension")
    too_large([n_z, zh, zw], MAX_TENSOR)
    too_large([2, m, sh, sw], MAX_TENSOR)
    z, y = r.stream(), r.stream()
    r.end()
    return dict(m=m, sh=sh, sw=sw, n_z=n_z, zh=zh, zw=zw, z_rans=z, rans=y)


def parse_wia1(p):
    r = opening(p, b"WIA1", 32)
    afrac, res0 = r.u8(), r.u16()
    sr, n_ch, lat_len, t_out = r.u32(), r.u32(), r.u32(), r.u32()
    offset, cdf_len, res1 = r.i32(), r.u16(), r.u16()
    if res0 or res1:
        raise Refused("reserved_nonzero", "reserved")
    if not 1 <= afrac <= 30:
        raise Refused("field_range", "afrac")
    if 0 in (sr, n_ch, lat_len, t_out):
        raise Refused("field_range", "a dimension")
    if not 3 <= cdf_len <= 4097:
        raise Refused("field_range", "cdf_len")
    too_large([n_ch, lat_len], MAX_TENSOR)
    too_large([t_out], MAX_TENSOR)
    med = list(struct.unpack(f"<{n_ch}q", r.take(8 * n_ch)))
    if any(abs(x) > MAX_MEDIAN_Q for x in med):
        raise Refused("field_range", "median")
    cdfs = [r.cdf(cdf_len) for _ in range(n_ch)]
    rans = r.stream()
    r.end()
    return dict(afrac=afrac, sr=sr, n_ch=n_ch, lat_len=lat_len, t_out=t_out, offset=offset,
                median_q=med, cdfs=cdfs, rans=rans)


def parse_wis1(p):
    r = opening(p, b"WIS1", 32)
    afrac, res = r.u8(), r.u16()
    c, sh, sw, s, attr, n_g = (r.u32() for _ in range(6))
    if res:
        raise Refused("reserved_nonzero", "reserved")
    if not 1 <= afrac <= 30:
        raise Refused("field_range", "afrac")
    if 0 in (c, sh, sw, s, attr, n_g):
        raise Refused("field_range", "a dimension")
    if n_g > s * s:
        raise Refused("field_range", "n_gaussians")
    too_large([c, sh, sw], MAX_TENSOR)
    too_large([attr, s, s], MAX_TENSOR)
    raw = r.take(8 * attr)
    ranges = [struct.unpack_from("<ff", raw, 8 * i) for i in range(attr)]
    for lo, hi in ranges:
        if not (math.isfinite(lo) and math.isfinite(hi)) or lo > hi:
            raise Refused("field_range", "a range")
    lens = list(struct.unpack(f"<{c}H", r.take(2 * c)))
    if any(not 3 <= n <= 4097 for n in lens):
        raise Refused("field_range", "a CDF length")
    offs = list(struct.unpack(f"<{c}i", r.take(4 * c)))
    med = list(struct.unpack(f"<{c}q", r.take(8 * c)))
    if any(abs(x) > MAX_MEDIAN_Q for x in med):
        raise Refused("field_range", "median")
    cdfs = [r.cdf(n) for n in lens]
    rans = r.stream()
    r.end()
    return dict(afrac=afrac, c=c, sh=sh, sw=sw, s=s, attr=attr, n_gaussians=n_g, ranges=ranges,
                cdfs=cdfs, offsets=offs, median_q=med, rans=rans)


def parse(cap, p):
    m = CAP_MAGIC[cap]
    if m == b"WIV1":
        return parse_wiv1(p, 1 if cap == "wai.video.int_hyper" else 0)
    return {b"WIH1": parse_wih1, b"WIA1": parse_wia1, b"WIS1": parse_wis1}[m](p)


# ---- writers: the JSON sources converted independently (extensions/integer-payloads.md §10) ----

def put_stream(b):
    return struct.pack("<I", len(b)) + b


class NotConvertible(Exception):
    """A JSON source the converters of §10 refuse. Raised explicitly, never through
    `assert`, so `python3 -O` checks the same rules."""


def require(ok, why):
    if not ok:
        raise NotConvertible(why)


def wiv1_from_json(j, keyframe=None):
    nf, h, w, block = j["n_frames"], j["h"], j["w"], j["block"]
    mvp = j.get("mv_precision", 1)
    require(mvp in (1, 2), "mv_precision is neither 1 nor 2")
    if nf == 1:
        # A one-frame clip's header fixes mv_precision 1 and block 1 (§3, §10).
        mvp, block = 1, 1
    nb = -(-h // block) * -(-w // block)
    mvs = j["mvs"]
    require(len(mvs) == nf and not mvs[0], "mvs is not one list per frame, the first empty")
    out = b"".join([b"WIV1", bytes([1, mvp, 1 if keyframe else 0, 0]),
                    struct.pack("<IIIIIiHH", nf, h, w, block, j["q"], j["offset"], len(j["cdf"]), 0),
                    struct.pack(f"<{len(j['cdf'])}I", *j["cdf"])])
    for fr in mvs[1:]:
        require(isinstance(fr, list), "a frame's motion vectors are not a list")
        require(len(fr) <= nb, "a frame carries more motion vectors than blocks")
        for dy, dx in list(fr) + [[0, 0]] * (nb - len(fr)):
            out += struct.pack("<hh", dy, dx)
    out += put_stream(base64.b64decode(j["rans_b64"]))
    if keyframe:
        out += put_stream(keyframe)
    return out


def wih1_from_json(j):
    return (b"WIH1" + bytes([1, 0, 0, 0]) + struct.pack("<6I", j["m"], j["sh"], j["sw"], j["n_z"], j["zh"], j["zw"])
            + put_stream(bytes(j["z_rans"])) + put_stream(bytes(j["rans"])))


def wia1_from_json(j, sr):
    n = len(j["cdfs"][0])
    require(all(len(t) == n for t in j["cdfs"]), "the channels' tables are not of one length")
    return (b"WIA1" + bytes([1, j["afrac"], 0, 0])
            + struct.pack("<IIIIiHH", sr, j["n_ch"], j["lat_len"], j["t_out"], j["lo"], n, 0)
            + struct.pack(f"<{j['n_ch']}q", *j["median_q"])
            + b"".join(struct.pack(f"<{n}I", *t) for t in j["cdfs"])
            + put_stream(base64.b64decode(j["rans_b64"])))


def f32(x):
    """`x` as IEEE 754 binary32, refused unless it is exactly a binary32 value (§10)."""
    require(isinstance(x, (int, float)) and math.isfinite(x), f"{x!r} is not a finite number")
    try:
        b = struct.pack("<f", x)
    except OverflowError:
        raise NotConvertible(f"{x!r} is not a binary32 value")
    require(struct.unpack("<f", b)[0] == x, f"{x!r} is not a binary32 value")
    return b


def wis1_from_json(j):
    c = j["c"]
    return (b"WIS1" + bytes([1, j["afrac"], 0, 0])
            + struct.pack("<6I", c, j["sh"], j["sw"], j["s"], j["attr"], j["n_gaussians"])
            + b"".join(f32(lo) + f32(hi) for lo, hi in zip(j["lo"], j["hi"]))
            + struct.pack(f"<{c}H", *(len(t) for t in j["cdfs"]))
            + struct.pack(f"<{c}i", *j["offsets"])
            + struct.pack(f"<{c}q", *j["median_q"])
            + b"".join(struct.pack(f"<{len(t)}I", *t) for t in j["cdfs"])
            + put_stream(base64.b64decode(j["rans_b64"])))


# ---- the rANS decoder (precision 16, bypass 4) and the WIV1 keyframe-0 decode ----

I32_MIN, I32_MAX = -(1 << 31), (1 << 31) - 1


def sat(v):
    return max(I32_MIN, min(I32_MAX, v))


class Rans:
    def __init__(self, b):
        self.w = [struct.unpack_from("<I", b, i)[0] for i in range(0, len(b), 4)]
        self.i, self.x, self.exhausted = 2, (self.w[1] << 32) | self.w[0], False

    def renorm(self):
        if self.x < (1 << 31):
            if self.i >= len(self.w):
                self.exhausted = True
                return
            self.x = (self.x << 32) | self.w[self.i]
            self.i += 1

    def bits(self, n):
        v = self.x & ((1 << n) - 1)
        self.x >>= n
        self.renorm()
        return v

    def decode(self, cdf, offset):
        n = len(cdf)
        max_value = n - 2
        cum = self.x & 0xFFFF
        lo, hi = 1, n - 1
        while lo < hi:
            mid = (lo + hi) // 2
            if cdf[mid] > cum:
                hi = mid
            else:
                lo = mid + 1
        s = lo - 1
        self.x = (cdf[s + 1] - cdf[s]) * (self.x >> 16) + (self.x & 0xFFFF) - cdf[s]
        self.renorm()
        value = s
        if value == max_value:
            val = self.bits(4)
            count, groups = val, 0
            while val == 15 and groups < 8 and not self.exhausted:
                val = self.bits(4)
                count += val
                groups += 1
            raw = 0
            for j in range(min(count, 8)):
                raw |= self.bits(4) << (4 * j)
            value = raw >> 1
            value = sat(-value - 1) if raw & 1 else sat(value + max_value)
        return sat(value + offset)


def decode_wiv1_k0(v):
    h, w, block, q, prec = v["h"], v["w"], v["block"], v["q"], v["mv_precision"]
    nbx = -(-w // block)
    nb = nbx * -(-h // block)
    d = Rans(v["rans"])
    frames, prev = [], None
    for f in range(v["n_frames"]):
        mvs = v["mvs"][(f - 1) * nb:f * nb] if f else None
        # The frame grows a sample at a time as its symbols decode: nothing is
        # allocated ahead of the stream that fills it.
        cur = bytearray()
        for y in range(h):
            for x in range(w):
                dy, dx = mvs[(y // block) * nbx + x // block] if f else (0, 0)
                for c in range(3):
                    if f == 0:
                        pred = 128
                    else:
                        def p(yy, xx):
                            return prev[(min(max(yy, 0), h - 1) * w + min(max(xx, 0), w - 1)) * 3 + c]
                        if prec == 1:
                            pred = p(y + dy, x + dx)
                        else:
                            iy, fy, ix, fx = y + (dy >> 1), dy & 1, x + (dx >> 1), dx & 1
                            if not fy and not fx:
                                pred = p(iy, ix)
                            elif not fy:
                                pred = (p(iy, ix) + p(iy, ix + 1) + 1) >> 1
                            elif not fx:
                                pred = (p(iy, ix) + p(iy + 1, ix) + 1) >> 1
                            else:
                                pred = (p(iy, ix) + p(iy, ix + 1) + p(iy + 1, ix) + p(iy + 1, ix + 1) + 2) >> 2
                    r = d.decode(v["cdf"], v["offset"])
                    if d.exhausted:
                        # §2: the first symbol during whose decode the stream runs out.
                        raise Refused("stream_exhausted", "the residual stream ran out")
                    cur.append(min(max(pred + r * q, 0), 255))
        frames.append(bytes(cur))
        prev = cur
    if d.i != len(d.w) or d.x != 1 << 31:
        # §2: after the last symbol every word has been read and the state is 2^31.
        raise Refused("stream_noncanonical", "the residual stream does not end canonically")
    return frames


def main(cdir):
    wai = os.path.normpath(os.path.join(cdir, ".."))

    def rd(rel):
        with open(os.path.join(wai, rel), "rb") as f:
            return f.read()

    def jl(rel):
        return json.loads(rd(rel))

    idx = jl("int-payload-conformance/vectors.json")
    fails, decoded, parsed_only = [], 0, []
    for v in idx["vectors"]:
        name, cap = v["name"], v["capability"]
        p = rd(v["payload"])
        try:
            fields = parse(cap, p)
        except Refused as e:
            fails.append(f"{name}: does not parse: {e}")
            continue
        if hashlib.sha256(p).hexdigest() != v["payload_sha256"]:
            fails.append(f"{name}: payload SHA-256 differs from the index")
        src = v["converted_from"]
        j = jl(src["json"])
        m = CAP_MAGIC[cap]
        try:
            if m == b"WIV1":
                conv = wiv1_from_json(j, rd(src["keyframe"]) if "keyframe" in src else None)
            elif m == b"WIH1":
                conv = wih1_from_json(j)
            elif m == b"WIA1":
                conv = wia1_from_json(j, jl(src["sample_rate"])["sample_rate"])
            else:
                conv = wis1_from_json(j)
        except NotConvertible as e:
            fails.append(f"{name}: {src['json']} is not convertible: {e}")
            continue
        if conv != p:
            fails.append(f"{name}: the independent conversion of {src['json']} differs from the payload")
        if m == b"WIV1" and fields["keyframe"] == 0:
            frames = decode_wiv1_k0(fields)
            side = jl(v["sidecar"])[name]
            digests = [hashlib.sha256(f).digest() for f in frames]
            if [d.hex() for d in digests] != side["frames_sha256"] or hashlib.sha256(b"".join(digests)).hexdigest() != side["clip_sha256"]:
                fails.append(f"{name}: the independent decode differs from {v['sidecar']}")
            else:
                decoded += 1
                print(f"ok   {name}: parsed, converted, decoded to its sidecar's SHA-256 ({len(frames)} frames)")
        else:
            parsed_only.append(name)
            print(f"ok   {name}: parsed and converted (not decoded independently)")
    checked, decode_layer = 0, []
    for f in idx["refusals"]:
        name, cap, code, layer = f["name"], f["capability"], f["code"], f["layer"]
        p = rd(f["payload"])
        try:
            fields = parse(cap, p)
            got = None
        except Refused as e:
            got = e.code
        if layer == "parse":
            if got != code:
                fails.append(f"{name}: parsed with {got}, want {code}")
                continue
        else:
            if got is not None:
                fails.append(f"{name}: refused at parse with {got}, want a decode-layer {code}")
                continue
            if CAP_MAGIC[cap] == b"WIV1" and fields["keyframe"] == 0:
                try:
                    decode_wiv1_k0(fields)
                    fails.append(f"{name}: decoded, want {code}")
                    continue
                except Refused as e:
                    if e.code != code:
                        fails.append(f"{name}: decode refused with {e.code}, want {code}")
                        continue
            else:
                decode_layer.append(name)
        checked += 1
        print(f"ok   {name}: {code}" + ("" if layer == "parse" or name not in decode_layer else " (parses; the code is the decode's)"))
    # wiv1_end_state codes three symbols from the state 2^31 + 1, as the words [1, 4]; its
    # canonical twin, the words [0, 4], decodes to one grey pixel.
    twin = bytearray(rd("int-payload-conformance/refusals/wiv1_end_state.bin"))
    twin[-8:-4] = struct.pack("<I", 0)
    try:
        if decode_wiv1_k0(parse("wai.video.int_motion", bytes(twin))) != [bytes([128, 128, 128])]:
            fails.append("wiv1_end_state's canonical twin does not decode to one grey pixel")
        else:
            print("ok   wiv1_end_state's canonical twin decodes to one grey pixel")
    except Refused as e:
        fails.append(f"wiv1_end_state's canonical twin is refused: {e}")
    print(f"vectors: {decoded} decoded independently; parsed and converted only: {', '.join(parsed_only)}")
    print(f"refusal fixtures checked: {checked}; decode-layer codes left to wai-rs: {', '.join(decode_layer)}")
    for x in fails:
        print("FAIL", x)
    return 1 if fails else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1] if len(sys.argv) > 1 else os.path.dirname(os.path.abspath(__file__))))