#!/usr/bin/env python3
"""Independent verifier of the exact-upscale corpus (extensions/exact-upscale.md §9).

Written from the extension's text. It shares no code with wai-rs, which generates the corpus.
It runs the derivations of §3 with `_int_sr.py`, written from §3 in integer arithmetic of its
own; verifies the origin claims with the synthesis corpus's second implementation
(`synthesis_verify.py`, §7.2 step 2); reads and decodes residuals with the staged-delivery
corpus's (`staged_verify.py`, staged-delivery Appendix A); verifies every signature with WAI's
strict Ed25519 rule (`claim-conformance/_ed25519.py`) and every link hash with BLAKE3 written
from its specification (`claim-conformance/_blake3.py`); reads §8's codes from the extension;
and implements the rest: the records of §4.2 and §6.2, the steps of §7.2, the counts of §6.3.

For every vector it requires the outcome, the reason, the place and the statements the corpus
states. Then it mutates copies of the corpus, each with one defect, and requires every copy to
be refused somewhere.

    python3 -I upscale_verify.py <live-conformance/upscale>

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

HERE = os.path.dirname(os.path.abspath(__file__))
WAI = os.path.normpath(os.path.join(HERE, ".."))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(WAI, "claim-conformance"))
sys.path.insert(0, os.path.join(WAI, "staged-conformance"))
import _int_sr as SR  # noqa: E402
import synthesis_verify as SV  # noqa: E402
import staged_verify as STV  # noqa: E402
import _ed25519 as ED  # noqa: E402
from _blake3 import blake3  # noqa: E402

EXTENSION = os.path.join(WAI, "extensions", "exact-upscale.md")
REGISTRY = os.path.join(WAI, "registry", "capabilities.json")
SIGN_DOMAIN = b"wai:upscale\x01"           # §2
LINK_DOMAIN = b"wai:upscale-link\x01"      # §2
OBJECT_DOMAIN = b"moq-jwp:object\x01"      # jwp-receipts §2.1
CAP = 1 << 26
MADE = ["transmitted", "derived", "asserted", "synthesized"]   # §6.1, in a count's order
CLASSES = ["captured", "reconstructed", "assistive", "synthesized"]
LEAST = {"transmitted": "captured", "derived": "assistive", "asserted": "assistive", "synthesized": "synthesized"}

Refuse, Bad, need, hexb, integer, members = SV.Refuse, SV.Bad, SV.need, SV.hexb, SV.integer, SV.members


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


def opt(v, f):
    return None if v is None else f(v)


# ---------------------------------------------------------------- §2 records

def read(text, kind, names, fields):
    try:
        v = SV.strict_json(text)
        members(v, names)
        need(v["kind"] == kind)
        signer, sig = hexb(v["signer_pubkey"], 32), hexb(v["sig"], 64)
        SV.date_time(v["issued_at"])
        r = fields(v)
    except (Bad, KeyError, TypeError, AttributeError, ValueError):
        raise Refuse("record_malformed")
    bad = ED.verify(signer, SIGN_DOMAIN + SV.jcs({k: x for k, x in v.items() if k != "sig"}).encode("utf-8"), sig)
    if bad:
        raise Refuse(bad)
    r.update(link=blake3(LINK_DOMAIN + SV.jcs(v).encode("utf-8")))
    return r


# ---------------------------------------------------------------- §4.2 the upscale record

UPSCALE = ["kind", "object", "input", "derivation", "prior", "upscaled", "residual", "output", "target", "issued_at", "signer_pubkey", "sig"]


def upscale_fields(v):
    members(v["input"], ["width", "height", "samples"])
    u = {"object": hexb(v["object"], 32), "w": integer(v["input"]["width"]), "h": integer(v["input"]["height"]),
         "input": hexb(v["input"]["samples"], 32), "derivation": SV.one_of(v["derivation"], [SR.SR2, SR.SR4]),
         "prior": hexb(v["prior"], 32), "upscaled": hexb(v["upscaled"], 32), "output": hexb(v["output"], 32)}
    need(u["w"] >= 1 and u["h"] >= 1)
    r, t = v["residual"], v["target"]
    if r is not None:
        members(r, ["sha256", "q"])
        u["residual"] = (hexb(r["sha256"], 32), integer(r["q"]))
        need(1 <= u["residual"][1] <= 65535)
    else:
        u["residual"] = None
    if t is not None:
        members(t, ["samples", "bound"])
        u["target"] = (hexb(t["samples"], 32), integer(t["bound"]))
        need(u["target"][1] <= 255)
    else:
        u["target"] = None
    f = SR.FACTOR[u["derivation"]]
    need(f * f * u["w"] * u["h"] * 3 <= CAP)
    need((r is None) == (t is None))
    need(r is not None or u["output"] == u["upscaled"])
    return u


# ---------------------------------------------------------------- §6.2 the presentation statement

STATEMENT = ["kind", "session", "from", "to", "frames", "totals", "issued_at", "signer_pubkey", "sig"]
ENTRY = ["at", "made", "width", "height", "samples", "object", "claim", "upscale", "made_by", "pixels"]


def count(v):
    members(v, MADE)
    return [integer(v[m]) for m in MADE]


def entry_fields(e):
    members(e, ENTRY)
    f = {"at": integer(e["at"]), "made": SV.one_of(e["made"], MADE), "w": integer(e["width"]), "h": integer(e["height"]),
         "samples": hexb(e["samples"], 32), "object": opt(e["object"], lambda x: hexb(x, 32)), "claim": hexb(e["claim"], 32),
         "upscale": opt(e["upscale"], lambda x: hexb(x, 32)), "made_by": opt(e["made_by"], lambda x: hexb(x, 32)),
         "pixels": count(e["pixels"])}
    need(f["w"] >= 1 and f["h"] >= 1 and f["w"] * f["h"] <= 1 << 40)
    need((f["object"] is None) == (f["made"] == "synthesized"))
    need(f["upscale"] is None or f["made"] in ("transmitted", "derived"))
    if f["made"] == "transmitted":
        need(f["made_by"] is None)
    elif f["made"] in ("derived", "synthesized"):
        need(f["made_by"] is not None)
    return f


def statement_fields(v):
    members(v["totals"], ["frames", "pixels"])
    need(isinstance(v["frames"], list) and 1 <= len(v["frames"]) <= 4096)
    s = {"session": hexb(v["session"], 32), "from": integer(v["from"]), "to": integer(v["to"]),
         "frames": [entry_fields(e) for e in v["frames"]],
         "totals": (count(v["totals"]["frames"]), count(v["totals"]["pixels"]))}
    need(s["from"] < s["to"])
    need(all(s["from"] <= e["at"] < s["to"] for e in s["frames"]))
    need(all(a["at"] < b["at"] for a, b in zip(s["frames"], s["frames"][1:])))
    return s


# ---------------------------------------------------------------- §4.1 the residual

def residual_reads(payload, q, ow, oh):
    """§4.1: a WIR1 payload of the RGB8 unit kind, identity, one unit of the output's size, with
    the record's q, that decodes to its end."""
    try:
        w = STV.parse_wir1(payload, "rgb8")
        if w["pred"] != 0 or w["n_units"] != 1 or w["d"] != (oh, ow, 3) or w["q"] != q:
            return None
        STV.check_streams(w)
        return w
    except STV.Refused:
        return None


