"""Verify saved inference attestation evidence, following the provider's public guide
(github.com/chutesai/chutes-api/blob/main/docs/tee-verification.md), and print a report.

Per instance:
  1. Intel TDX quote: signature chain to the Intel root, TCB status, revocation
     (voltage-verify quote, Intel PCS collateral fetched now).
  2. Nonce binding: first 32 bytes of report_data == SHA256(nonce + e2e_pubkey).
  3. Debug bit off (td_attributes, little-endian).
  4. Key possession: RSA signature of attested_body by the instance certificate,
     and attested_body carries our nonce and this very quote.
  5. NVIDIA GPU evidence: nv-attestation-sdk, remote verifier (NVIDIA NRAS),
     nonce = the same SHA256(nonce + e2e_pubkey).

Usage: python verify_inference_evidence.py <dir> <slug> <voltage-verify.exe>
"""
import base64
import hashlib
import json
import subprocess
import sys
import tempfile
from datetime import datetime, timezone
from pathlib import Path

from cryptography import x509
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding


def tdx_check(verifier: str, quote_b64: str) -> dict:
    with tempfile.NamedTemporaryFile("w", suffix=".b64", delete=False) as f:
        f.write(quote_b64)
    out = subprocess.run([verifier, "quote", f.name, "--json"], capture_output=True, text=True)
    Path(f.name).unlink(missing_ok=True)
    return json.loads(out.stdout)


def key_possession(inst: dict, nonce: str) -> str:
    if not inst.get("signature") or not inst.get("attested_body"):
        return "SKIPPED (no signature)"
    cert = x509.load_der_x509_certificate(base64.b64decode(inst["certificate"]))
    body = base64.b64decode(inst["attested_body"])
    try:
        cert.public_key().verify(base64.b64decode(inst["signature"]), body, padding.PKCS1v15(), hashes.SHA256())
    except Exception:
        return "FAIL (signature)"
    parsed = json.loads(body)
    if parsed.get("nonce") != nonce:
        return "FAIL (signed body has another nonce)"
    signed = parsed.get("evidence", {})
    signed_quote = signed.get("tdx_quote") or signed.get("quote")  # field name differs across proxy versions
    if not signed_quote or base64.b64decode(signed_quote) != base64.b64decode(inst["quote"]):
        return "FAIL (signed body has another quote)"
    return "PASS"


def gpu_check(gpu_evidence: list, expected: str) -> str:
    from nv_attestation_sdk.attestation import Attestation, Devices, Environment

    client = Attestation()
    client.add_verifier(Devices.GPU, Environment.REMOTE, "", "")
    client.set_nonce(expected)
    try:
        return "PASS" if client.attest(gpu_evidence) else "FAIL"
    except Exception as e:  # report, never hide
        return f"ERROR ({type(e).__name__})"


def main() -> int:
    folder, slug, verifier = Path(sys.argv[1]), sys.argv[2], sys.argv[3]
    nonce = (folder / f"{slug}.nonce.txt").read_text().strip()
    evidence = json.loads((folder / f"{slug}.evidence.json").read_text())
    e2e = json.loads((folder / f"{slug}.e2e-instances.json").read_text())
    pubkeys = {i["instance_id"]: i["e2e_pubkey"] for i in e2e.get("instances", [])}
    version = subprocess.run([verifier, "--version"], capture_output=True, text=True).stdout.strip()

    lines = [
        f"verified_at: {datetime.now(timezone.utc).isoformat(timespec='seconds')}",
        f"nonce (ours, random 32 bytes): {nonce}",
        f"verifiers: {version}; nv-attestation-sdk (NVIDIA NRAS, remote)",
        f"instances returned: {len(evidence.get('evidence', []))}",
        "",
    ]
    all_ok = True
    full = quote_only = 0
    for inst in evidence.get("evidence", []):
        iid = inst["instance_id"]
        tdx = tdx_check(verifier, inst["quote"])
        rd = tdx["facts"]["tdx"]["report_data"]
        attrs = tdx["facts"]["tdx"].get("td_attributes", "")
        debug_off = bool(attrs) and (int(attrs[:2], 16) & 1) == 0
        pk = pubkeys.get(iid)
        expected = hashlib.sha256((nonce + pk).encode()).hexdigest() if pk else None
        bound = expected is not None and rd[:64] == expected
        poss = key_possession(inst, nonce)
        gpu = gpu_check(inst["gpu_evidence"], expected) if expected else "SKIPPED (no e2e key)"
        ok = tdx["ok"] and bound and debug_off and poss == "PASS" and gpu == "PASS"
        all_ok = all_ok and ok
        full += ok
        quote_only += (not ok) and pk is None and tdx["ok"] and debug_off and poss == "PASS"
        failed = [c["name"] for c in tdx["checks"] if c["result"] != "PASS"]
        lines += [
            f"instance {iid}",
            f"  intel tdx quote  : {'PASS' if tdx['ok'] else 'FAIL ' + ','.join(failed)} ({tdx['checks'][3]['detail'] if len(tdx['checks']) > 3 else ''})",
            f"  nonce binding    : {'PASS' if bound else 'FAIL' if pk else 'NOT CHECKABLE (instance not in the E2E key list)'} report_data[0:32] = {rd[:64]}",
            f"  debug bit        : {'OFF' if debug_off else 'ON or unknown'} (td_attributes {attrs})",
            f"  key possession   : {poss}",
            f"  nvidia gpus      : {gpu} ({len(inst['gpu_evidence'])} GPU(s), arch {inst['gpu_evidence'][0].get('arch', '?') if inst['gpu_evidence'] else '?'})",
            f"  mr_td            : {tdx['facts']['tdx']['mr_td']}",
            "",
        ]
    total = len(evidence.get("evidence", []))
    lines.append(f"RESULT: {full} of {total} instance(s) fully verified (Intel quote, our nonce, key possession, NVIDIA GPUs)")
    if quote_only:
        lines.append(
            f"        {quote_only} instance(s) not in the provider's public E2E key list: genuine up-to-date TDX quote and"
            " signed body carrying our nonce, but the report_data binding cannot be recomputed without that key"
        )
    if total - full - quote_only:
        lines.append(f"        {total - full - quote_only} instance(s) FAILED a check, see above")
    report = "\n".join(lines) + "\n"
    print(report)
    (folder / f"{slug}.verify.txt").write_text(report)
    return 0 if all_ok else 2


if __name__ == "__main__":
    sys.exit(main())