#!/usr/bin/env python3
"""OER/2 reference verifier (draft 0.1 as profiled by extensions/oer-receipts.md).

Verifies receipts, chains, records, bundles and ledger identities against the
canonical-bytes, number, hashing and composition rules of the extension. Every failure is
an explicit reason; nothing is reinterpreted silently.
"""
import hashlib, io, json, math, re, sys
from decimal import (Context, Decimal, DivisionByZero, Inexact, InvalidOperation, MAX_EMAX,
                     MAX_PREC, MIN_EMIN, Overflow, localcontext)

FIELD_ORDER = {
    "": ["oer", "receipt_id", "kind", "capability", "intent", "subject", "work", "energy", "provenance",
         "settlement", "floor", "efficiency", "bundle", "chain", "hash_alg", "sig_alg"],
    "subject": ["inputs", "output"],
    "work": ["status", "outcome", "latency_s", "proficiency"],
    "energy": ["joules_total", "boundary", "breakdown", "method", "temperature_k"],
    "energy.breakdown": ["compute_j", "movement_j", "idle_alloc_j", "shared_alloc_j", "upstream_j"],
    "provenance": ["how", "device", "weakest_input", "hardware_profile", "corpus"],
    "provenance.weakest_input": ["how", "input"],
    "settlement": ["question_key", "served_from_record", "record_id", "record_hash", "reuse_count", "derivation_j", "valid_until"],
    "floor": ["bits_committed", "bits_delivered", "floor_j", "decades_above_floor"],
    "efficiency": ["completed_units", "wasted_units", "joules_per_completed", "waste_fraction", "value_per_joule"],
    "bundle": ["stages", "provenance_root"],
    "bundle.stages[]": ["kind", "receipt_hash", "joules"],
    "chain": ["prev_hash", "receipt_hash", "signer", "signature", "anchor"],
    "chain.anchor": ["path", "delivered_at", "root"],
}
KINDS = {"decision", "task", "reconstruction", "calibration", "job", "compile", "decode", "mitigate", "rearrange", "anchor", "bundle"}
GRADES = ["projected", "simulated", "derived", "measured"]
KBOLTZ, LN2 = Decimal("1.380649e-23"), Decimal("0.6931471805599453")

def _order(obj, path=""):
    if isinstance(obj, dict):
        keys = FIELD_ORDER.get(path, sorted(obj.keys()))
        out = {}
        for k in keys:
            if k in obj:
                out[k] = _order(obj[k], (path + "." + k) if path else k)
        for k in obj:
            if k not in out:
                raise ValueError("unknown field %s at %s" % (k, path or "root"))
        return out
    if isinstance(obj, list):
        return [_order(x, path + "[]") for x in obj]
    return obj

def _dumps(v):
    """The one serialisation: compact separators, non-ASCII unescaped, UTF-8."""
    return json.dumps(v, separators=(",", ":"), ensure_ascii=False).encode("utf-8")

def canonical_bytes(receipt):
    body = {k: v for k, v in receipt.items() if k not in ("chain", "receipt_id")}
    return _dumps(_order(body))

def stage_list_bytes(stages):
    """The bytes `bundle.provenance_root` addresses: the stage list in member order,
    serialised exactly as the body is."""
    return _dumps([_order(s, "bundle.stages[]") for s in stages])

# The number profile (extension §4). A receipt body carries integers within I-JSON's
# +-(2**53 - 1) and decimal strings, nothing else: no fraction or exponent, no negative
# zero. Parsers disagree about the rest -- Python reads -0 and 2**64 as exact integers,
# serde_json reads both as floats -- so two verifiers would hash one receipt two ways.
MAX_SAFE = 2**53 - 1

def _int(s):
    if s == "-0":
        return -0.0  # a float, as serde_json reads it, so the number rule sees it
    v = int(s)
    float(v)  # OverflowError past binary64, where serde_json refuses the document
    return v

def _float(s):
    v = float(s)
    if not math.isfinite(v):
        raise ValueError("number out of range: " + s)
    return v

def _constant(s):
    raise ValueError("not JSON: " + s)