# ---------------------------------------------------------------- §7 verification

def verify(case, classes, priors_by_name):
    """§7.2: the statements of §7.3, or Refuse with its reason and place."""
    def each(key, kind, names, fields, place):
        out = []
        for i, t in enumerate(case[key]):
            try:
                out.append(read(t, kind, names, fields))
            except Refuse as e:
                raise Refuse(e.reason, [place, i])
        return out

    # Step 1.
    upscales = each("upscales", "wai-upscale", UPSCALE, upscale_fields, "upscale")
    statements = each("presentations", "wai-presentation", STATEMENT, statement_fields, "presentation")

    # Step 2: the claims, by synthesis's second implementation.
    pol = case["policy"]
    prior_bytes = [priors_by_name[n] for n in case["priors"]]
    syn = {"claims": case["claims"], "markings": [], "detections": [], "reviews": [], "label_policy": None, "player": [],
           "objects": case["objects"], "decoded": case["decoded"], "priors": [p.hex() for p in prior_bytes], "unmarked": [],
           "room": None, "policy": {"dual_layer": pol["dual_layer"], "assistive": pol["assistive"], "second_layer": pol["second_layer"]}}
    try:
        syn_st = SV.verify(syn, classes)
    except Refuse as e:
        raise Refuse(e.reason, ["synthesis", e.at])
    st = {"claims": [{"link": c["link"], "class": c["class"], "assistive": c["assistive"]} for c in syn_st["claims"]],
          "upscales": [], "presentations": []}
    claims = {c["link"]: c for c in (SV.read(t, "wai-origin", SV.ORIGIN, SV.origin_fields) for t in case["claims"])}

    decoded = {bytes.fromhex(d["object"]): bytes.fromhex(d["samples"]) for d in case["decoded"]}
    priors = {sha(p): p for p in prior_bytes}
    residuals = {sha(bytes.fromhex(r)): bytes.fromhex(r) for r in case["residuals"]}
    targets = {sha(bytes.fromhex(t)): bytes.fromhex(t) for t in case["targets"]}

    # Step 4: the upscale records.
    for i, u in enumerate(upscales):
        def no(reason):
            raise Refuse(reason, ["upscale", i])
        prior = priors.get(u["prior"])
        if prior is not None:
            try:
                SR.read_prior(prior)
            except SR.Unreadable:
                no("prior_invalid")
        px = decoded.get(u["object"])
        if px is not None and (len(px) != u["w"] * u["h"] * 3 or sha(px) != u["input"]):
            no("input_samples_mismatch")
        f = SR.FACTOR[u["derivation"]]
        ow, oh = f * u["w"], f * u["h"]
        upscaled = None
        if prior is not None and px is not None:
            try:
                upscaled = SR.derive(u["derivation"], prior, px, u["w"], u["h"])
            except SR.Unreadable:
                no("upscale_mismatch")
            if sha(upscaled) != u["upscaled"]:
                no("upscale_mismatch")
        output = upscaled
        if u["residual"] is not None:
            output = None
            payload = residuals.get(u["residual"][0])
            if payload is not None:
                w = residual_reads(payload, u["residual"][1], ow, oh)
                if w is None:
                    no("residual_invalid")
                if upscaled is not None:
                    output = bytes(STV.apply_unit(w, 0, list(upscaled), (oh, ow, 3)))
                    if sha(output) != u["output"]:
                        no("output_mismatch")
        diff = None
        if u["target"] is not None:
            t, bound = u["target"]
            if bound == 0 and u["output"] != t:
                no("residual_short")
            target = targets.get(t)
            if output is not None and target is not None:
                if len(target) != len(output):
                    no("residual_short")
                diff = max(abs(a - b) for a, b in zip(output, target))
                if diff > bound:
                    no("residual_short")
        if pol["upscale"] == "require-proof" and output is None:
            no("upscale_not_proved")
        st["upscales"].append({"link": u["link"].hex(), "status": "proved" if output is not None else "not-recomputed", "max_difference": diff})

    # Step 5: the presentation statements.
    by_link = {u["link"]: u for u in upscales}
    for i, s in enumerate(statements):
        frames, pixels = [0] * 4, [0] * 4
        for j, f in enumerate(s["frames"]):
            check_entry(f, claims, by_link, decoded, ["presentation", i, "frame", j])
            frames[MADE.index(f["made"])] += 1
            pixels = [a + b for a, b in zip(pixels, f["pixels"])]
        if (frames, pixels) != s["totals"]:
            raise Refuse("fractions_mismatch", ["presentation", i])
        st["presentations"].append({"link": s["link"].hex(), "totals": {"frames": dict(zip(MADE, frames)), "pixels": dict(zip(MADE, pixels))}})
    return st


