"""The second implementation of LOC carriage (extensions/loc-carriage.md), written from that
text, moq-streaming-format §5, jwp-receipts §2.1 and MOQT -22's encoding of properties, sharing
no code with joule-moq.

    python3 -I loc_verify.py <dir-holding-vectors.json>

For every object of vectors.json it decodes the MOQ Object Properties, reads the object as
loc-carriage §5.1 says, requires the envelope, LOC properties and receipt property stated, and
writes the object again from them as §4.1 says, requiring the same bytes; it requires every
refused object and envelope to be refused with its code, every code of §6's table given, and
the codes to be §6's; then it requires every mutated copy of the corpus, and the corpus under a
copy of this verifier with one of its own checks dropped, to fail. Content hashes are BLAKE3,
computed with ../claim-conformance/_blake3.py. Standard library only.
"""

import copy
import json
import os
import re
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "claim-conformance"))
import _blake3 as B3  # noqa: E402

IMMUTABLE = 0x0B
TIMESCALE, FRAME_MARKING, AUDIO_LEVEL, VIDEO_CONFIG, AUDIO_CONFIG, TIMESTAMP = 0x08, 0x09, 0x0C, 0x0D, 0x0F, 0x10
WAI_MANIFEST, WAI_RECEIPT = 0x3A61, 0x3A41
MAPPED = ("wai.audio.opus.packet", "wai.video.avc", "wai.video.hevc", "wai.audio.aac")
DROPPED = set()  # checks a mutated verifier leaves out


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


# ---- MOQT -22: the variable-length integer (§8.1) and Key-Value-Pairs (§8.3) -------

def get_varint(b, i):
    if i >= len(b):
        raise ValueError("truncated")
    first = b[i]
    ones = 0
    while ones < 8 and first & (0x80 >> ones):
        ones += 1
    n = 9 if ones >= 8 else ones + 1
    if i + n > len(b):
        raise ValueError("truncated")
    if n == 9:
        return int.from_bytes(b[i + 1:i + 9], "big"), i + 9
    v = first & ((0xFF >> n) if n > 1 else 0x7F)
    for x in b[i + 1:i + n]:
        v = (v << 8) | x
    return v, i + n


def put_varint(v):
    for n in range(1, 9):
        bits = 56 if n == 8 else 7 * n
        if v < (1 << bits):
            prefix = 0 if n == 1 else (0xFF << (8 - (n - 1))) & 0xFF
            raw = v.to_bytes(n, "big")
            return bytes([prefix | raw[0]]) + raw[1:]
    return b"\xff" + v.to_bytes(8, "big")


def decode_pairs(b):
    out, i, prev = [], 0, 0
    while i < len(b):
        d, i = get_varint(b, i)
        t = prev + d
        prev = t
        if t % 2 == 0:
            v, i = get_varint(b, i)
        else:
            n, i = get_varint(b, i)
            if n > 0xFFFF or i + n > len(b):
                raise ValueError("a value past its bound")
            v, i = bytes(b[i:i + n]), i + n
        out.append((t, v))
    return out


def encode_pairs(pairs):
    out, prev = b"", 0
    for t, v in sorted(pairs, key=lambda p: p[0]):
        out += put_varint(t - prev)
        prev = t
        out += put_varint(v) if t % 2 == 0 else put_varint(len(v)) + v
    return out


# ---- loc-carriage §3.2: one chunk --------------------------------------------------

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:
        return None
    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 one_chunk(cap, p):
    if cap == "wai.audio.opus.packet":
        return len(p) >= 1
    if cap == "wai.video.avc":
        us = nal_units(p)
        return us is not None and all(u and not u[0] & 0x80 for u in us) and \
            sum(1 for u in us if 1 <= (u[0] & 0x1F) <= 5 and len(u) > 1 and u[1] & 0x80) == 1
    if cap == "wai.video.hevc":
        us = nal_units(p)
        return us is not None and all(len(u) >= 2 and not u[0] & 0x80 and ((u[0] & 1) << 5 | u[1] >> 3) == 0 for u in us) and \
            sum(1 for u in us if len(u) >= 3 and (u[0] >> 1) & 0x3F < 32 and u[2] & 0x80) == 1
    if cap == "wai.audio.aac":
        return len(p) >= 7 and p[0] == 0xFF and p[1] >> 4 == 0xF and (p[1] >> 1) & 3 == 0 and \
            ((p[3] & 3) << 11 | p[4] << 3 | p[5] >> 5) == len(p)
    raise Refuse("loc_capability_unmapped")


def manifest_capability(m):
    try:
        v = json.loads(m.decode("utf-8"))
    except (ValueError, UnicodeDecodeError):
        raise Refuse("loc_manifest_invalid")
    if not isinstance(v, dict):
        raise Refuse("loc_manifest_invalid")
    wai = v.get("wai")
    cap = (v.get("model_requirement") or {}).get("capability") if isinstance(v.get("model_requirement"), dict) else None
    sec = v.get("security")
    enc = sec.get("encryption") if isinstance(sec, dict) else None
    if not (isinstance(wai, str) and wai.split(".")[0] == "1") or not isinstance(cap, str) or enc not in (None, "none"):
        raise Refuse("loc_manifest_invalid")
    return cap


