#!/usr/bin/env python3
"""Independent verifier of the prior-forms corpus (extensions/prior-carriage.md §14).

Written from the extension's text: §9 (the safetensors and GGUF layouts), §10 (the
integer-quantised profile and its bindings), §11 (the canonical forms), §12 (the NNC
subset) and Appendix A (the integer payload's arithmetic coding). It shares no code with
wai-rs, which generates the corpus, and reads §13's table of codes from the extension.

For every vector it requires the outcome the corpus states: for an encode, the file's
bytes or the code; for a load, the binding and the native artifact's bytes or the code.
Then it flips one byte of each accepted file, in several places, and requires every copy
to be refused or to load to other bytes.

    python3 -I prior_forms_verify.py <live-conformance/prior-forms>

Exit status 0 only if every vector verifies and every mutated copy is caught.
"""
import hashlib
import json
import os
import re
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
EXTENSION = os.path.join(HERE, "..", "extensions", "prior-carriage.md")


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


def sha(b):
    return hashlib.sha256(b).hexdigest()


# ------------------------------------------------------------------ §13 codes

def codes():
    text = open(EXTENSION, encoding="utf-8").read()
    sec = text.split("\n## 13.")[1].split("\n## ")[0]
    return [m.group(1) for m in re.finditer(r"^\| `([a-z_]+)` \|", sec, re.M)]


# ------------------------------------------------------------------ strict JSON (§9.1 step 3)

def strict_json(b):
    """RFC 8259 over UTF-8, no member named twice, at most 64 levels; numbers kept as spelt."""
    try:
        text = b.decode("utf-8")
    except UnicodeDecodeError:
        raise Refuse("header_malformed", "not UTF-8")

    def pairs(ps):
        seen = set()
        for k, _ in ps:
            if k in seen:
                raise Refuse("header_malformed", f"member {k!r} twice")
            seen.add(k)
        return ("obj", ps)

    def bad_const(c):
        raise Refuse("header_malformed", f"{c} is not JSON")

    try:
        v = json.loads(text, object_pairs_hook=pairs, parse_int=lambda s: ("num", s),
                       parse_float=lambda s: ("num", s), parse_constant=bad_const)
    except Refuse:
        raise
    except (ValueError, RecursionError) as e:
        raise Refuse("header_malformed", str(e))

    def walk(x, d):
        if d > 64:
            raise Refuse("header_malformed", "nested too deeply")
        if isinstance(x, str):
            try:
                x.encode("utf-8")
            except UnicodeEncodeError:
                raise Refuse("header_malformed", "an unpaired surrogate")
        elif isinstance(x, list):
            for y in x:
                walk(y, d + 1)
        elif isinstance(x, tuple) and x[0] == "obj":
            for k, y in x[1]:
                walk(k, d)
                walk(y, d + 1)

    walk(v, 1)
    return v


def plain_uint(x):
    """An integer in plain decimal, at most 2^64 - 1 (§9.1 step 5), or None."""
    if not (isinstance(x, tuple) and x[0] == "num"):
        return None
    s = x[1]
    if not re.fullmatch(r"0|[1-9][0-9]*", s):
        return None
    v = int(s)
    return v if v < 1 << 64 else None


def jquote(s):
    """A JSON string as RFC 8785 §3.2.2.2 escapes it."""
    return json.dumps(s, ensure_ascii=False)


# ------------------------------------------------------------------ element types (§9.1, §9.2)

ST_SIZE = {"BOOL": 1, "U8": 1, "I8": 1, "U16": 2, "I16": 2, "F16": 2, "BF16": 2, "U32": 4, "I32": 4,
           "F32": 4, "U64": 8, "I64": 8, "F64": 8, "F8_E5M2": 1, "F8_E4M3": 1}
GG_TYPES = {0: "F32", 1: "F16", 8: "Q8_0", 24: "I8", 25: "I16", 26: "I32", 27: "I64", 28: "F64", 30: "BF16"}
GG_NUMBER = {v: k for k, v in GG_TYPES.items()}
INT_WIDTH = {"I8": 1, "I16": 2, "I32": 4, "I64": 8}


def nbytes(dtype, n):
    if dtype == "Q8_0":
        return None if n % 32 else n // 32 * 34
    return n * ST_SIZE[dtype]


def product(shape):
    p = 1
    for d in shape:
        p *= d
    return p


class TSet:
    """Metadata (str -> str) and tensors name -> (dtype, shape, bytes)."""

    def __init__(self):
        self.meta = {}
        self.t = {}


# ------------------------------------------------------------------ §9.1 safetensors

def st_read(b):
    if len(b) < 8:
        raise Refuse("header_truncated", "under 8 bytes")
    n = struct.unpack_from("<Q", b, 0)[0]
    if n > 100_000_000:
        raise Refuse("header_too_large", str(n))
    if 8 + n > len(b):
        raise Refuse("header_truncated", "header past the end")
    doc = strict_json(b[8:8 + n])
    if not (isinstance(doc, tuple) and doc[0] == "obj"):
        raise Refuse("header_malformed", "not an object")
    data = b[8 + n:]
    s = TSet()
    spans = []
    for name, v in doc[1]:
        if name == "__metadata__":
            if not (isinstance(v, tuple) and v[0] == "obj"):
                raise Refuse("header_malformed", "__metadata__")
            for k, mv in v[1]:
                if not isinstance(mv, str):
                    raise Refuse("header_malformed", "metadata not a string")
                s.meta[k] = mv
            continue
        if name == "":
            raise Refuse("tensor_malformed", "empty name")
        if not (isinstance(v, tuple) and v[0] == "obj") or sorted(k for k, _ in v[1]) != ["data_offsets", "dtype", "shape"]:
            raise Refuse("tensor_malformed", name)
        m = dict(v[1])
        if not isinstance(m["dtype"], str):
            raise Refuse("tensor_malformed", "dtype not a string")
        if m["dtype"] not in ST_SIZE:
            raise Refuse("dtype_unknown", m["dtype"])
        if not isinstance(m["shape"], list) or any(plain_uint(d) is None for d in m["shape"]):
            raise Refuse("tensor_malformed", "shape")
        shape = [plain_uint(d) for d in m["shape"]]
        offs = m["data_offsets"]
        if not isinstance(offs, list) or len(offs) != 2 or any(plain_uint(o) is None for o in offs):
            raise Refuse("tensor_malformed", "data_offsets")
        b0, e0 = plain_uint(offs[0]), plain_uint(offs[1])
        if b0 > e0:
            raise Refuse("tensor_malformed", "begin past end")
        cnt = product(shape)
        if cnt >= 1 << 64 or nbytes(m["dtype"], cnt) >= 1 << 64:
            raise Refuse("tensor_malformed", "count overflows")
        if e0 - b0 != nbytes(m["dtype"], cnt):
            raise Refuse("offset_mismatch", name)
        spans.append((b0, e0, name))
        s.t[name] = [m["dtype"], shape, None]
    at = 0
    for b0, e0, name in sorted(spans):
        if b0 != at:
            raise Refuse("offset_mismatch", f"{name} at {b0}, data so far {at}")
        at = e0
    if at != len(data):
        raise Refuse("offset_mismatch", "uncovered data")
    for b0, e0, name in spans:
        s.t[name][2] = data[b0:e0]
    return s