def check_entry(f, claims, upscales, decoded, at):
    """§7.2 step 5, checks 1 to 7."""
    def no(reason):
        raise Refuse(reason, at)
    c = claims.get(f["claim"])
    if c is None:
        no("claim_unknown")
    u = None
    if f["upscale"] is not None:
        u = upscales.get(f["upscale"])
        if u is None:
            no("upscale_unknown")
    made = f["made"]
    if made == "transmitted" and u is None:
        subject_ok = c["subject"] == ("object", f["object"])
    else:
        subject_ok = c["subject"] == ("samples", f["samples"])
    if not subject_ok:
        no("frame_subject_mismatch")
    if CLASSES.index(c["class"]) < CLASSES.index(LEAST[made]):
        no("made_understated")
    e = c["evidence"]
    one_input = len(c["inputs"]) == 1 and c["inputs"][0]["content_hash"] == f["object"]
    if made == "transmitted":
        ok = u is None or (u["residual"] is not None and u["object"] == f["object"] and u["output"] == f["samples"]
                           and c["step"] == "convert" and one_input)
    elif made == "derived":
        ok = (c["step"] == "assist" and e is not None and e["tier"] == "replicate" and one_input and e["derivation"] in SR.FACTOR
              and SR.FACTOR[e["derivation"]] * e["width"] == f["w"] and SR.FACTOR[e["derivation"]] * e["height"] == f["h"]
              and (u is None or (u["object"] == f["object"] and u["derivation"] == e["derivation"] and u["upscaled"] == f["samples"])))
    elif made == "asserted":
        ok = c["step"] == "assist" and e is not None and e["tier"] == "attest" and one_input
    else:
        ok = c["step"] == "generate"
    if not ok:
        no("made_mismatch")
    if made == "derived":
        prior = SV.evidence_prior(e)
        if f["made_by"] != prior or (u is not None and u["prior"] != f["made_by"]):
            no("prior_mismatch")
    elif made in ("asserted", "synthesized"):
        if f["made_by"] is not None and f["made_by"] not in [m["digest"] for m in c["models"]]:
            no("made_by_mismatch")
    if made == "transmitted" and u is None and f["object"] in decoded and sha(decoded[f["object"]]) != f["samples"]:
        no("presented_samples_mismatch")
    # §6.3.
    n_all = f["w"] * f["h"]
    n = min(e["width"] * e["height"], n_all) if e is not None else 0
    want = {"transmitted": [n_all, 0, 0, 0], "derived": [n, n_all - n, 0, 0], "asserted": [n, 0, n_all - n, 0],
            "synthesized": [0, 0, 0, n_all]}[made]
    if f["pixels"] != want:
        no("pixels_mismatch")


