#!/usr/bin/env python3
"""verify_receipts.py -- offline (L1) verification of the SZL-Khipu signed receipts.

Self-contained: given owner_pubkey.json + the two *.signed.json receipts, it
re-derives the keyId, verifies each Ed25519 signature over the receipt's
canonical JSON, confirms the embedded canonical string matches an independently
recomputed canonicalization of the payload, and recomputes the
evaluation->training hash chain. No network, no trust in the repo card.

keyId semantics (identical to szl-forge khipu/sign_receipt.py):
    keyId = first 16 hex chars of sha256(SPKI DER of the public key)

Canonicalization (must match the TypeScript canonicalJson() in the server and
the Python canonical_json() in szl-forge): recursively sorted object keys,
compact separators, floats REFUSED (doctrine: strings + ints only).

Usage:
    python3 verify_receipts.py \
        --pubkey owner_pubkey.json \
        --training training_receipt.signed.json \
        --eval eval_receipt.signed.json

Exit code 0 and a final "PASS" line iff every check passes; non-zero + "FAIL"
otherwise. Requires the `cryptography` package.
"""
import argparse
import base64
import hashlib
import json
import sys

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from cryptography.exceptions import InvalidSignature


def canonical_json(value):
    """Byte-for-byte canonical JSON: sorted keys, compact separators, no floats."""
    if isinstance(value, bool) or value is None or isinstance(value, (str, int)):
        return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
    if isinstance(value, float):
        raise TypeError(
            "canonical_json refuses float values (doctrine: strings+ints only)"
        )
    if isinstance(value, list):
        return "[" + ",".join(canonical_json(v) for v in value) + "]"
    if isinstance(value, dict):
        items = []
        for k in sorted(value.keys()):
            items.append(json.dumps(k, ensure_ascii=False) + ":" + canonical_json(value[k]))
        return "{" + ",".join(items) + "}"
    raise TypeError(f"canonical_json: unsupported type {type(value)!r}")


def sha256_hex(data) -> str:
    if isinstance(data, str):
        data = data.encode("utf-8")
    return hashlib.sha256(data).hexdigest()


def spki_der_from_b64(b64: str) -> bytes:
    der = base64.b64decode(b64)
    # Round-trip through cryptography to confirm it is a real Ed25519 SPKI.
    pub = serialization.load_der_public_key(der)
    if not isinstance(pub, Ed25519PublicKey):
        raise TypeError("public key is not Ed25519")
    return der


def key_id_from_spki(spki_der: bytes) -> str:
    return sha256_hex(spki_der)[:16]


def load(path):
    with open(path, "r", encoding="utf-8") as fh:
        return json.load(fh)


def check(label, ok):
    print(f"  [{'ok' if ok else 'XX'}] {label}")
    return ok


def verify_one(name, signed, pub_spki_der, expected_key_id):
    print(f"[{name}]")
    ok = True

    # 1. Embedded SPKI matches the owner pubkey file.
    embedded_der = base64.b64decode(signed["publicKeySpkiBase64"])
    ok &= check("receipt SPKI == owner_pubkey SPKI", embedded_der == pub_spki_der)

    # 2. keyId re-derived from SPKI matches owner file, receipt wrapper, payload.
    derived = key_id_from_spki(pub_spki_der)
    ok &= check(
        f"keyId re-derived from SPKI = {derived}",
        derived == expected_key_id
        and derived == signed.get("keyId")
        and derived == signed["payload"].get("keyId"),
    )

    # 3. Embedded canonical string == independent canonicalization of payload.
    recomputed = canonical_json(signed["payload"])
    ok &= check(
        "canonical string == canonical_json(payload) recomputed independently",
        recomputed == signed["canonical"],
    )

    # 4. Ed25519 signature verifies over the canonical bytes.
    pub = Ed25519PublicKey.from_public_bytes(
        embedded_der[-32:]  # raw key is last 32 bytes of the 44-byte Ed25519 SPKI
    )
    sig = base64.b64decode(signed["signatureBase64"])
    try:
        pub.verify(sig, signed["canonical"].encode("utf-8"))
        sig_ok = True
    except InvalidSignature:
        sig_ok = False
    ok &= check("Ed25519 signature over canonical JSON", sig_ok)

    return ok, sha256_hex(signed["canonical"])


def main():
    ap = argparse.ArgumentParser(description="Offline verify SZL-Khipu signed receipts.")
    ap.add_argument("--pubkey", default="owner_pubkey.json")
    ap.add_argument("--training", default="training_receipt.signed.json")
    ap.add_argument("--eval", default="eval_receipt.signed.json")
    args = ap.parse_args()

    pub = load(args.pubkey)
    pub_spki_der = spki_der_from_b64(pub["publicKeySpkiBase64"])
    file_key_id = pub["keyId"]

    derived_key_id = key_id_from_spki(pub_spki_der)
    print(f"owner_pubkey keyId (declared) = {file_key_id}")
    print(f"owner_pubkey keyId (derived)  = {derived_key_id}")
    header_ok = derived_key_id == file_key_id
    print(f"  [{'ok' if header_ok else 'XX'}] owner_pubkey keyId self-consistent\n")

    training = load(args.training)
    ev = load(args.eval)

    t_ok, training_canonical_sha = verify_one(
        "training_receipt", training, pub_spki_der, file_key_id
    )
    print()
    e_ok, _ = verify_one("eval_receipt", ev, pub_spki_der, file_key_id)

    # 5. eval->training hash chain: eval.trainingReceiptSha256 == sha256(training canonical).
    print("\n[chain]")
    claimed = ev["payload"].get("trainingReceiptSha256")
    chain_ok = check(
        f"eval.trainingReceiptSha256 == sha256(training canonical) = {training_canonical_sha}",
        claimed == training_canonical_sha,
    )

    all_ok = header_ok and t_ok and e_ok and chain_ok
    print()
    if all_ok:
        print("PASS: both signatures valid, keyId re-derived from SPKI, "
              "canonical strings reproduced, eval->training chain intact.")
        sys.exit(0)
    print("FAIL: one or more checks did not pass (see [XX] lines above).")
    sys.exit(1)


if __name__ == "__main__":
    main()