def st_write(s):
    """§11.1."""
    parts = []
    if s.meta:
        parts.append('"__metadata__":{' + ",".join(f"{jquote(k)}:{jquote(s.meta[k])}" for k in sorted(s.meta, key=lambda x: x.encode())) + "}")
    at = 0
    names = sorted(s.t, key=lambda x: x.encode())
    for name in names:
        dt, shape, data = s.t[name]
        parts.append(f'{jquote(name)}:{{"dtype":"{dt}","shape":[{",".join(str(d) for d in shape)}],"data_offsets":[{at},{at + len(data)}]}}')
        at += len(data)
    h = ("{" + ",".join(parts) + "}").encode()
    while (8 + len(h)) % 8:
        h += b" "
    return struct.pack("<Q", len(h)) + h + b"".join(s.t[n][2] for n in names)


# ------------------------------------------------------------------ §9.2 GGUF

FIXED = {0: 1, 1: 1, 7: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 10: 8, 11: 8, 12: 8}


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

    def take(self, n):
        if self.p + n > len(self.b):
            raise Refuse("header_truncated", f"{n} bytes at {self.p}")
        s = self.b[self.p:self.p + n]
        self.p += n
        return s

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

    def u64(self):
        return struct.unpack("<Q", self.take(8))[0]

    def string(self):
        raw = self.take(self.u64())
        try:
            return raw.decode("utf-8")
        except UnicodeDecodeError:
            raise Refuse("header_malformed", "a string that is not UTF-8")


def gg_value(c, t):
    if t in FIXED:
        v = c.take(FIXED[t])
        if t == 7 and v[0] > 1:
            raise Refuse("header_malformed", "boolean")
        return None
    if t == 8:
        return c.string()
    if t == 9:
        et = c.u32()
        if et == 9:
            raise Refuse("header_malformed", "array of arrays")
        if et != 8 and et not in FIXED:
            raise Refuse("header_malformed", f"array of {et}")
        n = c.u64()
        if n * FIXED.get(et, 8) > len(c.b) - c.p:
            raise Refuse("header_truncated", "array past the end")
        for _ in range(n):
            gg_value(c, et)
        return None
    raise Refuse("header_malformed", f"value type {t}")


KEY = re.compile(r"[a-z0-9_]+(\.[a-z0-9_]+)*")


