prover.py

#!/usr/bin/env python3
import os
import argparse
import json
import re
import datetime
import sys
import time
import zlib
from pathlib import Path
from typing import Optional
try:
    import resource  # POSIX-only
except ModuleNotFoundError:
    resource = None

from util.zk_pipeline import (
    generate_witness,
    groth16_prove,
    read_allowed_from_public,
)

from xrpl.clients import JsonRpcClient
from xrpl.wallet import Wallet
from xrpl.models.transactions import AccountSet, Memo
from xrpl.transaction import submit_and_wait  # using same helper

DEFAULT_RPC = "https://s.altnet.rippletest.net:51234"
DEFAULT_SEED = "sEdTTH9np17Br7zJkkE5qbQVnxtntFc"  # testnet seed example

MEMO_CHUNK_BYTES = 512  # keep memo hex payload <= 1024 chars (~512 bytes)
CIRCUIT_MODES = ("AllowedOnly", "VerifySignOnly", "FreshnessOnly", "Full")
CANONICAL_CIRCUIT = {
    "AllowedOnly": "Allowed",
    "VerifySignOnly": "VerifySign",
    "FreshnessOnly": "Freshness",
    "Full": "Full",
}


# ---------------- Utility ----------------
def hex_utf8(s: str) -> str:
    return s.encode("utf-8").hex()

def ensure_parent_dir(path: str) -> None:
    """Create parent directory for a file path if it does not exist."""
    Path(path).parent.mkdir(parents=True, exist_ok=True)

def current_memory_mb() -> Optional[float]:
    if resource is None:
        return None
    usage = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    if sys.platform == "darwin":
        return usage / (1024 * 1024)
    return usage / 1024

def chunk_hex_payload(label: str, hex_payload: str, chunk_bytes: int = MEMO_CHUNK_BYTES) -> list[Memo]:
    if chunk_bytes <= 0:
        raise ValueError("chunk_bytes must be positive")
    chunk_chars = chunk_bytes * 2
    chunks = [hex_payload[i:i + chunk_chars] for i in range(0, len(hex_payload), chunk_chars)] or [""]
    total = len(chunks)
    memos = []
    for idx, chunk in enumerate(chunks, start=1):
        memo_label = f"{label}[{idx}/{total}]"
        memos.append(Memo(memo_type=hex_utf8(memo_label), memo_data=chunk))
    return memos

def gzip_to_hex(data: bytes, level: int = 9) -> str:
    """
    Compress bytes with gzip-compatible zlib and return hex string.
    """
    return zlib.compress(data, level).hex()

def default_artifact_paths(circuit_mode: str, setup_root: str) -> dict[str, str]:
    """
    Derive artifact paths from build_setup_multi.py output layout.
    Example:
      setup/Allowed/Allowed_js/Allowed.wasm, setup/Allowed/Allowed_final.zkey, etc.
    """
    canonical = CANONICAL_CIRCUIT.get(circuit_mode, circuit_mode)
    base = Path(setup_root) / canonical
    return {
        "base": str(base),
        "canonical": canonical,
        "wasm": str(base / f"{canonical}_js" / f"{canonical}.wasm"),
        "zkey": str(base / f"{canonical}_final.zkey"),
        "vk": str(base / f"{canonical}_verification_key.json"),
        "wtns": str(base / "witness.wtns"),
        "proof": str(base / "proof.json"),
        "public": str(base / "public.json"),
        "winput": str(base / "witness_input.json"),
    }


# ---------------- Witness Builder ----------------
def _build_allowed_dict(raw_id_path: str) -> dict:
    with open(raw_id_path, "r", encoding="utf-8") as f:
        raw = json.load(f)

    elems = (((raw or {}).get("subfiles") or {}).get("elements")) or {}
    dbb = elems.get("DBB", "")
    if not dbb or not re.fullmatch(r"\d{8}", dbb):
        raise ValueError("DBB missing or not 8 digits in raw_id.json")

    # Detect format: MMDDYYYY vs YYYYMMDD
    if int(dbb[:2]) > 12:
        year, month, day = int(dbb[:4]), int(dbb[4:6]), int(dbb[6:])
    else:
        month, day, year = int(dbb[:2]), int(dbb[2:4]), int(dbb[4:])

    today = datetime.date.today()
    current_year, current_month, current_day = today.year, today.month, today.day

    return {
        "currentYear":  current_year,
        "currentMonth": current_month,
        "currentDay":   current_day,
        "birthYear":    year,
        "birthMonth":   month,
        "birthDay":     day,
    }