MAX_DEPTH = 64

def _walk(v):
    """Every value in a document with its depth, iteratively: a document nested past
    the interpreter's recursion limit is refused, not a crash."""
    stack = [(v, 1)]
    while stack:
        x, d = stack.pop()
        yield x, d
        if isinstance(x, dict):
            stack.extend((y, d + 1) for y in x.values())
            stack.extend((k, d + 1) for k in x)
        elif isinstance(x, list):
            stack.extend((y, d + 1) for y in x)

def _depth(v):
    return max((d for x, d in _walk(v) if isinstance(x, (dict, list))), default=0)

def _surrogate(v):
    return any(isinstance(x, str) and any("\ud800" <= c <= "\udfff" for c in x) for x, _ in _walk(v))

def _no_duplicates(pairs):
    seen = {}
    for k, v in pairs:
        if k in seen:
            raise ValueError("duplicate member %r" % k)
        seen[k] = v
    return seen

def load(fp):
    """Read a vector document under the document profile (§4), or refuse it whole
    (ValueError): a number whose nearest binary64 is infinite, NaN or Infinity (not
    JSON), an unpaired surrogate, a member name repeated in one object (RFC 7493
    §2.3), or nesting past 64. `-0` is read as a float."""
    try:
        v = json.load(fp, parse_int=_int, parse_float=_float, parse_constant=_constant,
                      object_pairs_hook=_no_duplicates)
    except OverflowError as e:
        raise ValueError("number out of range") from e
    except RecursionError as e:
        raise ValueError("nested deeper than %d" % MAX_DEPTH) from e
    if _depth(v) > MAX_DEPTH:
        raise ValueError("nested deeper than %d" % MAX_DEPTH)
    if _surrogate(v):
        raise ValueError("unpaired surrogate")
    return v

# A decimal string (§4): JSON's number grammar in ASCII -- no sign but `-`, no
# whitespace, no bare point, no leading zero -- with at most 34 significand digits and
# a one-to-three-digit exponent, so every one is exact in IEEE 754 decimal128.
DECIMAL = re.compile(r"-?(0|[1-9][0-9]*)(\.[0-9]+)?([eE][+-]?[0-9]{1,3})?", re.ASCII)

def dec(v):
    """A quantity a check reads: a decimal string, or an integer in I-JSON's range."""
    if isinstance(v, str) and DECIMAL.fullmatch(v) and len(re.sub(r"[eE].*|[-.]", "", v)) <= 34:
        return Decimal(v)
    if isinstance(v, int) and not isinstance(v, bool) and -MAX_SAFE <= v <= MAX_SAFE:
        return Decimal(v)
    raise TypeError("decimal")

# Every composition check is exact (§6): no operation here may round. A check that
# ever did would trap Inexact rather than compare rounding noise.
EXACT = Context(prec=MAX_PREC, Emax=MAX_EMAX, Emin=MIN_EMIN,
                traps=[InvalidOperation, DivisionByZero, Overflow, Inexact])

def _forbidden_number(v):
    if isinstance(v, dict):
        return any(_forbidden_number(x) for x in v.values())
    if isinstance(v, list):
        return any(_forbidden_number(x) for x in v)
    if isinstance(v, bool):
        return False
    if isinstance(v, float):
        return True
    return isinstance(v, int) and not -MAX_SAFE <= v <= MAX_SAFE

def h(alg, data):
    if alg == "sha-256":
        return "sha256:" + hashlib.sha256(data).hexdigest()
    if alg == "blake3":
        try:
            import blake3
            return "blake3:" + blake3.blake3(data).hexdigest()
        except ImportError:
            raise KeyError("missing_capability: blake3")
    raise KeyError("missing_capability: " + alg)

