#!/usr/bin/env python3
"""project_verify.py: a second reader of wai.lock and of a build record.

Written from wai/schema/README.md, the two schemas beside it and the texts they cite, sharing no
code with wai-rs, so the formats are not defined only by the Rust that writes them. Standard
library only; run it isolated (python3 -I).

    project_verify.py vectors <vectors.json>
        each case's bytes read: an accepted one is canonical and has its stated digest, and a
        refused one is refused with its stated code; and the watch room's tone, made by its
        integer rule, has the digest the corpus pins

    project_verify.py selftest <vectors.json>
        drop each of this reader's rules in turn, and require the corpus to catch it

    project_verify.py project <dir> --wai <wai-dir>
        recompute every digest <dir>/wai.lock pins from the files it pins, and hold the file to
        them byte for byte; and, where <dir>/dist/wai-build.json is there, hold the build record to
        the project, the lock and every file of the tree

Exit 0 when everything holds, 1 with one line a finding otherwise, 2 on usage or for a suite this
reader does not know.
"""
import hashlib
import json
import os
import re
import struct
import sys

def plain(p):
    """A path as a person writes it: on Windows, a verbatim path (\\\\?\\C:\\...) takes no "/" and no
    "..", so its prefix goes."""
    if p.startswith("\\\\?\\") and not p.startswith("\\\\?\\UNC\\"):
        return p[4:]
    return p


HERE = os.path.dirname(os.path.abspath(plain(__file__)))
SCHEMA = "https://wai.transaction.science/schema/wai.lock/1"
SCHEMA_OF = "https://wai.transaction.science/schema/wai.lock/"
RECORD_SCHEMA = "https://wai.transaction.science/schema/wai-build.json/1"
MEMBERS = {"$schema", "wai", "toolchain", "template", "registry", "priors", "policy", "conformance"}
HEX = re.compile(r"[0-9a-f]{64}\Z")

# A selftest's run drops one rule; None keeps them all.
DROPPED = None
RULES = ("repeated", "json", "version", "unknown", "shape", "canonical")


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


class Repeated(Exception):
    pass


def distinct(pairs):
    keys = [k for k, _ in pairs]
    if DROPPED != "repeated" and len(keys) != len(set(keys)):
        raise Repeated(keys)
    return dict(pairs)


def jcs(v):
    """RFC 8785 over strings, booleans, null, whole numbers, arrays and objects."""
    if isinstance(v, bool) or v is None:
        return json.dumps(v)
    if isinstance(v, int):
        return str(v)
    if isinstance(v, float):
        raise ValueError("a number that is not whole")
    if isinstance(v, str):
        return json.dumps(v, ensure_ascii=False)
    if isinstance(v, list):
        return "[" + ",".join(jcs(x) for x in v) + "]"
    if isinstance(v, dict):
        keys = sorted(v, key=lambda k: k.encode("utf-16-be"))
        return "{" + ",".join(json.dumps(k, ensure_ascii=False) + ":" + jcs(v[k]) for k in keys) + "}"
    raise ValueError(repr(v))


def whole_numbers(v):
    """A whole number however it is written, as wai.json's reader takes it."""
    if isinstance(v, float) and v == v and v not in (float("inf"), float("-inf")) and v == int(v):
        return int(v)
    if isinstance(v, list):
        return [whole_numbers(x) for x in v]
    if isinstance(v, dict):
        return {k: whole_numbers(x) for k, x in v.items()}
    return v


def is_hex(v):
    return isinstance(v, str) and HEX.match(v) is not None


def is_text(v):
    return isinstance(v, str) and v != ""


def exact(v, required, optional=()):
    """Every problem of an object with these members: (required, optional) as {name: rule}."""
    if not isinstance(v, dict):
        return ["lock-member-invalid"]
    out = []
    rules = dict(required)
    rules.update(dict(optional))
    for k in v:
        if k not in rules and DROPPED != "unknown":
            out.append("lock-unknown-member")
    for k, rule in dict(required).items():
        if k not in v:
            out.append("lock-member-invalid")
    for k, x in v.items():
        if k in rules and not rules[k](x):
            out.append("lock-member-invalid")
    return out