def build_witness_from_raw_id(raw_id_path: str, out_witness_path: str) -> dict:
    """
    Reads AAMVA-style raw_id.json, extracts DOB,
    builds witness.json for the Allowed circuit:
      {
        "currentYear": ..., "currentMonth": ..., "currentDay": ...,
        "birthYear": ..., "birthMonth": ..., "birthDay": ...
      }
    Returns the witness dict.
    """
    witness = _build_allowed_dict(raw_id_path)
    ensure_parent_dir(out_witness_path)
    with open(out_witness_path, "w", encoding="utf-8") as f:
        json.dump(witness, f)

    print(f"witness written → {out_witness_path}")
    print("witness preview:", witness)
    return witness


def _limbs_6x43_le(x: int) -> list[int]:
    mask = (1 << 43) - 1
    return [(x >> (43 * i)) & mask for i in range(6)]


def _build_verifysign_dict_from_payload(signature_json_path: str) -> dict:
    with open(signature_json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    def limbs_from_nested(obj: dict, path: list[str]) -> list[int]:
        cur = obj
        for key in path:
            cur = cur[key]
        if not isinstance(cur, list) or len(cur) != 6:
            raise ValueError(f"Expected 6-limb array at {'.'.join(path)}")
        return cur

    if all(k in data for k in ("msgHash", "Qx", "Qy", "r", "s")):
        return {
            "msgHash": data["msgHash"],
            "Qx": data["Qx"],
            "Qy": data["Qy"],
            "r": data["r"],
            "s": data["s"],
        }
    return {
        "msgHash": limbs_from_nested(data, ["message_hash", "hash_limbs_6x43_le"]),
        "Qx":      limbs_from_nested(data, ["public_key", "Qx_limbs_6x43_le"]),
        "Qy":      limbs_from_nested(data, ["public_key", "Qy_limbs_6x43_le"]),
        "r":       limbs_from_nested(data, ["signature", "r_limbs_6x43_le"]),
        "s":       limbs_from_nested(data, ["signature", "s_limbs_6x43_le"]),
    }


def _build_verifysign_dict_from_keys(raw_id_path: str, priv_pem_path: str, pub_pem_path: str) -> dict:
    try:
        from cryptography.hazmat.primitives import hashes, serialization
        from cryptography.hazmat.primitives.asymmetric import ec
        from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature
    except ModuleNotFoundError as e:
        raise RuntimeError("cryptography package is required for signing; install it to use key-based witness generation") from e

    with open(raw_id_path, "r", encoding="utf-8") as f:
        payload = json.load(f)

    # Canonicalize payload (sorted keys, no whitespace)
    message_bytes = json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8")

    with open(priv_pem_path, "rb") as f:
        priv = serialization.load_pem_private_key(f.read(), password=None)
    with open(pub_pem_path, "rb") as f:
        pub = serialization.load_pem_public_key(f.read())

    # Sign payload (ECDSA P-256 / SHA-256)
    msg_hash = hashes.Hash(hashes.SHA256())
    msg_hash.update(message_bytes)
    digest = msg_hash.finalize()
    sig_der = priv.sign(message_bytes, ec.ECDSA(hashes.SHA256()))
    r, s = decode_dss_signature(sig_der)

    # Public key coordinates
    nums = pub.public_numbers()
    Qx, Qy = nums.x, nums.y

    msg_hash_int = int.from_bytes(digest, "big")
    return {
        "msgHash": _limbs_6x43_le(msg_hash_int),
        "Qx": _limbs_6x43_le(Qx),
        "Qy": _limbs_6x43_le(Qy),
        "r": _limbs_6x43_le(r),
        "s": _limbs_6x43_le(s),
    }


def build_verifysign_input(signature_json_path: str, out_input_path: str) -> dict:
    """
    Normalize signature payload into VerifySign circuit input schema.
    Returns the witness dict.
    """
    witness = _build_verifysign_dict_from_payload(signature_json_path)

    ensure_parent_dir(out_input_path)
    with open(out_input_path, "w", encoding="utf-8") as f:
        json.dump(witness, f)
    print(f"VerifySign input written → {out_input_path}")
    return witness


def _build_freshness_dict(
    source_json_path: str | None,
    *,
    nonce: int | None = None,
    valid_until: int | None = None,
    holder_ts: int | None = None,
    expected_nonce: int | None = None,
) -> dict:
    required = [
        "holderTimestamp",
        "nonce",
        "expectedNonce",
        "validUntil",
    ]

    if source_json_path:
        with open(source_json_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        witness = {}
        for key in required:
            if key not in data:
                raise ValueError(f"Freshness input missing field '{key}'")
            witness[key] = int(data[key])
        return witness

    if nonce is None or valid_until is None:
        raise ValueError("nonce and valid_until are required when no freshness JSON is provided")

    ts = int(holder_ts if holder_ts is not None else time.time())
    exp_nonce = int(expected_nonce if expected_nonce is not None else nonce)
    return {
        "holderTimestamp": ts,
        "nonce": int(nonce),
        "expectedNonce": exp_nonce,
        "validUntil": int(valid_until),
    }


def build_freshness_input(
    source_json_path: str | None,
    out_input_path: str,
    *,
    nonce: int | None = None,
    valid_until: int | None = None,
    holder_ts: int | None = None,
    expected_nonce: int | None = None,
) -> dict:
    """
    Normalize freshness witness file (holderTimestamp/nonce/expectedNonce/validUntil).
    Returns the witness dict.
    """
    witness = _build_freshness_dict(
        source_json_path,
        nonce=nonce,
        valid_until=valid_until,
        holder_ts=holder_ts,
        expected_nonce=expected_nonce,
    )

    ensure_parent_dir(out_input_path)
    with open(out_input_path, "w", encoding="utf-8") as f:
        json.dump(witness, f)
    print(f"Freshness input written → {out_input_path}")
    return witness


def build_full_input(
    raw_id_path: str,
    signature_json_path: str | None,
    priv_pem_path: str,
    pub_pem_path: str,
    freshness_json_path: str | None,
    *,
    fresh_nonce: int | None = None,
    fresh_valid_until: int | None = None,
    fresh_holder_ts: int | None = None,
    fresh_expected_nonce: int | None = None,
    out_input_path: str,
) -> dict:
    """
    Combine Allowed + VerifySign + Freshness inputs into one witness_input.json
    expected by Full.circom.
    """
    allowed = _build_allowed_dict(raw_id_path)
    if signature_json_path:
        verifysign = _build_verifysign_dict_from_payload(signature_json_path)
    else:
        verifysign = _build_verifysign_dict_from_keys(raw_id_path, priv_pem_path, pub_pem_path)
    fresh = _build_freshness_dict(
        freshness_json_path,
        nonce=fresh_nonce,
        valid_until=fresh_valid_until,
        holder_ts=fresh_holder_ts,
        expected_nonce=fresh_expected_nonce,
    )

    # Merge all inputs; keys are disjoint across the three sections.
    witness = {**allowed, **verifysign, **fresh}

    ensure_parent_dir(out_input_path)
    with open(out_input_path, "w", encoding="utf-8") as f:
        json.dump(witness, f)
    print(f"Full input written → {out_input_path}")
    return witness


# ---------------- AllowedOnly flow ----------------
def run_allowed_only(
    rpc: str,
    seed: str,
    raw_id_path: str = "test_inputs/dummy_id1/raw_id.json",
    wasm_path: str = "Allowed_js/Allowed.wasm",
    witness_input_json: str = "witness.json",
    zkey_path: str = "Allowed_final.zkey",
    vk_path: str = "verification_key.json",
    witness_path: str = "witness.wtns",
    proof_path: str = "proof.json",
    public_path: str = "public.json",
    proof_label: str = "Groth16-Proof",
):
    perf_start = time.perf_counter()

    # 0) Build witness.json from raw_id.json
    build_witness_from_raw_id(raw_id_path, witness_input_json)

    # 1) Generate witness (circom wasm)
    ensure_parent_dir(witness_path)
    generate_witness(wasm_path, witness_input_json, witness_path)

    # 2) Prove
    ensure_parent_dir(proof_path)
    ensure_parent_dir(public_path)
    groth16_prove(zkey_path, witness_path, proof_path, public_path)

    allowed = read_allowed_from_public(public_path)

    # 4) Post to XRPL (AccountSet + memos)
    client = JsonRpcClient(rpc)
    wallet = Wallet.from_seed(seed)

    proof_json  = open(proof_path,  "r", encoding="utf-8").read()
    public_json = open(public_path, "r", encoding="utf-8").read()
    proof_bytes = proof_json.encode("utf-8")
    # Compress full payloads to fit XRPL memo limits, still sending full data on-chain
    proof_hex_gz = gzip_to_hex(proof_bytes)
    proof_chunks = chunk_hex_payload("ProofGZ", proof_hex_gz)

    public_bytes = public_json.encode("utf-8")
    public_hex_gz = gzip_to_hex(public_bytes)
    public_chunks = chunk_hex_payload("PublicGZ", public_hex_gz)

    memos = [
        Memo(memo_type=hex_utf8("ZKProof"),   memo_data=hex_utf8(proof_label)),
        Memo(memo_type=hex_utf8("Allowed"),   memo_data=hex_utf8(allowed)),
    ]
    memos.extend(proof_chunks)
    memos.extend(public_chunks)

    tx = AccountSet(account=wallet.address, memos=memos)
    res = submit_and_wait(tx, client, wallet)

    tx_hash = res.result.get("tx_json", {}).get("hash") or res.result.get("hash")
    print("XRPL tx hash:", tx_hash)
    runtime = time.perf_counter() - perf_start
    proof_chunk_count = len(proof_chunks)
    proof_memo_bytes = max((len(m.memo_data) // 2 for m in proof_chunks), default=0)
    mem_mb = current_memory_mb()
    print("--- Performance Metrics ---")
    print(f"Prover runtime       : {runtime:.2f}s")
    print(f"Proof JSON size      : {len(proof_bytes)} bytes ({proof_chunk_count} memo chunk(s), max {proof_memo_bytes} bytes each)")
    if mem_mb is not None:
        print(f"Process max RSS      : {mem_mb:.2f} MB")
    else:
        print("Process max RSS      : unavailable (resource module missing)")


# ---------------- Placeholder stubs ----------------
def run_verifysign_only(
    rpc: str,
    seed: str,
    signature_json_path: str,
    raw_id_path: str,
    priv_pem_path: str,
    pub_pem_path: str,
    wasm_path: str,
    witness_input_json: str,
    zkey_path: str,
    vk_path: str,
    witness_path: str,
    proof_path: str,
    public_path: str,
    proof_label: str,
):
    perf_start = time.perf_counter()
    if signature_json_path:
        build_verifysign_input(signature_json_path, witness_input_json)
    else:
        print("Generating witness")
        witness = _build_verifysign_dict_from_keys(raw_id_path, priv_pem_path, pub_pem_path)
        ensure_parent_dir(witness_input_json)
        with open(witness_input_json, "w", encoding="utf-8") as f:
            json.dump(witness, f)
        print(f"VerifySign input (signed from keys) written → {witness_input_json}")

    ensure_parent_dir(witness_path)
    generate_witness(wasm_path, witness_input_json, witness_path)

    ensure_parent_dir(proof_path)
    ensure_parent_dir(public_path)
    groth16_prove(zkey_path, witness_path, proof_path, public_path)

    client = JsonRpcClient(rpc)
    wallet = Wallet.from_seed(seed)

    proof_json  = open(proof_path,  "r", encoding="utf-8").read()
    public_json = open(public_path, "r", encoding="utf-8").read()
    proof_bytes = proof_json.encode("utf-8")
    # Compress full payloads to fit XRPL memo limits, still sending full data on-chain
    proof_hex_gz = gzip_to_hex(proof_bytes)
    public_hex_gz = gzip_to_hex(public_json.encode("utf-8"))
    proof_chunks = chunk_hex_payload("ProofGZ", proof_hex_gz)
    public_chunks = chunk_hex_payload("PublicGZ", public_hex_gz)

    memos = [
        Memo(memo_type=hex_utf8("ZKProof"),    memo_data=hex_utf8(proof_label)),
    ]
    memos.extend(proof_chunks)
    memos.extend(public_chunks)

    tx = AccountSet(account=wallet.address, memos=memos)
    res = submit_and_wait(tx, client, wallet)

    tx_hash = res.result.get("tx_json", {}).get("hash") or res.result.get("hash")
    print("XRPL tx hash:", tx_hash)
    runtime = time.perf_counter() - perf_start
    proof_chunk_count = len(proof_chunks)
    proof_memo_bytes = max((len(m.memo_data) // 2 for m in proof_chunks), default=0)
    mem_mb = current_memory_mb()
    print("--- Performance Metrics ---")
    print(f"Prover runtime       : {runtime:.2f}s")
    print(f"Proof JSON size      : {len(proof_bytes)} bytes ({proof_chunk_count} memo chunk(s), max {proof_memo_bytes} bytes each)")
    if mem_mb is not None:
        print(f"Process max RSS      : {mem_mb:.2f} MB")
    else:
        print("Process max RSS      : unavailable (resource module missing)")


def run_freshness_only(
    rpc: str,
    seed: str,
    freshness_json_path: str | None,
    fresh_nonce: int | None,
    fresh_valid_until: int | None,
    fresh_holder_ts: int | None,
    fresh_expected_nonce: int | None,
    wasm_path: str,
    witness_input_json: str,
    zkey_path: str,
    vk_path: str,
    witness_path: str,
    proof_path: str,
    public_path: str,
    proof_label: str,
):
    perf_start = time.perf_counter()
    build_freshness_input(
        freshness_json_path,
        witness_input_json,
        nonce=fresh_nonce,
        valid_until=fresh_valid_until,
        holder_ts=fresh_holder_ts,
        expected_nonce=fresh_expected_nonce,
    )

    ensure_parent_dir(witness_path)
    generate_witness(wasm_path, witness_input_json, witness_path)

    ensure_parent_dir(proof_path)
    ensure_parent_dir(public_path)
    groth16_prove(zkey_path, witness_path, proof_path, public_path)

    client = JsonRpcClient(rpc)
    wallet = Wallet.from_seed(seed)

    proof_json  = open(proof_path,  "r", encoding="utf-8").read()
    public_json = open(public_path, "r", encoding="utf-8").read()
    proof_bytes = proof_json.encode("utf-8")
    proof_hex_gz = gzip_to_hex(proof_bytes)
    proof_chunks = chunk_hex_payload("ProofGZ", proof_hex_gz)

    public_hex_gz = gzip_to_hex(public_json.encode("utf-8"))
    public_chunks = chunk_hex_payload("PublicGZ", public_hex_gz)

    memos = [
        Memo(memo_type=hex_utf8("ZKProof"),    memo_data=hex_utf8(proof_label)),
    ]
    memos.extend(proof_chunks)
    memos.extend(public_chunks)

    tx = AccountSet(account=wallet.address, memos=memos)
    res = submit_and_wait(tx, client, wallet)

    tx_hash = res.result.get("tx_json", {}).get("hash") or res.result.get("hash")
    runtime = time.perf_counter() - perf_start
    proof_chunk_count = len(proof_chunks)
    proof_memo_bytes = max((len(m.memo_data) // 2 for m in proof_chunks), default=0)
    mem_mb = current_memory_mb()

    print("XRPL tx hash:", tx_hash)
    print("--- Performance Metrics ---")
    print(f"Prover runtime       : {runtime:.2f}s")
    print(f"Proof JSON size      : {len(proof_bytes)} bytes ({proof_chunk_count} memo chunk(s), max {proof_memo_bytes} bytes each)")
    if mem_mb is not None:
        print(f"Process max RSS      : {mem_mb:.2f} MB")
    else:
        print("Process max RSS      : unavailable (resource module missing)")

def run_full(
    rpc: str,
    seed: str,
    raw_id_path: str,
    signature_json_path: str | None,
    priv_pem_path: str,
    pub_pem_path: str,
    freshness_json_path: str | None,
    fresh_nonce: int | None,
    fresh_valid_until: int | None,
    fresh_holder_ts: int | None,
    fresh_expected_nonce: int | None,
    wasm_path: str,
    witness_input_json: str,
    zkey_path: str,
    vk_path: str,
    witness_path: str,
    proof_path: str,
    public_path: str,
    proof_label: str,
):
    perf_start = time.perf_counter()
    build_full_input(
        raw_id_path,
        signature_json_path,
        priv_pem_path,
        pub_pem_path,
        freshness_json_path,
        fresh_nonce=fresh_nonce,
        fresh_valid_until=fresh_valid_until,
        fresh_holder_ts=fresh_holder_ts,
        fresh_expected_nonce=fresh_expected_nonce,
        out_input_path=witness_input_json,
    )

    ensure_parent_dir(witness_path)
    generate_witness(wasm_path, witness_input_json, witness_path)

    ensure_parent_dir(proof_path)
    ensure_parent_dir(public_path)
    groth16_prove(zkey_path, witness_path, proof_path, public_path)

    client = JsonRpcClient(rpc)
    wallet = Wallet.from_seed(seed)

    proof_json  = open(proof_path,  "r", encoding="utf-8").read()
    public_json = open(public_path, "r", encoding="utf-8").read()
    proof_bytes = proof_json.encode("utf-8")
    proof_hex_gz = gzip_to_hex(proof_bytes)
    proof_chunks = chunk_hex_payload("ProofGZ", proof_hex_gz)

    public_hex_gz = gzip_to_hex(public_json.encode("utf-8"))
    public_chunks = chunk_hex_payload("PublicGZ", public_hex_gz)

    memos = [
        Memo(memo_type=hex_utf8("ZKProof"),    memo_data=hex_utf8(proof_label)),
    ]
    memos.extend(proof_chunks)
    memos.extend(public_chunks)

    tx = AccountSet(account=wallet.address, memos=memos)
    res = submit_and_wait(tx, client, wallet)

    tx_hash = res.result.get("tx_json", {}).get("hash") or res.result.get("hash")
    runtime = time.perf_counter() - perf_start
    proof_chunk_count = len(proof_chunks)
    proof_memo_bytes = max((len(m.memo_data) // 2 for m in proof_chunks), default=0)
    mem_mb = current_memory_mb()

    print("XRPL tx hash:", tx_hash)
    print("--- Performance Metrics ---")
    print(f"Prover runtime       : {runtime:.2f}s")
    print(f"Proof JSON size      : {len(proof_bytes)} bytes ({proof_chunk_count} memo chunk(s), max {proof_memo_bytes} bytes each)")
    if mem_mb is not None:
        print(f"Process max RSS      : {mem_mb:.2f} MB")
    else:
        print("Process max RSS      : unavailable (resource module missing)")


# ---------------- CLI ----------------
def parse_args():
    p = argparse.ArgumentParser(description="Prover: prove circuits and post proof to XRPL.")
    p.add_argument("-c", "--circuit", choices=CIRCUIT_MODES, required=True,
                   help="Which proving flow to run. Maps to Allowed / VerifySign / Freshness / Full circuits.")
    p.add_argument("--rpc",  default=os.environ.get("XRPL_RPC_URL", DEFAULT_RPC))
    p.add_argument("--seed", default=os.environ.get("XRPL_TESTNET_SEED", DEFAULT_SEED))
    p.add_argument("--setup-root", default="setup",
                   help="Root directory containing build_setup_multi.py outputs (default: setup/).")
    p.add_argument("--raw-id", default="test_inputs/dummy_id1/raw_id.json")
    p.add_argument("--vs-json", default=None,
                   help="Optional precomputed signature witness JSON for VerifySign.")
    p.add_argument("--vs-priv", default="test_inputs/dummy_id1/private.pem",
                   help="Private key PEM for VerifySign witness generation.")
    p.add_argument("--vs-pub", default="test_inputs/dummy_id1/public.pem",
                   help="Public key PEM for VerifySign witness generation.")
    p.add_argument("--fresh-json", default=None,
                   help="Optional source JSON containing freshness inputs "
                        "(holderTimestamp/nonce/expectedNonce/validUntil). "
                        "If omitted, nonce/validUntil must be provided via flags.")
    p.add_argument("--fresh-nonce", type=int, default=None,
                   help="Nonce challenge used for Freshness inputs (sets both nonce and expectedNonce when no JSON provided).")
    p.add_argument("--fresh-valid-until", type=int, default=None,
                   help="Expiration timestamp/slot for Freshness inputs when no JSON provided.")
    p.add_argument("--fresh-holder-ts", type=int, default=None,
                   help="Override holderTimestamp for Freshness inputs; defaults to current unix time if omitted.")
    p.add_argument("--wasm",   default=None, help="Override path to circuit WASM (default inferred from setup-root).")
    p.add_argument("--winput", default=None, help="Override path to witness input JSON.")
    p.add_argument("--zkey",   default=None, help="Override path to final zkey.")
    p.add_argument("--vk",     default=None, help="Override path to verification key JSON.")
    p.add_argument("--wtns",   default=None, help="Override path to witness.wtns output.")
    p.add_argument("--proof",  default=None, help="Override path to proof.json output.")
    p.add_argument("--public", default=None, help="Override path to public.json output.")
    p.add_argument("--label",  default="Groth16-Proof")
    return p.parse_args()


def main():
    args = parse_args()
    defaults = default_artifact_paths(args.circuit, args.setup_root)

    def require_file(path: str, label: str) -> None:
        if not Path(path).exists():
            raise SystemExit(f"Required {label} not found: {path}")

    # Validate required inputs per mode
    def require_freshness_inputs() -> None:
        if args.fresh_json:
            require_file(args.fresh_json, "freshness input JSON")
        elif args.fresh_nonce is None or args.fresh_valid_until is None:
            raise SystemExit("Freshness inputs required: provide --fresh-json or both --fresh-nonce and --fresh-valid-until")

    if args.circuit == "AllowedOnly":
        require_file(args.raw_id, "raw ID JSON")
    elif args.circuit == "FreshnessOnly":
        require_freshness_inputs()
    elif args.circuit == "VerifySignOnly":
        if args.vs_json:
            require_file(args.vs_json, "VerifySign witness JSON")
        else:
            # Require all inputs to build witness from keys
            require_file(args.raw_id, "raw ID JSON")
            require_file(args.vs_priv, "private key PEM (--vs-priv)")
            require_file(args.vs_pub, "public key PEM (--vs-pub)")
    elif args.circuit == "Full":
        require_file(args.raw_id, "raw ID JSON")
        require_freshness_inputs()
        if args.vs_json:
            require_file(args.vs_json, "VerifySign witness JSON")
        else:
            require_file(args.vs_priv, "private key PEM (--vs-priv)")
            require_file(args.vs_pub, "public key PEM (--vs-pub)")

    wasm_path   = args.wasm   or defaults["wasm"]
    winput_path = args.winput or defaults["winput"]
    zkey_path   = args.zkey   or defaults["zkey"]
    vk_path     = args.vk     or defaults["vk"]
    wtns_path   = args.wtns   or defaults["wtns"]
    proof_path  = args.proof  or defaults["proof"]
    public_path = args.public or defaults["public"]

    if args.circuit == "AllowedOnly":
        run_allowed_only(
            rpc=args.rpc,
            seed=args.seed,
            raw_id_path=args.raw_id,
            wasm_path=wasm_path,
            witness_input_json=winput_path,
            zkey_path=zkey_path,
            vk_path=vk_path,
            witness_path=wtns_path,
            proof_path=proof_path,
            public_path=public_path,
            proof_label=args.label,
        )
    elif args.circuit == "VerifySignOnly":
        run_verifysign_only(
            rpc=args.rpc,
            seed=args.seed,
            signature_json_path=args.vs_json,
            raw_id_path=args.raw_id,
            priv_pem_path=args.vs_priv,
            pub_pem_path=args.vs_pub,
            wasm_path=wasm_path,
            witness_input_json=winput_path,
            zkey_path=zkey_path,
            vk_path=vk_path,
            witness_path=wtns_path,
            proof_path=proof_path,
            public_path=public_path,
            proof_label=args.label,
        )
    elif args.circuit == "FreshnessOnly":
        run_freshness_only(
            rpc=args.rpc,
            seed=args.seed,
            freshness_json_path=args.fresh_json,
            fresh_nonce=args.fresh_nonce,
            fresh_valid_until=args.fresh_valid_until,
            fresh_holder_ts=args.fresh_holder_ts,
            fresh_expected_nonce=None,
            wasm_path=wasm_path,
            witness_input_json=winput_path,
            zkey_path=zkey_path,
            vk_path=vk_path,
            witness_path=wtns_path,
            proof_path=proof_path,
            public_path=public_path,
            proof_label=args.label,
        )
    elif args.circuit == "Full":
        run_full(
            rpc=args.rpc,
            seed=args.seed,
            raw_id_path=args.raw_id,
            signature_json_path=args.vs_json,
            priv_pem_path=args.vs_priv,
            pub_pem_path=args.vs_pub,
            freshness_json_path=args.fresh_json,
            fresh_nonce=args.fresh_nonce,
            fresh_valid_until=args.fresh_valid_until,
            fresh_holder_ts=args.fresh_holder_ts,
            fresh_expected_nonce=None,
            wasm_path=wasm_path,
            witness_input_json=winput_path,
            zkey_path=zkey_path,
            vk_path=vk_path,
            witness_path=wtns_path,
            proof_path=proof_path,
            public_path=public_path,
            proof_label=args.label,
        )
    else:
        raise SystemExit(f"Unknown circuit mode: {args.circuit}")


if __name__ == "__main__":
    main()