#!/usr/bin/env python3
"""
JCR-1 reference verifier.

Independent of any JouleClaw code. Reads jcr1-vectors.json + jcr1-test-keys.json
and verifies each vector against the JCR-1 spec:
  1 structural / re-encode determinism (RFC 8949 4.2.1 canonical CBOR)
  2 COSE_Sign1 envelope signature (RFC 9052, EdDSA/Ed25519)
  3 gateway tool-receipt HMAC-SHA-256
  4 honesty consistency (worst-counter rule, joule sum, mJ resolution floor)

A second implementation in any language is "vector-conformant" if it reproduces
these PASS results. Deps: cbor2, cryptography (both pip-installable).

Usage: python3 verify.py [vectors.json] [test-keys.json]
"""
import sys, json, hmac, hashlib
import cbor2
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey

PROV_ORDER = {"HwShunt": 3, "ModelBased": 2, "Estimator": 1}
TOKEN_FIELD = "token_attributions"

def signed_core_bytes(receipt: dict) -> bytes:
    core = {k: v for k, v in receipt.items()
            if k != TOKEN_FIELD and v is not None}
    return cbor2.dumps(core, canonical=True)

def hmac_payload(tr: dict) -> bytes:
    return b"\x00".join([
        tr["tool_id"].encode(), str(tr["joules_uj"]).encode(),
        tr["energy_provenance"].encode(), tr["minted_at"].encode(),
        tr["gateway_id"].encode(),
    ])

def worst(provs):
    return min(provs, key=lambda p: PROV_ORDER[p])

def verify_vector(v, pubkey, gw_keys):
    r = v["receipt"]
    errs = []

    # 1 — re-encode determinism
    core = bytes.fromhex(v["signed_core_cbor_hex"])
    if signed_core_bytes(r) != core:
        errs.append("signed core not reproducible")

    # 2 — COSE_Sign1 envelope
    try:
        tag = cbor2.loads(bytes.fromhex(v["cose_sign1_hex"]))
        protected_b, _unprotected, payload, signature = tag.value
        if payload != core:
            errs.append("envelope payload != signed core")
        ph = cbor2.loads(protected_b)
        if ph.get(1) != -8:
            errs.append(f"unexpected alg {ph.get(1)} (want -8 EdDSA)")
        sig_struct = cbor2.dumps(["Signature1", protected_b, b"", payload], canonical=True)
        pubkey.verify(signature, sig_struct)
    except Exception as e:
        errs.append(f"envelope verify failed: {e}")

    # 3 — gateway tool-receipt HMACs
    for tr in r.get("tool_receipts", []):
        key = gw_keys.get(tr["gateway_id"])
        if key is None:
            errs.append(f"no key for gateway {tr['gateway_id']}")
            continue
        expect = hmac.new(key, hmac_payload(tr), hashlib.sha256).hexdigest()
        if not hmac.compare_digest(expect, tr["hmac_sha256_hex"]):
            errs.append(f"HMAC mismatch {tr['tool_id']}")

    # 4 — honesty consistency
    provs = [t["energy_provenance"] for t in r.get("tools_touched", [])]
    provs += [t["energy_provenance"] for t in r.get("tool_receipts", [])]
    if provs and r["energy_provenance"] != worst(provs):
        errs.append("energy_provenance != worst counter")
    tsum = sum(t["joules_uj"] for t in r.get("tools_touched", []))
    if r["joules_uj"] < tsum:
        errs.append("joules_uj < sum(tools_touched)")
    if r["energy_provenance"] in ("ModelBased", "Estimator") and r["joules_uj"] % 1000 != 0:
        errs.append("mJ resolution floor violated")

    return errs

def main():
    vec_path = sys.argv[1] if len(sys.argv) > 1 else "/mnt/user-data/outputs/jcr1-vectors.json"
    key_path = sys.argv[2] if len(sys.argv) > 2 else "/mnt/user-data/outputs/jcr1-test-keys.json"
    data = json.load(open(vec_path))
    keys = json.load(open(key_path))

    pubkey = Ed25519PublicKey.from_public_bytes(bytes.fromhex(keys["ed25519_public_raw_hex"]))
    gw_keys = {k: val.encode() for k, val in keys["gateway_hmac_keys"].items()}

    print(f"JCR-1 verifier — {data['vector_count']} vectors — {data['envelope']}\n")
    npass = 0
    for v in data["vectors"]:
        errs = verify_vector(v, pubkey, gw_keys)
        ok = not errs
        npass += ok
        print(f"  [{('PASS' if ok else 'FAIL')}] vector {v['n']:>2}: {v['description']}")
        for e in errs:
            print(f"         - {e}")
    print(f"\n{npass}/{data['vector_count']} vectors verified")
    sys.exit(0 if npass == data["vector_count"] else 1)

if __name__ == "__main__":
    main()