def read_lock(b):
    """A lock's bytes: (the object, []) or (None, [each refusal's code])."""
    try:
        v = json.loads(b.decode("utf-8"), object_pairs_hook=distinct)
    except Repeated:
        return None, ["lock-repeated-member"]
    except (ValueError, RecursionError):
        if DROPPED == "json":
            return {}, []
        return None, ["lock-not-json"]
    if not isinstance(v, dict):
        if DROPPED == "json":
            return {}, []
        return None, ["lock-not-json"]
    if DROPPED != "version":
        w = v.get("wai")
        if isinstance(w, (int, float)) and not isinstance(w, bool) and w == int(w) and w > 1:
            return None, ["lock-version"]
        s = v.get("$schema")
        if isinstance(s, str) and s.startswith(SCHEMA_OF) and s[len(SCHEMA_OF):].isdigit() and int(s[len(SCHEMA_OF):]) > 1:
            return None, ["lock-version"]
    out = []
    for k in v:
        if k not in MEMBERS and DROPPED != "unknown":
            out.append("lock-unknown-member")
    if DROPPED != "shape":
        if v.get("wai") != 1 or "wai" not in v:
            out.append("lock-member-invalid")
        if v.get("$schema") != SCHEMA:
            out.append("lock-member-invalid")
    catalogue = lambda x: isinstance(x, dict) and set(x) == {"sha256"} and is_hex(x["sha256"])
    shapes = []
    if "toolchain" in v or DROPPED != "shape":
        shapes.append(exact(v.get("toolchain"), {"version": is_text, "catalogue": catalogue}))
    if "registry" in v or DROPPED != "shape":
        shapes.append(exact(v.get("registry"), {"sha256": is_hex}))
    if "template" in v:
        shapes.append(exact(v["template"], {"name": is_text}, {"sha256": is_hex}))
    if "policy" in v:
        shapes.append(exact(v["policy"], {"schema": is_text, "blake3": is_hex}))
    if "priors" in v:
        m = v["priors"]
        if not isinstance(m, dict) or not m:
            shapes.append(["lock-member-invalid"])
        else:
            shapes.extend(exact(x, {"sha256": is_hex}, {"form": is_text}) for x in m.values())
    if "conformance" in v:
        m = v["conformance"]
        digest = lambda x: isinstance(x, str) and x.startswith("sha256:") and is_hex(x[7:])
        if not isinstance(m, dict) or not m:
            shapes.append(["lock-member-invalid"])
        else:
            shapes.extend(exact(x, {"digest": digest}) for x in m.values())
    for s in shapes:
        for code in s:
            if code == "lock-unknown-member" or DROPPED != "shape":
                out.append(code)
    if DROPPED != "canonical":
        try:
            if jcs(v).encode("utf-8") != b:
                out.append("lock-not-canonical")
        except ValueError:
            out.append("lock-not-canonical")
    return (None, out) if out else (v, [])


def run_vectors(path):
    with open(path, encoding="utf-8") as f:
        doc = json.load(f)
    fails = []
    if DROPPED is None and sha256(tone()) != doc["tone"]["sha256"]:
        fails.append(f"the tone: its SHA-256 is {sha256(tone())}, the corpus says {doc['tone']['sha256']}")
    for c in doc["cases"]:
        b = c["bytes"].encode("utf-8")
        v, codes = read_lock(b)
        if "refused" in c:
            if c["refused"] not in codes:
                fails.append(f"{c['name']}: refused with {codes}, not {c['refused']}")
        else:
            if codes:
                fails.append(f"{c['name']}: refused with {codes}")
            elif sha256(b) != c["sha256"]:
                fails.append(f"{c['name']}: its SHA-256 is {sha256(b)}, the corpus says {c['sha256']}")
            elif DROPPED is None and jcs(v).encode("utf-8") != b:
                fails.append(f"{c['name']}: it does not write back byte for byte")
    return fails, len(doc["cases"])