def gg_read(b):
    if len(b) < 4:
        raise Refuse("header_truncated", "under 4 bytes")
    if b[:4] != b"GGUF":
        raise Refuse("bad_magic")
    c = Cur(b)
    c.p = 4
    if c.u32() != 3:
        raise Refuse("version_unsupported")
    nt, nkv = c.u64(), c.u64()
    if nkv * 12 > len(b) or nt * 28 > len(b):
        raise Refuse("header_truncated", "counts")
    s = TSet()
    keys = set()
    align = 32
    for _ in range(nkv):
        k = c.string()
        if not (len(k.encode()) <= 65535 and KEY.fullmatch(k)):
            raise Refuse("header_malformed", f"key {k!r}")
        if k in keys:
            raise Refuse("header_malformed", f"key {k!r} twice")
        keys.add(k)
        t = c.u32()
        at = c.p
        v = gg_value(c, t)
        if k == "general.alignment":
            a = struct.unpack_from("<I", b, at)[0] if t == 4 else 0
            if t != 4 or a == 0 or a % 8:
                raise Refuse("header_malformed", "general.alignment")
            align = a
        if v is not None:
            s.meta[k] = v
    infos = []
    for _ in range(nt):
        name = c.string()
        if not 1 <= len(name.encode()) <= 64:
            raise Refuse("tensor_malformed", "name length")
        if name in s.t:
            raise Refuse("tensor_malformed", f"{name} twice")
        nd = c.u32()
        if not 1 <= nd <= 4:
            raise Refuse("tensor_malformed", "dimensions")
        dims = [c.u64() for _ in range(nd)]
        if 0 in dims:
            raise Refuse("tensor_malformed", "a dimension of 0")
        ty = c.u32()
        if ty not in GG_TYPES:
            raise Refuse("dtype_unknown", str(ty))
        off = c.u64()
        if off % align:
            raise Refuse("offset_mismatch", "unaligned offset")
        dt = GG_TYPES[ty]
        if dt == "Q8_0" and dims[0] % 32:
            raise Refuse("tensor_malformed", "Q8_0 rows")
        cnt = product(dims)
        if cnt >= 1 << 64 or nbytes(dt, cnt) >= 1 << 64:
            raise Refuse("tensor_malformed", "count")
        infos.append((off, nbytes(dt, cnt), name))
        s.t[name] = [dt, list(reversed(dims)), None]
    hend = c.p
    dstart = -(-hend // align) * align
    if dstart > len(b):
        raise Refuse("header_truncated", "data begins past the end")
    if any(b[hend:dstart]):
        raise Refuse("header_malformed", "padding")
    data = b[dstart:]
    at = 0
    for off, ln, name in sorted(infos):
        if off < at:
            raise Refuse("offset_mismatch", "overlap")
        if off + ln > len(data):
            raise Refuse("offset_mismatch", "past the end")
        s.t[name][2] = data[off:off + ln]
        at = off + ln
    return s


def gg_write(s):
    """§11.2."""
    def st(x):
        x = x.encode()
        return struct.pack("<Q", len(x)) + x
    names = sorted(s.t, key=lambda x: x.encode())
    o = b"GGUF" + struct.pack("<IQQ", 3, len(names), len(s.meta))
    for k in sorted(s.meta, key=lambda x: x.encode()):
        o += st(k) + struct.pack("<I", 8) + st(s.meta[k])
    at = 0
    offs = []
    for n in names:
        dt, shape, data = s.t[n]
        at = -(-at // 32) * 32
        offs.append(at)
        o += st(n) + struct.pack("<I", len(shape)) + b"".join(struct.pack("<Q", d) for d in reversed(shape))
        o += struct.pack("<IQ", GG_NUMBER[dt], at)
        at += len(data)
    o += b"\0" * ((-len(o)) % 32)
    base = len(o)
    o = bytearray(o)
    for n, off in zip(names, offs):
        o += b"\0" * (base + off - len(o))
        o += s.t[n][2]
    return bytes(o)


# ------------------------------------------------------------------ Appendix A

T = [
    128, 112, 97, 84, 74, 65, 57, 50, 45, 39, 34, 30, 27, 23, 20, 18, 15, 14, 12, 11, 10, 9, 7, 7, 5, 5, 4, 4, 3, 3, 2, 2,
    142, 125, 108, 93, 82, 72, 63, 56, 50, 43, 38, 33, 30, 26, 22, 20, 17, 16, 13, 12, 11, 10, 8, 8, 6, 6, 5, 5, 3, 3, 2, 2,
    156, 137, 119, 103, 90, 79, 70, 61, 55, 48, 42, 37, 33, 28, 24, 22, 19, 17, 15, 13, 12, 11, 9, 9, 6, 6, 5, 5, 4, 4, 2, 2,
    171, 150, 130, 112, 99, 87, 76, 67, 60, 52, 46, 40, 36, 31, 27, 24, 21, 19, 16, 15, 13, 12, 10, 10, 7, 7, 6, 6, 4, 4, 3, 3,
    185, 162, 141, 121, 107, 94, 82, 73, 65, 56, 50, 43, 39, 34, 29, 26, 22, 21, 17, 16, 14, 13, 11, 11, 8, 8, 6, 6, 4, 4, 3, 3,
    199, 175, 152, 131, 115, 101, 89, 78, 70, 61, 54, 47, 42, 36, 31, 28, 24, 22, 19, 17, 15, 14, 12, 12, 8, 8, 7, 7, 5, 5, 3, 3,
    213, 187, 163, 140, 123, 108, 95, 84, 75, 65, 58, 50, 45, 39, 33, 30, 26, 24, 20, 18, 16, 15, 13, 13, 9, 9, 7, 7, 5, 5, 3, 3,
    228, 200, 174, 150, 132, 116, 102, 90, 80, 70, 62, 54, 48, 42, 36, 32, 28, 26, 22, 20, 18, 16, 14, 14, 10, 10, 8, 8, 6, 6, 4, 4,
]
C = [2512, 2288, 2064, 1840, 1616, 1392, 1168, 944, 720, 560, 464, 368, 272, 208, 144, 80,
     64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 0]
INIT = [[1, 4, 0, 0], [1, 4, -1310, -41], [1, 4, 3039, 95], [0, 5, 0, 0], [2, 6, 962, 30],
        [2, 6, 3039, 95], [2, 6, -673, -21], [3, 5, 0, 0], [3, 5, 962, 30]]


def wrap(x, bits):
    x &= (1 << bits) - 1
    return x - (1 << bits) if x >> (bits - 1) else x


class Ctx:
    def __init__(self, j=0):
        b0, b1, v01, v0 = INIT[j]
        self.s0, self.s01, self.h0, self.h1 = v0, v01, b0 + 4, b1

    def mps(self):
        return 1 if self.s01 >= 0 else 0

    def rlps(self, rng):
        i = abs(self.s01 >> 7) + (rng & 0xE0)
        if not 0 <= i < len(T):
            raise Refuse("nnc_malformed", "state outside table A.1")
        return T[i]

    def update(self, b):
        sg = 1 if b == 1 else -1
        s1 = self.s01 - 16 * self.s0

        def cv(x, k):
            i = 16 + ((sg * x) >> k)
            if not 0 <= i < len(C):
                raise Refuse("nnc_malformed", "state outside table A.2")
            return C[i]
        s0 = wrap(self.s0 + sg * (cv(self.s0, 3) >> self.h0), 8)
        s1 = s1 + sg * (cv(s1, 7) >> self.h1)
        self.s0, self.s01 = s0, wrap(s1 + 16 * s0, 16)


def lz(x):
    return 32 - x.bit_length()


class Dec:
    def __init__(self, b):
        if len(b) < 2:
            raise Refuse("nnc_truncated", "payload under 2 bytes")
        self.b, self.p = b, 2
        self.rng, self.val, self.need = 510, 256 * b[0] + b[1], -8

    def byte(self):
        if self.p >= len(self.b):
            raise Refuse("nnc_truncated", "payload ends inside its unit")
        v = self.b[self.p]
        self.p += 1
        return v

    def bin(self, ctx):
        r = ctx.rlps(self.rng)
        m = self.rng - r
        if self.val < m * 128:
            b, self.rng = ctx.mps(), m
        else:
            self.val -= m * 128
            self.rng, b = r, 1 - ctx.mps()
        n = lz(self.rng) - 23
        self.rng <<= n
        self.val <<= n
        self.need += n
        if self.need >= 0:
            self.val += self.byte() << self.need
            self.need -= 8
        ctx.update(b)
        return b

    def bypass(self, k=1):
        v = 0
        for _ in range(k):
            self.val *= 2
            self.need += 1
            if self.need >= 0:
                self.val += self.byte()
                self.need = -8
            bit = 0
            if self.val >= self.rng * 128:
                self.val -= self.rng * 128
                bit = 1
            v = 2 * v + bit
        return v

    def finish(self):
        self.rng -= 2
        if self.val < self.rng * 128:
            raise Refuse("nnc_malformed", "no terminating bin")
        if ((self.b[self.p - 1] << (8 + self.need)) & 0xFF) != 0x80:
            raise Refuse("nnc_malformed", "no stop pattern")
        return self.p


def cpos(p):
    return 0 if p == 0 else (1 if p < 0 else 2)


def a_decode(b, n, g):
    d = Dec(b)
    L = 24 + 3 + 2 * g + 32 + 4
    ctx = [Ctx() for _ in range(L)]
    for i in range(L - 4):
        if 3 <= i <= 23:
            continue
        j = 1 + d.bypass(3) if d.bin(ctx[L - 2]) == 1 else 0
        ctx[i] = Ctx(j)
    out, p = [], 0
    for _ in range(n):
        v = 0
        if d.bin(ctx[cpos(p)]) == 1:
            v = 1
            s = d.bin(ctx[24 + cpos(p)])
            gg = d.bin(ctx[27 + s])
            k = 1
            while gg == 1 and k < g:
                v += 1
                gg = d.bin(ctx[27 + s + 2 * k])
                k += 1
            if gg == 1:
                v += 1
                q = 27 + 2 * g
                rem = 0
                if d.bin(ctx[q]) == 1:
                    rem, q, e = 1, q + 1, 1
                    while True:
                        if q >= 59 + 2 * g:
                            raise Refuse("nnc_malformed", "remainder prefix")
                        if d.bin(ctx[q]) == 0:
                            break
                        rem += 1 << e
                        q += 1
                        e += 1
                    rem += d.bypass(e)
                v += rem
            if v > (1 << 31) - 1:
                raise Refuse("nnc_malformed", "past 2^31 - 1")
            if s == 1:
                v = -v
        out.append(v)
        p = v
    return out, d.finish()


class Enc:
    M = 0xFFFFFFFF

    def __init__(self):
        self.low, self.rng, self.left, self.B, self.n, self.o = 0, 510, 23, 0, 0, bytearray()

    def out(self):
        lead = self.low >> (24 - self.left)
        self.left += 8
        self.low &= self.M >> self.left
        if lead == 0xFF:
            self.n += 1
        elif self.n > 0:
            c = lead >> 8
            self.o.append((self.B + c) % 256)
            self.B = lead % 256
            for _ in range(self.n - 1):
                self.o.append((0xFF + c) % 256)
            self.n = 1
        else:
            self.n, self.B = 1, lead

    def bin(self, b, ctx):
        r = ctx.rlps(self.rng)
        self.rng -= r
        if b == ctx.mps():
            if self.rng < 256:
                self.rng *= 2
                self.low = (2 * self.low) & self.M
                self.left -= 1
                if self.left < 12:
                    self.out()
        else:
            k = lz(r) - 23
            self.low = (self.low + self.rng) & self.M
            self.rng = r << k
            self.low = (self.low << k) & self.M
            self.left -= k
            if self.left < 12:
                self.out()
        ctx.update(b)

    def bypass(self, v, k):
        for i in reversed(range(k)):
            self.low = (2 * self.low) & self.M
            if (v >> i) & 1:
                self.low = (self.low + self.rng) & self.M
            self.left -= 1
            if self.left < 12:
                self.out()

    def finish(self):
        self.rng -= 2
        self.low = ((self.low + self.rng) << 7) & self.M
        self.rng = 256
        self.left -= 7
        if self.left < 12:
            self.out()
        if self.low >> (32 - self.left):
            self.o.append((self.B + 1) % 256)
            self.o += b"\0" * (self.n - 1)
            self.low -= 1 << (32 - self.left)
        else:
            if self.n > 0:
                self.o.append(self.B)
            self.o += b"\xff" * max(self.n - 1, 0)
        self.low = 2 * (self.low >> 8) + 1
        self.left -= 1
        self.low = (self.low << (32 - (24 - self.left))) & self.M
        i = 0
        while i < 24 - self.left:
            self.o.append((self.low >> (24 - i)) & 0xFF)
            i += 8
        return bytes(self.o)


def a_encode(vals, g):
    e = Enc()
    L = 24 + 3 + 2 * g + 32 + 4
    ctx = [Ctx() for _ in range(L)]
    for i in range(L - 4):
        if 3 <= i <= 23:
            continue
        e.bin(0, ctx[L - 2])
    p = 0
    for v in vals:
        if abs(v) > (1 << 31) - 1:
            raise Refuse("value_out_of_range", "no NNR_PT_INT form")
        e.bin(int(v != 0), ctx[cpos(p)])
        if v != 0:
            s = 1 if v < 0 else 0
            e.bin(s, ctx[24 + cpos(p)])
            r = abs(v) - 1
            gg = int(r > 0)
            e.bin(gg, ctx[27 + s])
            k = 1
            while gg == 1 and k < g:
                r -= 1
                gg = int(r > 0)
                e.bin(gg, ctx[27 + s + 2 * k])
                k += 1
            if gg == 1:
                r -= 1
                q = 27 + 2 * g
                if r == 0:
                    e.bin(0, ctx[q])
                else:
                    e.bin(1, ctx[q])
                    base, ex = 1, 1
                    q += 1
                    while r > base + (1 << ex) - 1:
                        e.bin(1, ctx[q])
                        base += 1 << ex
                        q += 1
                        ex += 1
                    e.bin(0, ctx[q])
                    e.bypass(r - base, ex)
        p = v
    return e.finish()


# ------------------------------------------------------------------ §12 NNC

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

    def u(self, n):
        v = 0
        for _ in range(n):
            if self.bit // 8 >= len(self.b):
                raise Refuse("nnc_truncated", "a field past its unit")
            v = 2 * v + ((self.b[self.bit // 8] >> (7 - self.bit % 8)) & 1)
            self.bit += 1
        return v

    def i(self, n):
        v = self.u(n)
        return v - (1 << n) if v >> (n - 1) else v

    def ue(self, k):
        x = 0
        while True:
            if k >= 40:
                raise Refuse("nnc_malformed", "exp-Golomb too long")
            if self.u(1) == 1:
                break
            x += 1 << k
            k += 1
        return x + self.u(k)

    def align(self):
        if self.u(1) != 1:
            raise Refuse("nnc_malformed", "alignment bit")
        while self.bit % 8:
            if self.u(1) != 0:
                raise Refuse("nnc_malformed", "alignment padding")

    def st(self):
        if self.bit % 8:
            raise Refuse("nnc_malformed", "string not on a byte")
        at = self.bit // 8
        z = self.b.find(b"\0", at)
        if z < 0:
            raise Refuse("nnc_truncated", "string without zero byte")
        try:
            s = self.b[at:z].decode("utf-8")
        except UnicodeDecodeError:
            raise Refuse("nnc_malformed", "string not UTF-8")
        self.bit = (z + 1) * 8
        return s


def nnc_read(b):
    s = TSet()
    at, mps, tpl, ndu, first = 0, False, False, False, True
    names = []
    while at < len(b):
        h = Bits(b[at:])
        flag = h.u(1)
        size = h.u(31 if flag else 15)
        if size < 3 or at + size > len(b):
            raise Refuse("nnc_truncated", "unit size")
        unit = b[at:at + size]
        h = Bits(unit, h.bit)
        ty = h.u(6)
        h.u(1)
        partial = h.u(1)
        if first and ty != 0:
            raise Refuse("nnc_malformed", "first unit")
        if partial:
            raise Refuse("nnc_not_read", "partial data counter")

        def done():
            if h.bit != 8 * len(unit):
                raise Refuse("nnc_malformed", "unit size")
        if ty == 0:
            if not first:
                raise Refuse("nnc_malformed", "second STR")
            h.u(8)
            done()
        elif ty == 1:
            if mps or ndu:
                raise Refuse("nnc_malformed", "MPS")
            mps = True
            h.u(1)
            maps, q, idx = h.u(4), h.u(3), h.u(1)
            if h.u(7):
                raise Refuse("nnc_malformed", "reserved")
            if maps:
                raise Refuse("nnc_not_read", "performance map")
            if idx:
                raise Refuse("nnc_not_read", "indexed reference")
            if q:
                h.u(3)
                h.i(13)
            h.align()
            done()
        elif ty == 3:
            if tpl or not mps:
                raise Refuse("nnc_malformed", "TPL")
            tpl = True
            fmt, comp = h.u(8), h.u(8)
            if fmt not in (0, 3, 4) or comp:
                raise Refuse("nnc_not_read", "topology format")
            text = h.st()
            done()
            if fmt == 0:
                try:
                    obj = strict_json(text.encode())
                except Refuse:
                    obj = None
                if isinstance(obj, tuple) and obj[0] == "obj" and all(isinstance(v, str) for _, v in obj[1]):
                    for k, v in obj[1]:
                        s.meta[k] = v
        elif ty == 5:
            if not mps:
                raise Refuse("nnc_malformed", "NDU before MPS")
            ndu = True
            name, t = nnc_ndu(h, unit)
            names.append(name)
            s.t[name] = t
        elif ty in (2, 4, 6):
            raise Refuse("nnc_not_read", f"unit type {ty}")
        else:
            raise Refuse("nnc_malformed", f"unit type {ty}")
        first = False
        at += size
    if first:
        raise Refuse("nnc_truncated", "empty")
    if len(set(names)) != len(names):
        raise Refuse("nnc_malformed", "a name twice")
    return s


def nnc_ndu(h, unit):
    pt, multi, fmtp, inband = h.u(5), h.u(1), h.u(1), h.u(1)
    if pt not in (0, 2):
        raise Refuse("nnc_not_read", f"payload type {pt}")
    if multi:
        raise Refuse("nnc_not_read", "several elements")
    name = h.st()
    if name == "":
        raise Refuse("nnc_malformed", "empty name")
    if pt == 0 and h.u(1):
        raise Refuse("nnc_not_read", "dq_flag")
    if fmtp and h.u(7) != (0 if pt == 0 else 1):
        raise Refuse("nnc_not_read", "decompressed data format")
    if not inband:
        raise Refuse("nnc_not_read", "out of band")
    dflag, uflag, cpt = h.u(1), h.u(1), h.u(4)
    if cpt:
        raise Refuse("nnc_not_read", "compressed parameter types")
    if not dflag:
        raise Refuse("nnc_not_read", "dimensions out of band")
    nd = h.ue(1)
    if not 1 <= nd <= 16:
        raise Refuse("nnc_malformed", "dimension count")
    shape = [h.ue(7) for _ in range(nd)]
    cnt = product(shape)
    if cnt == 0 or cnt >= 1 << 64:
        raise Refuse("nnc_malformed", "count")
    if cnt > 1 << 28:
        raise Refuse("nnc_not_read", "too many elements")
    unary = h.u(8) if uflag else None
    if nd > 1 and h.u(4):
        raise Refuse("nnc_not_read", "scan order")
    h.align()
    rest = unit[h.bit // 8:]
    if pt == 2:
        if len(rest) != 4 * cnt:
            raise Refuse("nnc_malformed", "raw floats")
        return name, ["F32", shape, bytes(rest)]
    if unary is None:
        raise Refuse("nnc_not_read", "unary length out of band")
    vals, used = a_decode(rest, cnt, unary + 1)
    if used != len(rest):
        raise Refuse("nnc_malformed", "payload length")
    return name, ["I32", shape, b"".join(struct.pack("<i", v) for v in vals)]


class W:
    def __init__(self):
        self.o, self.bit = bytearray(), 0

    def u(self, n, v):
        for i in reversed(range(n)):
            if self.bit % 8 == 0:
                self.o.append(0)
            if (v >> i) & 1:
                self.o[-1] |= 1 << (7 - self.bit % 8)
            self.bit += 1

    def ue(self, k, x):
        while x >= 1 << k:
            self.u(1, 0)
            x -= 1 << k
            k += 1
        self.u(1, 1)
        self.u(k, x)

    def align(self):
        self.u(1, 1)
        while self.bit % 8:
            self.u(1, 0)

    def st(self, s):
        self.o += s.encode() + b"\0"
        self.bit = 8 * len(self.o)


def sized(body):
    w = W()
    if len(body) + 2 < 32768:
        w.u(1, 0)
        w.u(15, len(body) + 2)
    else:
        w.u(1, 1)
        w.u(31, len(body) + 4)
    return bytes(w.o) + bytes(body)


def head(w, ty):
    w.u(6, ty)
    w.u(1, 1)
    w.u(1, 0)


def ints_of(dt, data):
    k = INT_WIDTH[dt]
    return [int.from_bytes(data[i:i + k], "little", signed=True) for i in range(0, len(data), k)]


def nnc_write(s):
    """§12.3."""
    o = bytearray()
    w = W()
    head(w, 0)
    w.u(8, 0)
    o += sized(w.o)
    w = W()
    head(w, 1)
    w.u(1, 1 if s.meta else 0)
    w.u(4, 0)
    w.u(3, 0)
    w.u(1, 0)
    w.u(7, 0)
    w.align()
    o += sized(w.o)
    if s.meta:
        w = W()
        head(w, 3)
        w.u(8, 0)
        w.u(8, 0)
        w.st("{" + ",".join(f"{jquote(k)}:{jquote(s.meta[k])}" for k in sorted(s.meta, key=lambda x: x.encode())) + "}")
        o += sized(w.o)
    for name in sorted(s.t, key=lambda x: x.encode()):
        dt, shape, data = s.t[name]
        pt = 0 if dt in INT_WIDTH else 2
        w = W()
        head(w, 5)
        w.u(5, pt)
        w.u(1, 0)
        w.u(1, 0)
        w.u(1, 1)
        w.st(name)
        if pt == 0:
            w.u(1, 0)
        w.u(1, 1)
        w.u(1, 1)
        w.u(4, 0)
        w.ue(1, len(shape))
        for d in shape:
            w.ue(7, d)
        w.u(8, 10)
        if len(shape) > 1:
            w.u(4, 0)
        w.align()
        body = bytes(w.o) + (a_encode(ints_of(dt, data), 11) if pt == 0 else bytes(data))
        o += sized(body)
    return bytes(o)


# ------------------------------------------------------------------ §10 profile and bindings

READ = {"wai.prior.safetensors": st_read, "wai.prior.gguf": gg_read, "wai.prior.nnc": nnc_read}
WRITE = {"wai.prior.safetensors": st_write, "wai.prior.gguf": gg_write, "wai.prior.nnc": nnc_write}
U32 = (1 << 32) - 1


def profile_meta(binding, form):
    m = {"wai.binding": binding, "wai.profile": "wai-int/1"}
    if form == "wai.prior.gguf":
        m["general.architecture"] = "wai"
    return m


def ints_tensor(dt, shape, vals):
    k = INT_WIDTH[dt]
    lo, hi = -(1 << (8 * k - 1)), (1 << (8 * k - 1)) - 1
    for v in vals:
        if not lo <= v <= hi:
            raise Refuse("binding_mismatch", "a value wider than its type")
    return [dt, list(shape), b"".join(v.to_bytes(k, "little", signed=True) for v in vals)]


class N:
    """A cursor over a native artifact."""

    def __init__(self, b):
        self.b, self.p = b, 0

    def take(self, n, why="truncated"):
        if self.p + n > len(self.b):
            raise Refuse("native_malformed", why)
        s = self.b[self.p:self.p + n]
        self.p += n
        return s

    def u(self, k):
        return int.from_bytes(self.take(k), "little")

    def ints(self, n, k):
        raw = self.take(n * k)
        return [int.from_bytes(raw[i:i + k], "little", signed=True) for i in range(0, len(raw), k)]


def wsr_from(b):
    """exact-upscale §3.1 into the tensors of §10.2."""
    c = N(b)
    if c.take(4, "bad_magic") != b"WSR1":
        raise Refuse("native_malformed", "bad_magic")
    if c.u(1) != 1:
        raise Refuse("native_malformed", "version")
    n, act = c.u(1), c.u(2)
    if not (1 <= n <= 8 and 1 <= act <= 32767):
        raise Refuse("native_malformed", "field_range")
    s = TSet()
    s.t["act_hi"] = ints_tensor("I16", [1], [act])
    cin = 3
    for i in range(n):
        k, shift, cout = c.u(1), c.u(1), c.u(2)
        cok = cout == 12 if i == n - 1 else 1 <= cout <= 64
        if k not in (1, 3, 5) or shift > 31 or not cok:
            raise Refuse("native_malformed", "field_range")
        w = c.ints(cout * cin * k * k, 1)
        bias = c.ints(cout, 4)
        mult = [c.u(2) for _ in range(cout)]
        if 0 in mult:
            raise Refuse("native_malformed", "field_range")
        s.t[f"layer.{i}.weight"] = ints_tensor("I8", [cout, cin, k, k], w)
        s.t[f"layer.{i}.bias"] = ints_tensor("I32", [cout], bias)
        s.t[f"layer.{i}.mult"] = ints_tensor("I32", [cout], mult)
        s.t[f"layer.{i}.shift"] = ints_tensor("I8", [1], [shift])
        cin = cout
    if c.p != len(b):
        raise Refuse("native_malformed", "trailing")
    return s


def whm_from(b):
    """integer-payloads §9 into the tensors of §10.2."""
    c = N(b)
    if c.take(4) != b"WHM1":
        raise Refuse("native_malformed", "magic")
    af, gf, sh, n = c.u(4), c.u(4), c.u(4), c.u(4)
    if n > 1000:
        raise Refuse("binding_mismatch", "layers")
    s = TSet()
    s.t["header"] = ints_tensor("I64", [3], [af, gf, sh])
    for i in range(n):
        tag = c.u(1)
        p = f"layer.{i:03d}."
        if tag in (0, 1):
            geo = [c.u(4) for _ in range(6 if tag == 0 else 5)]
            inc, out, k = geo[:3]
            if 0 in (inc, out, k):
                raise Refuse("binding_mismatch", "a zero dimension")
            w = c.ints(inc * out * k * k, 2)
            mult, bias = c.ints(out, 8), c.ints(out, 8)
            kind = "convt" if tag == 0 else "conv"
            s.t[p + kind] = ints_tensor("I64", [len(geo)], geo)
            s.t[p + "weight"] = ints_tensor("I16", [inc, out, k, k] if tag == 0 else [out, inc, k, k], w)
            s.t[p + "mult"] = ints_tensor("I64", [out], mult)
            s.t[p + "bias"] = ints_tensor("I64", [out], bias)
        elif tag == 2:
            ch = c.u(4)
            if ch == 0:
                raise Refuse("binding_mismatch", "a zero dimension")
            s.t[p + "igdn"] = ints_tensor("I64", [1], [ch])
            s.t[p + "gamma"] = ints_tensor("I32", [ch, ch], c.ints(ch * ch, 4))
            s.t[p + "beta"] = ints_tensor("I64", [ch], c.ints(ch, 8))
        elif tag == 3:
            s.t[p + "relu"] = ints_tensor("I8", [1], [0])
        elif tag == 4:
            q = c.ints(1, 8)[0]
            s.t[p + "leaky"] = ints_tensor("I64", [2], [q, c.u(4)])
        else:
            raise Refuse("native_malformed", "tag")
    if c.p != len(b):
        raise Refuse("native_malformed", "trailing")
    return s


def tvals(s, name):
    dt, shape, data = s.t[name]
    return ints_of(dt, data)


def check_wants(s, form, wants, rng):
    """§10.3 step 4's names, types and shapes, then step 5's values."""
    if set(s.t) != {w[0] for w in wants}:
        raise Refuse("binding_mismatch", "names")
    for name, dt, shape in wants:
        want_dt = "I32" if form == "wai.prior.nnc" else dt
        if s.t[name][0] != want_dt or s.t[name][1] != shape:
            raise Refuse("binding_mismatch", name)
    for name, dt, shape in wants:
        k = INT_WIDTH[dt]
        tlo, thi = -(1 << (8 * k - 1)), (1 << (8 * k - 1)) - 1
        for i, v in enumerate(tvals(s, name)):
            r = rng(name, i)
            lo, hi = (max(tlo, r[0]), min(thi, r[1])) if r else (tlo, thi)
            if not lo <= v <= hi:
                raise Refuse("value_out_of_range", f"{name}[{i}]")


def indices(s, prefix, suffix, width):
    out = []
    for n in s.t:
        if n.startswith(prefix) and n.endswith(suffix) and len(n) > len(prefix) + len(suffix):
            d = n[len(prefix):len(n) - len(suffix)]
            if re.fullmatch(r"[0-9]+", d) and (len(d) == width if width else (d == "0" or d[0] != "0")):
                out.append(int(d))
    return sorted(out)


def wsr_to(s, form):
    idx = indices(s, "layer.", ".weight", None)
    n = len(idx)
    if not 1 <= n <= 8 or idx != list(range(n)):
        raise Refuse("binding_mismatch", "layers")
    couts, ks = [], []
    for i in range(n):
        shape = s.t[f"layer.{i}.weight"][1]
        cout, k = (shape[0], shape[-1]) if shape else (0, 0)
        cok = cout == 12 if i == n - 1 else 1 <= cout <= 64
        if k not in (1, 3, 5) or not cok:
            raise Refuse("binding_mismatch", "weight shape")
        couts.append(cout)
        ks.append(k)
    wants = [("act_hi", "I16", [1])]
    cin = 3
    for i in range(n):
        wants += [(f"layer.{i}.weight", "I8", [couts[i], cin, ks[i], ks[i]]), (f"layer.{i}.bias", "I32", [couts[i]]),
                  (f"layer.{i}.mult", "I32", [couts[i]]), (f"layer.{i}.shift", "I8", [1])]
        cin = couts[i]

    def rng(name, i):
        if name == "act_hi":
            return (1, 32767)
        if name.endswith(".mult"):
            return (1, 65535)
        if name.endswith(".shift"):
            return (0, 31)
        return None
    check_wants(s, form, wants, rng)
    o = bytearray(b"WSR1") + bytes([1, n]) + tvals(s, "act_hi")[0].to_bytes(2, "little")
    for i in range(n):
        o += bytes([ks[i], tvals(s, f"layer.{i}.shift")[0]]) + couts[i].to_bytes(2, "little")
        o += bytes(v & 0xFF for v in tvals(s, f"layer.{i}.weight"))
        o += b"".join(v.to_bytes(4, "little", signed=True) for v in tvals(s, f"layer.{i}.bias"))
        o += b"".join(v.to_bytes(2, "little") for v in tvals(s, f"layer.{i}.mult"))
    return bytes(o)


KINDS = ["convt", "conv", "igdn", "relu", "leaky"]


def whm_to(s, form):
    kinds = {}
    for kd in KINDS:
        for i in indices(s, "layer.", "." + kd, 3):
            if i in kinds:
                raise Refuse("binding_mismatch", "two kinds")
            kinds[i] = kd
    n = len(kinds)
    if sorted(kinds) != list(range(n)):
        raise Refuse("binding_mismatch", "layers")
    wants = [("header", "I64", [3])]
    for i in range(n):
        kd = kinds[i]
        p = f"layer.{i:03d}."
        if kd == "relu":
            wants.append((p + "relu", "I8", [1]))
            continue
        ln = {"convt": 6, "conv": 5, "igdn": 1, "leaky": 2}[kd]
        dt = "I32" if form == "wai.prior.nnc" else "I64"
        if s.t[p + kd][0] != dt or s.t[p + kd][1] != [ln]:
            raise Refuse("binding_mismatch", p + kd)
        g = tvals(s, p + kd)
        dims = g[:3] if kd in ("conv", "convt") else (g[:1] if kd == "igdn" else [])
        for x in dims:
            if not 1 <= x <= U32:
                raise Refuse("value_out_of_range", p + kd)
        wants.append((p + kd, "I64", [ln]))
        if kd in ("conv", "convt"):
            inc, out, k = g[:3]
            wants += [(p + "weight", "I16", [inc, out, k, k] if kd == "convt" else [out, inc, k, k]),
                      (p + "mult", "I64", [out]), (p + "bias", "I64", [out])]
        elif kd == "igdn":
            wants += [(p + "gamma", "I32", [g[0], g[0]]), (p + "beta", "I64", [g[0]])]

    def rng(name, i):
        if name == "header":
            return (0, U32)
        if name.endswith((".convt", ".conv")):
            return (1, U32) if i < 3 else (0, U32)
        if name.endswith(".igdn"):
            return (1, U32)
        if name.endswith(".relu"):
            return (0, 0)
        if name.endswith(".leaky") and i == 1:
            return (0, U32)
        return None
    check_wants(s, form, wants, rng)
    o = bytearray(b"WHM1") + b"".join(v.to_bytes(4, "little") for v in tvals(s, "header")) + n.to_bytes(4, "little")
    for i in range(n):
        kd = kinds[i]
        p = f"layer.{i:03d}."
        if kd in ("conv", "convt"):
            o.append(0 if kd == "convt" else 1)
            o += b"".join(v.to_bytes(4, "little") for v in tvals(s, p + kd))
            o += b"".join(v.to_bytes(2, "little", signed=True) for v in tvals(s, p + "weight"))
            o += b"".join(v.to_bytes(8, "little", signed=True) for v in tvals(s, p + "mult") + tvals(s, p + "bias"))
        elif kd == "igdn":
            o.append(2)
            o += tvals(s, p + "igdn")[0].to_bytes(4, "little")
            o += b"".join(v.to_bytes(4, "little", signed=True) for v in tvals(s, p + "gamma"))
            o += b"".join(v.to_bytes(8, "little", signed=True) for v in tvals(s, p + "beta"))
        elif kd == "relu":
            o.append(3)
        else:
            q, fr = tvals(s, p + "leaky")
            o.append(4)
            o += q.to_bytes(8, "little", signed=True) + fr.to_bytes(4, "little")
    return bytes(o)


BINDINGS = {"wsr1/1": (wsr_from, wsr_to), "whm1/1": (whm_from, whm_to)}


def load(form, b):
    """§10.3."""
    s = READ[form](b)
    if WRITE[form](s) != b:
        raise Refuse("header_not_canonical")
    bname = s.meta.get("wai.binding", "")
    want = profile_meta(bname, form)
    if s.meta.get("wai.profile") != "wai-int/1" or set(s.meta) != set(want) or (form == "wai.prior.gguf" and s.meta.get("general.architecture") != "wai"):
        raise Refuse("profile_absent")
    if bname not in BINDINGS:
        raise Refuse("binding_unknown", bname)
    for name, (dt, _, _) in s.t.items():
        if (dt != "I32") if form == "wai.prior.nnc" else (dt not in INT_WIDTH):
            raise Refuse("dtype_not_integer", name)
    return bname, BINDINGS[bname][1](s, form)


def encode(binding, native, form):
    s = BINDINGS[binding][0](native)
    s.meta = profile_meta(binding, form)
    return WRITE[form](s)


# ------------------------------------------------------------------ the corpus

def run(corpus):
    doc = json.load(open(os.path.join(corpus, "vectors.json"), encoding="utf-8"))
    known = codes()
    fails = []
    natives = {}
    for nat in doc["natives"]:
        b = open(os.path.join(corpus, nat["path"]), "rb").read()
        if sha(b) != nat["sha256"]:
            fails.append(f"native {nat['name']} does not match its pin")
        natives[nat["name"]] = b
    files = {}
    for f in doc["files"]:
        b = open(os.path.join(corpus, f["file"]), "rb").read()
        if sha(b) != f["sha256"]:
            fails.append(f"{f['file']} does not match its digest")
        files[f["file"]] = b
    listed = set(files) | {"vectors.json", "README.md", ".gitattributes"}
    for root, _, fs in os.walk(corpus):
        for x in fs:
            rel = os.path.relpath(os.path.join(root, x), corpus).replace(os.sep, "/")
            if rel not in listed:
                fails.append(f"{rel} is not listed")
    accepted = []
    for v in doc["vectors"]:
        name, op, form = v["name"], v["op"], v["form"]
        if v["outcome"] == "refuse" and v["code"] not in known:
            fails.append(f"{name}: code {v['code']} is not in §13")
        try:
            if op == "encode":
                native = natives[v["native"]] if "native" in v else files[v["native_file"]]
                got = encode(v["binding"], native, form)
                outcome = ("accept", got)
            else:
                b, native = load(form, files[v["file"]])
                outcome = ("accept", (b, native))
        except Refuse as r:
            outcome = ("refuse", r.code)
        if v["outcome"] == "refuse":
            if outcome != ("refuse", v["code"]):
                fails.append(f"{name}: wanted {v['code']}, got {outcome[0]} {outcome[1] if outcome[0] == 'refuse' else ''}")
        elif outcome[0] != "accept":
            fails.append(f"{name}: wanted accept, got {outcome[1]}")
        elif op == "encode" and outcome[1] != files[v["file"]]:
            fails.append(f"{name}: the encoded file differs from {v['file']}")
        elif op == "load":
            want_native = natives[v["native"]]
            if outcome[1] != (v["binding"], want_native):
                fails.append(f"{name}: loaded to another binding or other bytes")
            else:
                accepted.append((form, files[v["file"]], want_native))
    if set(known) - {v.get("code") for v in doc["vectors"]}:
        fails.append(f"codes no vector gives: {sorted(set(known) - {v.get('code') for v in doc['vectors']})}")
    # Mutations: one byte flipped at a few places of every accepted file.
    caught = missed = 0
    for form, b, native in accepted:
        for at in sorted({0, 7, len(b) // 3, len(b) // 2, len(b) - 1}):
            m = bytearray(b)
            m[at] ^= 0x01
            try:
                _, got = load(form, bytes(m))
                ok = got != native
            except Refuse:
                ok = True
            if ok:
                caught += 1
            else:
                missed += 1
                fails.append(f"a flip at byte {at} of a {form} file loaded to the same prior")
    return fails, len(doc["vectors"]), caught


def main():
    if len(sys.argv) != 2:
        print(__doc__)
        sys.exit(2)
    fails, n, caught = run(sys.argv[1])
    for f in fails:
        print("FAIL", f)
    if fails:
        sys.exit(1)
    print(f"prior-forms: every vector verifies ({n}), and every mutated copy is caught ({caught})")


if __name__ == "__main__":
    main()