def envelope(m, p):
    return b"WAI1" + len(m).to_bytes(4, "little") + m + len(p).to_bytes(4, "little") + p


def content_hash(env):
    return B3.blake3(b"moq-jwp:object\x01" + env)


def read_receipt(b):
    try:
        if len(b) < 32:
            raise ValueError
        h, i = b[:32], 32
        joules, i = get_varint(b, i)
        n, i = get_varint(b, i)
        acq = b[i:i + n].decode("utf-8")
        if i + n > len(b):
            raise ValueError
        i += n
        n, i = get_varint(b, i)
        if i + n > len(b):
            raise ValueError
        ev, i = b[i:i + n], i + n
        if len(b) != i + 32:
            raise ValueError
        group = b[i:i + 32]
    except (ValueError, UnicodeDecodeError):
        raise Refuse("property_invalid")
    # moq-streaming-format §5's rules on the property alone, for the classes this corpus uses.
    if joules == 0:
        ok = acq == "" and ev == b""
    elif acq == "HwShunt":
        ok = ev == b""
    else:
        raise RuntimeError("a class this verifier does not read: %r" % acq)
    if not ok:
        raise Refuse("property_invalid")
    return {"content_hash": h, "joules": joules, "acquisition": acq, "evidence": ev, "group": group}


# ---- loc-carriage §5.1: reading ------------------------------------------------------

LOC_RULES = {
    TIMESTAMP: lambda v: isinstance(v, int),
    TIMESCALE: lambda v: isinstance(v, int) and ("timescale" in DROPPED or v != 0),
    FRAME_MARKING: lambda v: isinstance(v, bytes) and 1 <= len(v) <= 4,
    AUDIO_LEVEL: lambda v: isinstance(v, int) and v <= 255,
    VIDEO_CONFIG: lambda v: isinstance(v, bytes) and len(v) > 0,
    AUDIO_CONFIG: lambda v: isinstance(v, bytes) and len(v) > 0,
}


def read_object(props_bytes, payload, track_cap):
    try:
        outer = decode_pairs(props_bytes)
        imm = [v for t, v in outer if t == IMMUTABLE]
        if len(imm) > 1:
            raise ValueError
        inner = decode_pairs(imm[0]) if imm else []
        if any(t == IMMUTABLE for t, _ in inner):
            raise ValueError
    except ValueError:
        raise Refuse("loc_property_invalid")
    mutable = [(t, v) for t, v in outer if t != IMMUTABLE]
    allp = mutable + inner
    loc = {}
    for t, rule in LOC_RULES.items():
        vs = [v for tt, v in allp if tt == t]
        if len(vs) > 1 or (vs and not rule(vs[0])):
            raise Refuse("loc_property_invalid")
        loc[t] = vs[0] if vs else None
    for t in (WAI_MANIFEST, WAI_RECEIPT):
        if any(tt == t for tt, _ in mutable) or sum(1 for tt, _ in inner if tt == t) > 1:
            raise Refuse("loc_property_invalid")
    mans = [v for t, v in inner if t == WAI_MANIFEST]
    if not mans:
        raise Refuse("loc_manifest_absent")
    cap = manifest_capability(mans[0])
    if cap not in MAPPED:
        raise Refuse("loc_capability_unmapped")
    if cap != track_cap:
        raise Refuse("loc_capability_mismatch")
    if not one_chunk(cap, payload):
        raise Refuse("loc_payload_form")
    env = envelope(mans[0], payload)
    rcs = [v for t, v in inner if t == WAI_RECEIPT]
    receipt = None
    if rcs:
        receipt = read_receipt(rcs[0])
        if receipt["content_hash"] != content_hash(env):
            raise Refuse("loc_receipt_mismatch")
        receipt = rcs[0]
    return env, loc, receipt


# ---- loc-carriage §4.1: writing ------------------------------------------------------

def write_object(env, loc, receipt):
    if len(env) < 8 or env[:4] != b"WAI1":
        raise Refuse("loc_manifest_invalid")
    ml = int.from_bytes(env[4:8], "little")
    m = env[8:8 + ml]
    pl = int.from_bytes(env[8 + ml:12 + ml], "little")
    p = env[12 + ml:12 + ml + pl]
    if len(m) != ml or len(p) != pl or len(env) != 12 + ml + pl or ml > 0xFFFF:
        raise Refuse("loc_manifest_invalid")
    cap = manifest_capability(m)
    if not one_chunk(cap, p):
        raise Refuse("loc_payload_form")
    for t, v in loc.items():
        if v is not None and not LOC_RULES[t](v):
            raise Refuse("loc_property_invalid")
    inner = [(t, v) for t, v in loc.items() if v is not None] + [(WAI_MANIFEST, m)]
    if receipt is not None:
        if read_receipt(receipt)["content_hash"] != content_hash(env):
            raise Refuse("loc_receipt_mismatch")
        inner.append((WAI_RECEIPT, receipt))
    return encode_pairs([(IMMUTABLE, encode_pairs(inner))]), p