def listing_digest(lines):
    """SHA-256 over `<hex>  <path>\\n` sorted by path as bytes."""
    text = "".join(f"{h}  {p}\n" for p, h in sorted(lines, key=lambda x: x[0].encode("utf-8")))
    return sha256(text.encode("utf-8"))


def walk(root, skip_top=()):
    out = []
    for d, dirs, files in os.walk(root):
        rel = os.path.relpath(d, root).replace(os.sep, "/")
        rel = "" if rel == "." else rel + "/"
        dirs[:] = [x for x in dirs if rel + x not in skip_top]
        out.extend(rel + f for f in files)
    return out


def read_file(path):
    with open(path, "rb") as f:
        return f.read()


def tone_sample(p, rate):
    """The sample at phase p of a cycle of `rate` samples: a sine of peak 6000, each half cycle by
    the rational form of Bhaskara I, 16 x (1 - x) / (5 - 4 x (1 - x)), rounded half up. Integers
    alone, so it is the same value on every host."""
    half = rate // 2
    q = p % half
    u = q * (half - q)
    num, den = 16 * u, 5 * half * half - 4 * u
    v = (2 * 6000 * num + den) // (2 * den)
    return v if p < half else -v


def tone():
    """A watch room's sample sound: twenty seconds of a tone that steps in pitch each second."""
    rate, seconds = 8000, 20
    n = rate * seconds
    b = bytearray(b"RIFF" + struct.pack("<I", 36 + 2 * n) + b"WAVEfmt ")
    b += struct.pack("<IHHIIHH", 16, 1, 1, rate, 2 * rate, 2, 16) + b"data" + struct.pack("<I", 2 * n)
    for i in range(n):
        f = 220 + 20 * ((i // rate) % 12)
        b += struct.pack("<h", tone_sample(f * i % rate, rate))
    return bytes(b)


def template_digest(wai, name):
    """Each file of the example but its test/, and what `wai create` makes for it."""
    root = os.path.join(wai, "examples", name)
    if not os.path.isdir(root):
        return None
    lines = [(p, sha256(read_file(os.path.join(root, p)))) for p in walk(root, skip_top=("test",))]
    if name == "watch-together":
        lines.append(("tone.wav", sha256(tone())))
    return listing_digest(lines)


def prior_pin(path):
    """SPEC §3.1: a directory's listing, every name that starts with a dot and the *.sha256 files
    directly in it left out; or one file, with no form."""
    if os.path.isdir(path):
        names = [p for p in walk(path) if not any(s.startswith(".") for s in p.split("/")) and not ("/" not in p and p.endswith(".sha256"))]
        return {"form": "wai.prior.set", "sha256": listing_digest([(p, sha256(read_file(os.path.join(path, p)))) for p in names])}
    return {"sha256": sha256(read_file(path))}


def suite_files(wai, suite):
    """The files of a suite's corpus, under wai/."""
    one = {
        "verifier/claim": "claim-conformance",
        "verifier/jwp": "jwp-conformance",
        "verifier/signer": "signer-conformance",
        "verifier/conform-result": "conform-result-conformance",
        "sink/native-sink": "live-conformance/native-sink",
        "sink/browser": "live-conformance/browser",
        "verifier/pq": "pq-conformance",
    }
    if suite in one:
        return [one[suite] + "/vectors.json"]
    if suite == "relay/moqt":
        return ["moqt-conformance/cases.json"]
    if suite == "decoder/neural":
        doc = json.loads(read_file(os.path.join(wai, "neural-conformance/sha256.json")))
        return ["neural-conformance/sha256.json"] + [f"neural-conformance/vectors/{k}/{f}" for k in doc for f in ("int_synth.json", "latent.bin")]
    if suite == "cache/cache":
        doc = json.loads(read_file(os.path.join(wai, "http-conformance/objects.json")))
        out = ["cache-conformance/tests.json", "cache-conformance/driver.py", "claim-conformance/_blake3.py",
               "wai-rs/tests/vectors/int_mlicpp/model.wmb", "http-conformance/objects.json"]
        return out + [f"http-conformance/{o['path']}{s}" for o in doc.values() for s in ("", ".hp")]
    if suite in ("decoder/floor", "decoder/video-floor"):
        doc = json.loads(read_file(os.path.join(wai, "floor-conformance/cases.json")))
        want = suite.split("/")[1]
        return ["floor-conformance/cases.json"] + [f"floor-conformance/{c['envelope']}" for c in doc["cases"] if c["set"] == want]
    if suite == "decoder/score":
        doc = json.loads(read_file(os.path.join(wai, "score-conformance/expected_hashes.json")))
        return ["score-conformance/expected_hashes.json"] + [f"score-conformance/vectors/{k}" for k in doc]
    if suite == "packager/policy":
        out = ["pack-conformance/policy/expected.json", "pack-conformance/refusals/expected.json"]
        doc = json.loads(read_file(os.path.join(wai, out[0])))
        out += [f"pack-conformance/policy/{c['file']}" for c in doc["cases"] if "overrides" not in c and "digest" in c]
        doc = json.loads(read_file(os.path.join(wai, out[1])))
        out += [f"pack-conformance/refusals/{c['file']}" for c in doc["cases"] if c.get("file", "").startswith("policy/") and "reason" in c]
        return out
    if suite == "origin/http":
        doc = json.loads(read_file(os.path.join(wai, "http-conformance/objects.json")))
        out = ["http-conformance/objects.json"]
        for o in doc.values():
            if isinstance(o, dict) and o.get("path", "").startswith("objects/"):
                out.append("http-conformance/" + o["path"])
                if os.path.isfile(os.path.join(wai, "http-conformance", o["path"] + ".hp")):
                    out.append("http-conformance/" + o["path"] + ".hp")
        return out
    raise SystemExit(f"project_verify: this reader does not know the suite {suite}")


def corpus_digest(wai, suite):
    """conformance-results §2.1."""
    lines = sorted({f"{sha256(read_file(os.path.join(wai, p)))}  {p}\n" for p in suite_files(wai, suite)})
    return "sha256:" + sha256("".join(lines).encode("utf-8"))


def policy_pin(path):
    """packaging §3, by the packager's own second implementation."""
    sys.path.insert(0, os.path.join(HERE, "..", "pack-conformance"))
    import pack_verify as PV
    P = PV.parse_policy_bytes(read_file(path))
    return {"schema": "wai.pack.policy/1", "blake3": PV.digest(P)}


def cargo_version(wai):
    for line in read_file(os.path.join(wai, "wai-rs/Cargo.toml")).decode("utf-8").splitlines():
        m = re.match(r'version = "([^"]+)"', line)
        if m:
            return m.group(1)
    raise SystemExit("project_verify: wai-rs/Cargo.toml names no version")


def project_expected(d, wai, project):
    """What the lock of this project says, worked out from the files it pins."""
    want = {
        "$schema": SCHEMA,
        "wai": 1,
        "toolchain": {"version": cargo_version(wai), "catalogue": {"sha256": sha256(read_file(os.path.join(wai, "catalogue/ops.json")))}},
        "registry": {"sha256": sha256(read_file(os.path.join(wai, "registry/capabilities.json")))},
    }
    if "template" in project:
        digest = template_digest(wai, project["template"])
        want["template"] = {"name": project["template"]} if digest is None else {"name": project["template"], "sha256": digest}
    if project.get("priors"):
        want["priors"] = {k: prior_pin(os.path.join(d, v)) for k, v in project["priors"].items()}
    if "package" in project:
        want["policy"] = policy_pin(os.path.join(d, project["package"]["policy"]))
    suites = project.get("conform", {}).get("suites")
    if suites:
        want["conformance"] = {s: {"digest": corpus_digest(wai, s)} for s in suites}
    return want


def check_record(d, project_text, lock_bytes, lock, fails):
    dist = os.path.join(d, "dist")
    raw = read_file(os.path.join(dist, "wai-build.json"))
    r = json.loads(raw.decode("utf-8"), object_pairs_hook=distinct)
    if jcs(r).encode("utf-8") != raw:
        fails.append("wai-build.json is not RFC 8785 on one line")
    if r.get("$schema") != RECORD_SCHEMA or r.get("wai") != 1:
        fails.append("wai-build.json does not name its schema and version")
    want = sha256(jcs(whole_numbers(json.loads(project_text, object_pairs_hook=distinct))).encode("utf-8"))
    if r["project"]["sha256"] != want:
        fails.append(f"the record names the project {r['project']['sha256']}, and wai.json's canonical digest is {want}")
    if r["lock"]["sha256"] != sha256(lock_bytes):
        fails.append(f"the record names the lock {r['lock']['sha256']}, and wai.lock's digest is {sha256(lock_bytes)}")
    if r["toolchain"] != lock["toolchain"]:
        fails.append("the record's toolchain is not the lock's")
    paths = [f["path"] for f in r["files"]]
    if paths != sorted(paths, key=lambda p: p.encode("utf-8")):
        fails.append("the record's files are not sorted by path as bytes")
    on_disk = sorted(p for p in walk(dist) if p != "wai-build.json")
    if sorted(paths) != on_disk:
        fails.append(f"the record lists {sorted(paths)}, and the tree holds {on_disk}")
    for f in r["files"]:
        p = os.path.join(dist, f["path"])
        if os.path.isfile(p):
            b = read_file(p)
            if sha256(b) != f["sha256"] or len(b) != f["bytes"]:
                fails.append(f"{f['path']}: not the bytes the record names")


def check_project(d, wai):
    fails = []
    project_text = read_file(os.path.join(d, "wai.json")).decode("utf-8")
    project = json.loads(project_text, object_pairs_hook=distinct)
    lock_bytes = read_file(os.path.join(d, "wai.lock"))
    lock, codes = read_lock(lock_bytes)
    if codes:
        return [f"wai.lock is refused: {sorted(set(codes))}"]
    want = project_expected(d, wai, project)
    for m in sorted(set(want) | set(lock)):
        if want.get(m) != lock.get(m):
            fails.append(f"{m}: wai.lock says {lock.get(m)}, and this reader works out {want.get(m)}")
    if not fails and jcs(want).encode("utf-8") != lock_bytes:
        fails.append("wai.lock is not, byte for byte, the canonical form of what this reader works out")
    if os.path.isfile(os.path.join(d, "dist", "wai-build.json")):
        check_record(d, project_text, lock_bytes, lock, fails)
    return fails


def main(argv):
    global DROPPED
    argv = [plain(a) for a in argv]
    if len(argv) >= 2 and argv[0] == "vectors":
        fails, n = run_vectors(argv[1])
        for f in fails:
            print(f"FAIL {f}")
        if not fails:
            print(f"project_verify: every case of {argv[1]} holds")
        return 1 if fails else 0
    if len(argv) >= 2 and argv[0] == "selftest":
        fails, _ = run_vectors(argv[1])
        if fails:
            print(f"FAIL the corpus does not hold with every rule on: {fails}")
            return 1
        uncaught = []
        for rule in RULES:
            DROPPED = rule
            fails, _ = run_vectors(argv[1])
            if not fails:
                uncaught.append(rule)
        DROPPED = None
        for rule in uncaught:
            print(f"FAIL with the rule {rule} dropped, every case still holds")
        if not uncaught:
            print("project_verify: the corpus catches each rule dropped in turn")
        return 1 if uncaught else 0
    if len(argv) == 4 and argv[0] == "project" and argv[2] == "--wai":
        fails = check_project(argv[1], argv[3])
        for f in fails:
            print(f"FAIL {f}")
        if not fails:
            print(f"project_verify: {argv[1]}: wai.lock holds, each digest recomputed")
        return 1 if fails else 0
    print(__doc__, file=sys.stderr)
    return 2


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