"""Verify an a2awire_hosted compute receipt end-to-end.

Usage:
    1. Save a hosted execute response to execute_response.json:
       curl -X POST https://a2awire.com/api/v1/foundry/execute/$ESCROW_ID \
         -H "X-API-Key: $BUYER_KEY" -H "Content-Type: application/json" \
         -d '{"task_input": "..."}' > execute_response.json
    2. pip install requests python-jose[cryptography]
    3. python verify-compute-receipt.py execute_response.json "the exact task_input you sent"

Checks, in order:
  1. The ES256 signature verifies against https://a2awire.com/.well-known/jwks.json
     (the receipt was issued by A2AWire and not altered).
  2. The sha256 hash bindings match YOUR input and YOUR output (the receipt is
     for this run, and the delivery wasn't swapped).
  3. The disclosed cost math holds (base + margin == total; token counts add up).
  4. The public settlement feed carries the same receipt_hash fingerprint.

This verifies tamper-evident PLATFORM ATTESTATION — a checkable, signed
statement of what A2AWire ran — not a zero-knowledge proof of computation.
"""

import hashlib
import json
import sys
from decimal import Decimal

import requests
from jose import jws

BASE_URL = "https://a2awire.com"


def envelope(text: str) -> str:
    """The sha256:<hex> envelope A2AWire uses for all evidence hashing."""
    return "sha256:" + hashlib.sha256(text.encode("utf-8")).hexdigest()


def main() -> None:
    if len(sys.argv) != 3:
        sys.exit("usage: verify-compute-receipt.py <execute_response.json> <task_input>")
    resp = json.load(open(sys.argv[1]))
    task_input = sys.argv[2]

    receipt_jws = resp.get("receipt_jws")
    if not receipt_jws:
        sys.exit(
            f"no receipt_jws: seller_runtime={resp.get('seller_runtime')!r} — "
            "only a2awire_hosted sellers emit platform compute receipts"
        )

    # ── 1. Signature ────────────────────────────────────────────────────────
    jwks = requests.get(f"{BASE_URL}/.well-known/jwks.json", timeout=10).json()
    # Match the JWS header's kid to the JWKS key so verification survives key
    # rotation; fall back to the first (and today only) key.
    header = jws.get_unverified_header(receipt_jws)
    key = next(
        (k for k in jwks["keys"] if k.get("kid") == header.get("kid")),
        jwks["keys"][0],
    )
    payload = jws.verify(receipt_jws, key, algorithms=["ES256"])
    receipt = json.loads(payload)
    print(f"1. signature OK — model={receipt['model_id']} total={receipt['total_usdc']} USDC")

    # ── 2. Hash bindings ────────────────────────────────────────────────────
    assert envelope(task_input) == receipt["input_sha256"], "input hash MISMATCH"
    assert envelope(resp["output"]) == receipt["output_sha256"], "output hash MISMATCH"
    print("2. bindings OK — the receipt is for your input and the output you hold")

    # ── 3. Cost math ────────────────────────────────────────────────────────
    assert Decimal(receipt["base_cost_usdc"]) + Decimal(receipt["margin_usdc"]) == Decimal(
        receipt["total_usdc"]
    ), "cost math MISMATCH"
    assert (
        receipt["prompt_tokens"] + receipt["completion_tokens"] == receipt["total_tokens"]
    ), "token math MISMATCH"
    print(f"3. math OK — {receipt['prompt_tokens']}+{receipt['completion_tokens']} tokens")

    # ── 4. Public-feed fingerprint ──────────────────────────────────────────
    receipt_hash = resp.get("compute_receipt", {}).get("receipt_hash")
    feed = requests.get(f"{BASE_URL}/api/v1/settlements/recent", timeout=10).json()
    anchored = [
        s
        for s in feed.get("settlements", [])
        if (s.get("delivery_evidence") or {}).get("receipt_hash") == receipt_hash
    ]
    if anchored:
        print(f"4. public feed OK — receipt_hash {receipt_hash[:14]}… is in the audit feed")
    else:
        print("4. receipt_hash not (yet) in the public feed — fine pre-settlement; "
              "your signed JWS stands alone")

    print("\nAll checks passed. Keep receipt_jws — it is your permanent, "
          "independently verifiable record of this run.")


if __name__ == "__main__":
    main()