# ---------------------------------------------------------------- the extension's table

with open(EXTENSION, encoding="utf-8") as _f:
    _TEXT = _f.read()
REASONS = [l.split("`")[1] for l in SV.section(_TEXT, "## 8. Reasons").splitlines() if l.startswith("| `")]
ACCEPTED = {"sr2_generic", "sr4_seeded", "residual_exact", "residual_bound", "per_title", "span"}


def outcome(v, classes, priors):
    try:
        return {"outcome": "accept", "statements": verify(v, classes, priors)}
    except Refuse as e:
        return {"outcome": "refuse", "reason": e.reason, "at": e.at}


def stated(v):
    return {k: v[k] for k in ("outcome", "reason", "at", "statements") if k in v}


def check(doc, classes, priors):
    """Failures of a corpus: a prior that is not the one the document lists, or an outcome
    other than the one stated."""
    fails = []
    listed = {p["file"][len("priors/"):]: p["sha256"] for p in doc["priors"]}
    if listed != {n: sha(b).hex() for n, b in priors.items()}:
        fails.append("the priors differ from the ones vectors.json lists")
    for v in doc["vectors"]:
        try:
            got = outcome(v, classes, priors)
        except Exception as e:  # noqa: BLE001 — a malformed copy fails, whatever it breaks
            got = {"outcome": "unreadable: %r" % (e,)}
        if got != stated(v):
            fails.append("%s: %s" % (v["name"], json.dumps(got)[:300]))
    return fails


def flip_hex_digit(s, at):
    return s[:at] + ("0" if s[at] != "0" else "1") + s[at + 1:]


def mutations(doc, priors):
    """Copies of the corpus, each with one defect a conforming verifier must catch."""
    vn = {v["name"]: k for k, v in enumerate(doc["vectors"])}

    def mutate(label, f, p=None):
        d = copy.deepcopy(doc)
        f(d["vectors"])
        return label, d, p or priors

    def prior_byte(name, at):
        p = dict(priors)
        b = bytearray(p[name])
        b[at] ^= 0x01
        p[name] = bytes(b)
        return p

    def target_sample(vs):
        t = vs[vn["residual_bound"]]["targets"][0]
        b = bytearray.fromhex(t)
        b[10] = (b[10] + 3) % 256
        vs[vn["residual_bound"]]["targets"][0] = b.hex()

    def relisted(label, name, at):
        """A prior's weight changed and the document's listing updated to the changed bytes, so
        that only the records, which name the old digest, can tell."""
        p = prior_byte(name, at)
        d = copy.deepcopy(doc)
        for e in d["priors"]:
            if e["file"] == "priors/" + name:
                e["sha256"] = sha(p[name]).hex()
        return label, d, p

    first_refusal = next(k for k, v in enumerate(doc["vectors"]) if v["outcome"] == "refuse")
    return [
        relisted("a weight of the generic prior, the listing updated", "generic.wsr", 200),
        relisted("a weight of the per-title prior, the listing updated", "title.wsr", 300),
        mutate("a weight of the generic prior", lambda vs: None, prior_byte("generic.wsr", 200)),
        mutate("a weight of the per-title prior", lambda vs: None, prior_byte("title.wsr", 300)),
        mutate("a weight of the seeded prior", lambda vs: None, prior_byte("seeded.wsr", 20)),
        mutate("a byte of the exact residual's stream",
               lambda vs: vs[vn["residual_exact"]]["residuals"].__setitem__(0, flip_hex_digit(vs[vn["residual_exact"]]["residuals"][0], -3))),
        mutate("a sample of the bounded residual's target three further off", target_sample),
        mutate("a sample of scene one's decoded picture",
               lambda vs: vs[vn["sr2_generic"]]["decoded"][0].update(samples=flip_hex_digit(vs[vn["sr2_generic"]]["decoded"][0]["samples"], 1))),
        mutate("the span's derived pixels stated one higher",
               lambda vs: vs[vn["span"]]["statements"]["presentations"][0]["totals"]["pixels"].update(derived=145)),
        mutate("the exact residual's record stated not recomputed",
               lambda vs: vs[vn["residual_exact"]]["statements"]["upscales"][0].update(status="not-recomputed")),
        mutate("the bounded residual's largest difference stated as 0",
               lambda vs: vs[vn["residual_bound"]]["statements"]["upscales"][0].update(max_difference=0)),
        mutate("the per-title prior withheld", lambda vs: vs[vn["per_title"]].update(priors=[])),
        mutate("a digit of the span statement's signature",
               lambda vs: vs[vn["span"]]["presentations"].__setitem__(0, json.dumps(dict(json.loads(vs[vn["span"]]["presentations"][0]),
                                                                                        sig=flip_hex_digit(json.loads(vs[vn["span"]]["presentations"][0])["sig"], 4))))),
        mutate("the first refusal stated at another place", lambda vs: vs[first_refusal].update(at=["upscale", 7])),
    ]