def stated_loc(o):
    s = o["loc"]
    hx = lambda v: bytes.fromhex(v) if v is not None else None
    return {TIMESTAMP: s["timestamp"], TIMESCALE: s["timescale"], FRAME_MARKING: hx(s["video_frame_marking"]),
            AUDIO_LEVEL: s["audio_level"], VIDEO_CONFIG: hx(s["video_config"]), AUDIO_CONFIG: hx(s["audio_config"])}


def check(doc):
    fails, given = [], set()
    for o in doc["objects"]:
        env, props, payload = bytes.fromhex(o["envelope"]), bytes.fromhex(o["properties"]), bytes.fromhex(o["payload"])
        rc = bytes.fromhex(o["receipt"]) if o["receipt"] else None
        try:
            got = read_object(props, payload, o["capability"])
            if got != (env, stated_loc(o), rc):
                fails.append("%s: reads back to another envelope, LOC properties or receipt" % o["name"])
        except Refuse as e:
            fails.append("%s: refused %s" % (o["name"], e.code))
        try:
            if write_object(env, stated_loc(o), rc) != (props, payload):
                fails.append("%s: written otherwise" % o["name"])
        except Refuse as e:
            fails.append("%s: the writer refuses it (%s)" % (o["name"], e.code))
    for c in doc["writer_refused"]:
        try:
            write_object(bytes.fromhex(c["envelope"]), {t: None for t in LOC_RULES}, None)
            got = None
        except Refuse as e:
            got = e.code
        given.add(got)
        if got != c["refuse"]:
            fails.append("writer %s: %s, stated %s" % (c["name"], got, c["refuse"]))
    for c in doc["refused"]:
        try:
            read_object(bytes.fromhex(c["properties"]), bytes.fromhex(c["payload"]), c["capability"])
            got = None
        except Refuse as e:
            got = e.code
        given.add(got)
        if got != c["refuse"]:
            fails.append("reader %s: %s, stated %s" % (c["name"], got, c["refuse"]))
    text = open(os.path.join(HERE, "..", "extensions", "loc-carriage.md"), encoding="utf-8").read()
    at = text.index("\n## 6. Refusals")
    codes = re.findall(r"^\| `([a-z_]+)` \|", text[at:text.index("\n## ", at + 1)], re.M)
    for c in codes:
        if c not in given:
            fails.append("no vector is refused with %s" % c)
    if set(given) - set(codes):
        fails.append("codes given that §6 does not list: %s" % sorted(set(given) - set(codes)))
    return fails


def main(argv):
    if len(argv) != 2:
        print(__doc__)
        return 2
    doc = json.load(open(os.path.join(argv[1], "vectors.json"), encoding="utf-8"))
    fails = check(doc)
    for f in fails:
        print("FAIL", f)
    print("%d objects, %d refused envelopes, %d refused objects: %d failures" % (len(doc["objects"]), len(doc["writer_refused"]), len(doc["refused"]), len(fails)))

    def flip(h, i):
        b = bytearray(bytes.fromhex(h))
        b[i] ^= 1
        return b.hex()

    muts = []

    def mutate(label, f):
        d = copy.deepcopy(doc)
        f(d)
        muts.append((label, d))

    obj = lambda d, n: next(o for o in d["objects"] if o["name"] == n)
    mutate("a byte of a payload", lambda d: obj(d, "opus_unmetered").__setitem__("payload", flip(obj(d, "opus_unmetered")["payload"], 2)))
    mutate("a byte of the manifest in the properties", lambda d: obj(d, "hevc_key_metered").__setitem__("properties", flip(obj(d, "hevc_key_metered")["properties"], 40)))
    mutate("a stated timestamp", lambda d: obj(d, "opus_unmetered")["loc"].__setitem__("timestamp", 961))
    mutate("a stated receipt's figure", lambda d: obj(d, "hevc_key_metered").__setitem__("receipt", flip(obj(d, "hevc_key_metered")["receipt"], 33)))
    mutate("a refusal stated with another code", lambda d: d["refused"][0].__setitem__("refuse", "loc_property_invalid"))
    mutate("the track's capability stated otherwise", lambda d: obj(d, "aac_frame").__setitem__("capability", "wai.audio.opus.packet"))
    ok = True
    for label, d in muts:
        caught = check(d)
        print(("PASS" if caught else "FAIL") + " mutated copy fails: " + label)
        ok &= bool(caught)
    DROPPED.add("timescale")
    caught = check(doc)
    DROPPED.clear()
    print(("PASS" if caught else "FAIL") + " the corpus fails a verifier without the Timescale rule")
    ok &= bool(caught)
    return 0 if not fails and ok else 1


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