# Ed25519 key decoding, strictly (RFC 8032 §5.1.3), with the RFC's own reference
# arithmetic (§6). OpenSSL accepts a key that is a non-canonical encoding, or a point
# of small order -- which verifies a signature over any message -- so the reference
# checks both before it asks OpenSSL anything.
_P = 2**255 - 19
_D = -121665 * pow(121666, _P - 2, _P) % _P
_SQRT_M1 = pow(2, (_P - 1) // 4, _P)

def _point_add(p, q):
    a, b = (p[1] - p[0]) * (q[1] - q[0]) % _P, (p[1] + p[0]) * (q[1] + q[0]) % _P
    c, d = 2 * p[3] * q[3] * _D % _P, 2 * p[2] * q[2] % _P
    e, f, g, h = b - a, d - c, d + c, b + a
    return (e * f % _P, g * h % _P, f * g % _P, e * h % _P)

def _point_mul(s, p):
    q = (0, 1, 1, 0)
    while s > 0:
        if s & 1:
            q = _point_add(q, p)
        p = _point_add(p, p)
        s >>= 1
    return q

def _decompress(b):
    """The point `b` encodes, or None when it is not a canonical encoding of one."""
    y = int.from_bytes(b, "little")
    sign, y = y >> 255, y & ((1 << 255) - 1)
    if y >= _P:
        return None
    x2 = (y * y - 1) * pow(_D * y * y + 1, _P - 2, _P) % _P
    if x2 == 0:
        if sign:
            return None
        x = 0
    else:
        x = pow(x2, (_P + 3) // 8, _P)
        if (x * x - x2) % _P:
            x = x * _SQRT_M1 % _P
        if (x * x - x2) % _P:
            return None
        if (x & 1) != sign:
            x = _P - x
    return (x, y, 1, x * y % _P)

def strict_key(b):
    """A published Ed25519 key a signature can rest on: 32 bytes, the canonical
    encoding of a point, and a point not of small order."""
    a = _decompress(b) if len(b) == 32 else None
    if a is None:
        return False
    x, y, z, _ = _point_mul(8, a)
    return not (x % _P == 0 and (y - z) % _P == 0)  # [8]A is not the identity

def verify_signature(receipt, keys):
    if receipt.get("sig_alg") != "ed25519":
        return "missing_capability: " + str(receipt.get("sig_alg"))
    signer = receipt["chain"]["signer"]
    pub = keys.get(signer) if isinstance(signer, str) else None
    if not isinstance(pub, str):
        return "signer_unknown"
    try:
        from cryptography.hazmat.primitives.asymmetric import ed25519
        if not strict_key(bytes.fromhex(pub)):
            return "signature_invalid"
        ed25519.Ed25519PublicKey.from_public_bytes(bytes.fromhex(pub)).verify(bytes.fromhex(receipt["chain"]["signature"]), receipt["chain"]["receipt_hash"].encode())
        return None
    except Exception:
        return "signature_invalid"

def verify_receipt(r, keys):
    """Returns a list of failure reasons (empty means valid).

    A missing or mistyped member is reported as `malformed: <member>` rather than
    raised: every failure is an explicit reason, including a receipt too broken to
    check. wai-rs's `oer` verifier and this one agree on that.
    """
    try:
        with localcontext(EXACT):
            return _verify_receipt(r, keys)
    except (KeyError, TypeError, AttributeError, InvalidOperation) as e:
        return ["malformed: " + str(e.args[0] if e.args else type(e).__name__)]

def _verify_receipt(r, keys):
    reasons = []
    if not isinstance(r, dict):
        return ["malformed: receipt"]
    # The major comes from the string only: a parsed float would make it depend on how
    # each JSON parser rounds (2.99999999999999977 is 3.0 to one and not the other).
    o = r.get("oer")
    if not isinstance(o, str) or o.split(".")[0] != "2":
        return ["unknown_major"]
    if not (isinstance(r.get("kind"), str) and r["kind"] in KINDS):
        reasons.append("unknown_kind")
    try:
        body = canonical_bytes(r)
        if "chain" in r:
            _order(r["chain"], "chain")  # outside the body, but its members are fixed too (§3)
    except ValueError as e:
        return ["canonical_failure: " + str(e)]
    if _forbidden_number({k: v for k, v in r.items() if k not in ("chain", "receipt_id")}):
        return ["malformed: number"]
    try:
        if not isinstance(r["hash_alg"], str):
            raise TypeError("hash_alg")
        if h(r["hash_alg"], body) != r["receipt_id"]:
            reasons.append("receipt_id_mismatch")
        expected = h(r["hash_alg"], r["chain"]["prev_hash"].encode() + body)
        if expected != r["chain"]["receipt_hash"]:
            reasons.append("receipt_hash_mismatch")
    except KeyError as e:
        # e.args[0], not str(e): str() of a KeyError is repr-quoted, which turned
        # 'missing_capability' into "'missing_capability" -- a reason no vector could name.
        # An absent member is not a capability gap; let it surface as `malformed`.
        if str(e.args[0]).startswith("missing_capability"):
            return [e.args[0]]
        raise
    e = r["energy"]
    total = sum(dec(v) for v in e["breakdown"].values())
    if abs(total - dec(e["joules_total"])) > Decimal("1e-9") * max(Decimal(1), abs(total)):
        reasons.append("sum_mismatch")
    b = e.get("boundary")
    if not (isinstance(b, str) and b):
        reasons.append("boundary_missing")
    f = r["floor"]
    floor = dec(f["bits_committed"]) * KBOLTZ * dec(e["temperature_k"]) * LN2
    if abs(floor - dec(f["floor_j"])) > abs(floor) * Decimal("0.01"):
        reasons.append("floor_mismatch")
    if r["provenance"]["how"] not in GRADES:
        reasons.append("grade_unknown")
    # An abstention makes no verdict claim: no answer, and no record anything could
    # later be served from (§3). Read leniently, so a strange shape is simply not an
    # abstention rather than a different reason.
    w = r.get("work")
    if isinstance(w, dict) and w.get("outcome") == "abstain":
        sub = r.get("subject") if isinstance(r.get("subject"), dict) else {}
        st = r.get("settlement") if isinstance(r.get("settlement"), dict) else {}
        settled = any(st.get(k) for k in ("served_from_record", "record_id", "record_hash"))
        if sub.get("output") or settled:
            reasons.append("abstention_settled")
    if r.get("kind") == "bundle":
        stages = r["bundle"]["stages"]
        if not isinstance(stages, list):
            raise TypeError("stages")
        root = h(r["hash_alg"], stage_list_bytes(stages))
        if root != r["bundle"]["provenance_root"]:
            reasons.append("bundle_root_mismatch")
        st = sum(dec(s["joules"]) for s in stages)
        if abs(st - dec(e["joules_total"])) > Decimal("1e-9") * max(Decimal(1), abs(st)):
            reasons.append("bundle_sum_mismatch")
        eff = r.get("efficiency")
        if not isinstance(eff, dict) or eff.get("completed_units") is None:
            reasons.append("efficiency_missing")
    sig = verify_signature(r, keys)
    if sig:
        reasons.append(sig)
    return reasons

def _link(r, k):
    ch = r.get("chain") if isinstance(r, dict) else None
    v = ch.get(k) if isinstance(ch, dict) else None
    return v if isinstance(v, str) else None

def verify_chain(receipts, keys):
    reasons = {}
    for i, r in enumerate(receipts):
        rr = verify_receipt(r, keys)
        # Every link after the first: a string equal to the previous receipt's hash,
        # which must itself be a string -- a hash that is not one links nothing.
        if i > 0:
            prev = _link(receipts[i - 1], "receipt_hash")
            if prev is None or _link(r, "prev_hash") != prev:
                rr.append("chain_break")
        if rr:
            # Accumulate: a receipt id that recurs must not replace, and so hide, the
            # reasons an earlier receipt with the same id already failed for.
            rid = r.get("receipt_id", "?") if isinstance(r, dict) else "?"
            rid = rid if isinstance(rid, str) else json.dumps(rid, separators=(",", ":"), ensure_ascii=False)
            reasons.setdefault(rid, []).extend(rr)
    return reasons

def record_hash(rec):
    """A record whose key or answer is not a string, or whose created is not a
    number, cannot be rehashed (§6)."""
    key, answer, created = rec["key"], rec["answer"], rec["created"]
    if not (isinstance(key, str) and isinstance(answer, str)) or isinstance(created, bool) \
            or not isinstance(created, (int, float)):
        raise TypeError("record")
    return "sha256:" + hashlib.sha256(("%s|%s|%.1f" % (key, answer, created)).encode()).hexdigest()

def verify_records(receipts, records):
    bad = []
    records = records if isinstance(records, dict) else {}
    for r in receipts:
        s = r.get("settlement") if isinstance(r, dict) else None
        if not isinstance(s, dict) or not s.get("served_from_record") or s.get("record_id") == "unverified":
            continue
        rid = s.get("record_id")
        rec = records.get(rid) if isinstance(rid, str) else None  # a non-string id names no record
        try:
            ok = isinstance(rec, dict) and record_hash(rec) == s.get("record_hash")
        except (KeyError, TypeError, ValueError):
            ok = False  # a record too broken to rehash cannot vouch for anything
        if not ok:
            bad.append(str(r.get("receipt_id")))
    return bad

def verify_ledger(receipts, ledger):
    try:
        with localcontext(EXACT):
            return _verify_ledger(receipts, ledger)
    except (KeyError, TypeError, AttributeError, InvalidOperation, IndexError):
        return False  # an identity that cannot be summed does not hold

def _verify_ledger(receipts, ledger):
    total = sum(dec(r["energy"]["joules_total"]) for r in receipts) + dec(ledger["unallocated_joules"]) + dec(ledger.get("carried_forward_joules", "0"))
    meter = dec(ledger["meter_joules"])
    return abs(total - meter) <= Decimal("1e-9") * max(Decimal(1), abs(meter))

# The canonical form, pinned independently of the code that produces it.
#
# vectors.json cannot check this on its own: generate_vectors.py imports this module
# for canonical_bytes/_order/h, so a defect in the field order is baked into every
# receipt id and then blessed by this verifier. Reversing the declared order and
# regenerating still yields nine passes, with different receipt ids.
#
# The fixture below closes that. Its input deliberately supplies keys in the WRONG
# order at two levels, so the ordering is proved rather than incidental, and the
# expected output is a literal — the one thing in this file that does not come from
# _order. If canonicalization changes, this fails and the vectors do not have to.
CANONICAL_FIXTURE = {
    "sig_alg": "ed25519", "kind": "decision", "oer": "2.0",
    "capability": "oer.canonical.fixture", "intent": "replicate",
    "energy": {
        "method": "fixture",
        "breakdown": {"upstream_j": 5, "compute_j": 1, "idle_alloc_j": 3,
                      "movement_j": 2, "shared_alloc_j": 4},
        "joules_total": 15, "boundary": "node",
    },
    "hash_alg": "sha-256",
}
CANONICAL_FIXTURE_BYTES = "{\"oer\":\"2.0\",\"kind\":\"decision\",\"capability\":\"oer.canonical.fixture\",\"intent\":\"replicate\",\"energy\":{\"joules_total\":15,\"boundary\":\"node\",\"breakdown\":{\"compute_j\":1,\"movement_j\":2,\"idle_alloc_j\":3,\"shared_alloc_j\":4,\"upstream_j\":5},\"method\":\"fixture\"},\"hash_alg\":\"sha-256\",\"sig_alg\":\"ed25519\"}"
CANONICAL_FIXTURE_ID = "sha256:364718c762d6de8afd6ad59dcc5023eafbf0023b6bc4acb9f4885de17084ed30"

# The stage list `bundle.provenance_root` addresses, pinned the same way. Every
# stage list in the corpus is ASCII, and on ASCII the escaped and unescaped
# serialisations agree, so no vector can tell them apart; this input is non-ASCII
# (a two-byte and a four-byte character) with its members out of order, against a
# literal.
STAGE_FIXTURE = [{"joules": "2.5", "receipt_hash": "sha256:00", "kind": "décision 😀"}]
STAGE_FIXTURE_BYTES = '[{"kind":"décision 😀","receipt_hash":"sha256:00","joules":"2.5"}]'
STAGE_FIXTURE_ID = "sha256:c1caa0cea7cf581f240346d9e9431d19cbe204c4daca157fea555fee6bf2638b"


def check_canonical_form():
    """Return a list of reasons the canonical form or the stage list no longer
    matches its fixture."""
    reasons = []
    got = canonical_bytes(CANONICAL_FIXTURE).decode("utf-8")
    if got != CANONICAL_FIXTURE_BYTES:
        reasons.append("canonical_form_mismatch")
    if h("sha-256", got.encode("utf-8")) != CANONICAL_FIXTURE_ID:
        reasons.append("canonical_id_mismatch")
    got = stage_list_bytes(STAGE_FIXTURE).decode("utf-8")
    if got != STAGE_FIXTURE_BYTES:
        reasons.append("stage_form_mismatch")
    if h("sha-256", got.encode("utf-8")) != STAGE_FIXTURE_ID:
        reasons.append("stage_id_mismatch")
    # The number rule reads -0 through load(); if that ever stops yielding a float,
    # -0 would read as the integer zero and verify. Fail here, loudly, instead.
    if not _forbidden_number(load(io.StringIO("-0"))):
        reasons.append("number_profile_unavailable")
    return reasons


def run(vectors):
    keys = vectors.get("keys") if isinstance(vectors, dict) else None
    if not isinstance(keys, dict):
        raise ValueError("vector file has no keys")
    if not isinstance(vectors.get("vectors"), list):
        raise ValueError("vector file has no vectors")
    results = []
    for v in vectors["vectors"]:
        v = v if isinstance(v, dict) else {}
        receipts = v.get("receipts") if isinstance(v.get("receipts"), list) else []
        got = set()
        chain = verify_chain(receipts, keys)
        for rs in chain.values():
            for x in rs:
                got.add(x.split(":")[0])
        if "records" in v and verify_records(receipts, v["records"]):
            got.add("record_mismatch")
        if "ledger" in v and not verify_ledger(receipts, v["ledger"]):
            got.add("ledger_mismatch")
        exp = v.get("expected")
        expected = {x for x in exp if isinstance(x, str)} if isinstance(exp, list) else set()
        results.append({"vector": str(v.get("name")), "expected": sorted(expected), "got": sorted(got),
                        "pass": got == expected})
    return results

def refuse_all(directory):
    """Every document in `directory` must be refused whole (§4)."""
    import os
    names = sorted(n for n in os.listdir(directory)
                   if n.endswith(".json") and os.path.isfile(os.path.join(directory, n)))
    ok = bool(names)
    for name in names:
        try:
            with open(os.path.join(directory, name), encoding="utf-8") as fp:
                load(fp)
            print("FAIL %s was read; it must be refused whole" % name)
            ok = False
        except ValueError as e:
            print("PASS %s refused  (%s)" % (name, e))
        except OSError as e:
            print("FAIL %s: %s" % (name, e))
            ok = False
    return ok

if __name__ == "__main__":
    if sys.argv[1:2] == ["--unreadable"]:
        sys.exit(0 if refuse_all(sys.argv[2] if len(sys.argv) > 2 else "unreadable") else 1)
    path = sys.argv[1] if len(sys.argv) > 1 else "vectors.json"
    try:
        vec = load(open(path, encoding="utf-8"))
        run(vec)  # a file without keys or vectors is refused here, not mid-report
    except ValueError as e:
        print("FAIL %s: %s" % (path, e))
        sys.exit(1)
    ok = True
    # The canonical form first: if it has moved, every id below is computed under a
    # different rule and the passes that follow would mean nothing.
    cf = check_canonical_form()
    print(("PASS " if not cf else "FAIL ") + "canonical form and stage list pinned to their fixtures"
          + "  expected=  got=" + ",".join(cf))
    ok = ok and not cf
    for res in run(vec):
        ok = ok and res["pass"]
        print(("PASS " if res["pass"] else "FAIL ") + res["vector"] + "  expected=" + ",".join(res["expected"]) + "  got=" + ",".join(res["got"]))
    sys.exit(0 if ok else 1)