def main(corpus_dir):
    fails = []
    st = ED._self_test()
    print(("PASS" if not st else "FAIL") + " Ed25519 strict rule self-test " + ",".join(st))
    if st:
        fails.append("ed25519 self-test")
    with open(REGISTRY, encoding="utf-8") as f:
        classes = {c["capability"]: c["replicate"] for c in json.load(f)["capabilities"]}
    with open(os.path.join(corpus_dir, "vectors.json"), encoding="utf-8") as f:
        doc = json.load(f)
    pdir = os.path.join(corpus_dir, "priors")
    priors = {}
    for n in sorted(os.listdir(pdir)):
        with open(os.path.join(pdir, n), "rb") as f:
            priors[n] = f.read()
    reasons_seen = set()
    for v in doc["vectors"]:
        try:
            got = outcome(v, classes, priors)
        except Exception as e:  # noqa: BLE001
            got = {"outcome": "unreadable: %r" % (e,)}
        ok = got == stated(v)
        if v.get("reason"):
            reasons_seen.add(v["reason"])
        print(("PASS" if ok else "FAIL") + " vector %s: %s" % (v["name"], got["outcome"] + (" %s at %s" % (got["reason"], got["at"]) if "reason" in got else "")))
        if not ok:
            fails.append("%s: %s, the corpus states %s" % (v["name"], json.dumps(got)[:400], json.dumps(stated(v))[:400]))
    if any("priors differ" in t for t in check(doc, classes, priors)):
        fails.append("the priors differ from the ones vectors.json lists")
    missing = [r for r in REASONS if r not in reasons_seen]
    accepted = {v["name"] for v in doc["vectors"] if v["outcome"] == "accept"}
    if missing:
        fails.append("no vector refused with %s" % ", ".join(missing))
    if not ACCEPTED <= accepted:
        fails.append("vectors not accepted: %s" % ", ".join(sorted(ACCEPTED - accepted)))
    print("%s every reason of §8 refuses a vector, and every case of §9 is accepted" % ("PASS" if not missing and ACCEPTED <= accepted else "FAIL"))
    sidecar = os.path.join(corpus_dir, "sha256.json")
    with open(sidecar, encoding="utf-8") as f:
        side = json.load(f)
    outs = {}
    for v in doc["vectors"]:
        if v["outcome"] == "accept":
            for i, t in enumerate(v["upscales"]):
                outs["%s/upscale/%d/output" % (v["name"], i)] = json.loads(t)["output"]
    ok = {k: e["sha256"] for k, e in side.items()} == outs
    print(("PASS" if ok else "FAIL") + " sha256.json names the output of every accepted upscale record")
    if not ok:
        fails.append("sha256.json differs from the accepted records' outputs")
    for label, d, p in mutations(doc, priors):
        caught = check(d, classes, p)
        print(("PASS" if caught else "FAIL") + " mutated copy fails: %s%s" % (label, " (%s)" % caught[0][:100] if caught else ""))
        if not caught:
            fails.append("a mutated copy verifies: %s" % label)
    for f in fails:
        print("FAIL " + f)
    print(("PASS" if not fails else "FAIL") + " exact-upscale corpus: %d vectors" % len(doc["vectors"]))
    return 1 if fails else 0


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("usage: upscale_verify.py <live-conformance/upscale>", file=sys.stderr)
        sys.exit(2)
    sys.exit(main(sys.argv[1]))