#!/usr/bin/env python3
"""Independent inclusion-proof verification. No Cyphrex imports.

Allowed: rfc8785, cryptography, hashlib, stdlib.

Checks, in order:
1. Ed25519 over RFC 8785 JCS of the signed core
2. Recompute the Merkle root from one event and its audit path
   (RFC 6962 promote, leaf 0x00, node 0x01, JCS of the closed leaf)
3. Confirm that root appears in the named Solana transaction
   (live RPC, or --tx-fixture when RPC is absent)

Default event: the first event whose inclusion proof has anchored=true.
Pass --event-index to override. The chosen index is always printed.
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import json
import sys
import urllib.error
import urllib.request

import rfc8785
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives.serialization import load_pem_public_key

STRIP = ("sha256", "signature", "algorithm", "publicKeyUrl")
LEAF_PREFIX = b"\x00"
NODE_PREFIX = b"\x01"
DEFAULT_RPC = "https://api.devnet.solana.com"


def canonical_core(report: dict) -> bytes:
    core = {k: v for k, v in report.items() if k not in STRIP}
    canonical = rfc8785.dumps(core)
    if isinstance(canonical, str):
        canonical = canonical.encode("utf-8")
    return canonical


def verify_signature(report: dict, pem: bytes) -> tuple[str, bool, str]:
    canonical = canonical_core(report)
    computed = hashlib.sha256(canonical).hexdigest()
    stored = str(report.get("sha256") or "")
    hash_ok = computed == stored.lower()
    signature_status = "ok"
    try:
        key = load_pem_public_key(pem)
        key.verify(base64.b64decode(report["signature"]), canonical)
    except InvalidSignature:
        signature_status = "InvalidSignature"
    except Exception as exc:
        signature_status = f"error:{exc}"
    return signature_status, hash_ok, computed


def closed_leaf(event: dict) -> dict:
    return {
        "id": "" if event.get("id") is None else str(event["id"]),
        "action_type": "" if event.get("action_type") is None else str(event["action_type"]),
        "created_at": "" if event.get("created_at") is None else str(event["created_at"]),
        "allowed": event.get("allowed") is True or event.get("allowed") == "true",
        "agent_id": "" if event.get("agent_id") is None else str(event["agent_id"]),
        "severity": str(event["severity"]) if event.get("severity") else "",
        "url": str(event["url"]) if event.get("url") else "",
    }


def hash_event(event: dict) -> str:
    canonical = rfc8785.dumps(closed_leaf(event))
    if isinstance(canonical, str):
        canonical = canonical.encode("utf-8")
    return hashlib.sha256(LEAF_PREFIX + canonical).hexdigest()


def hash_node(left_hex: str, right_hex: str) -> str:
    return hashlib.sha256(NODE_PREFIX + bytes.fromhex(left_hex) + bytes.fromhex(right_hex)).hexdigest()


def verify_inclusion(leaf_hash: str, proof: dict) -> bool:
    if not proof or proof.get("anchored") is False:
        return False
    path = proof.get("path")
    if not isinstance(path, list):
        return False
    try:
        index = int(proof["leafIndex"])
        size = int(proof["treeSize"])
    except (KeyError, TypeError, ValueError):
        return False
    root = str(proof.get("merkleRoot") or "")
    if size < 1 or index < 0 or index >= size or not root:
        return False
    digest = leaf_hash
    path_i = 0
    while size > 1:
        is_right = index % 2 == 1
        sibling_index = index - 1 if is_right else index + 1
        if sibling_index < size:
            if path_i >= len(path):
                return False
            step = path[path_i]
            path_i += 1
            if not isinstance(step, dict) or not step.get("hash"):
                return False
            expected = "left" if is_right else "right"
            if step.get("position") and step["position"] != expected:
                return False
            if is_right:
                digest = hash_node(step["hash"], digest)
            else:
                digest = hash_node(digest, step["hash"])
        index //= 2
        size = (size + 1) // 2
    if path_i != len(path):
        return False
    return digest == root


def tx_contains_root(tx_obj: object, root_hex: str) -> bool:
    root_hex = root_hex.lower()
    root_bytes = bytes.fromhex(root_hex)

    def walk(value: object) -> bool:
        if isinstance(value, str):
            if root_hex in value.lower():
                return True
            try:
                pad = value + "=" * ((4 - len(value) % 4) % 4)
                blob = base64.b64decode(pad, validate=False)
                if root_bytes in blob:
                    return True
            except Exception:
                pass
            return False
        if isinstance(value, list):
            return any(walk(item) for item in value)
        if isinstance(value, dict):
            return any(walk(item) for item in value.values())
        return False

    raw = json.dumps(tx_obj)
    if root_hex in raw.lower():
        return True
    return walk(tx_obj)


def fetch_transaction(signature: str, rpc: str) -> object:
    payload = json.dumps(
        {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "getTransaction",
            "params": [
                signature,
                {
                    "encoding": "base64",
                    "maxSupportedTransactionVersion": 0,
                    "commitment": "confirmed",
                },
            ],
        }
    ).encode("utf-8")
    req = urllib.request.Request(
        rpc,
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=20) as handle:
        return json.load(handle)


def check_chain(proof: dict, rpc: str, fixture: object | None) -> tuple[str, str]:
    sig = str(proof.get("txSignature") or "").strip()
    root = str(proof.get("merkleRoot") or "").strip().lower()
    if not sig:
        return "fail", "no transaction signature"
    if len(root) != 64:
        return "fail", "invalid merkle root"
    tx_obj = fixture
    if tx_obj is None:
        try:
            tx_obj = fetch_transaction(sig, rpc)
        except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc:
            return "skipped", str(exc)
    if tx_contains_root(tx_obj, root):
        return "ok", "root found in transaction"
    return "fail", "root not found in transaction"


def proof_for_event(events: list, proofs: list, event_index: int) -> dict | None:
    if event_index < 0 or event_index >= len(events):
        return None
    event = events[event_index]
    proof = next((p for p in proofs if str(p.get("eventId")) == str(event.get("id"))), None)
    if proof is None and event_index < len(proofs):
        proof = proofs[event_index]
    return proof


def first_anchored_index(events: list, proofs: list) -> int:
    for i, _event in enumerate(events):
        proof = proof_for_event(events, proofs, i)
        if proof and proof.get("anchored") is True:
            return i
    raise SystemExit("no anchored inclusion proof in this package")


def first_anchored_proof(proofs: list) -> dict | None:
    for proof in proofs:
        if proof and proof.get("anchored") is True:
            return proof
    return None


def replace_proof_for_index(report: dict, event_index: int, new_proof: dict) -> None:
    events = report.get("events") or []
    proofs = report.get("inclusionProofs") or []
    event = events[event_index]
    event_id = str(event.get("id"))
    for i, proof in enumerate(proofs):
        if str(proof.get("eventId")) == event_id:
            proofs[i] = new_proof
            return
    if event_index < len(proofs):
        proofs[event_index] = new_proof
        return
    proofs.append(new_proof)
    report["inclusionProofs"] = proofs


def main() -> int:
    parser = argparse.ArgumentParser(description="Verify Cyphrex inclusion proofs independently.")
    parser.add_argument("report", help="Path to signed JSON package")
    parser.add_argument("pem", help="Path to Ed25519 public key PEM")
    parser.add_argument(
        "--event-index",
        type=int,
        default=None,
        help="Event index to check. Default: first event with anchored=true",
    )
    parser.add_argument("--tamper-event-url", default=None)
    parser.add_argument("--reorder-sibling", action="store_true")
    parser.add_argument("--foreign-proof", default=None, help="Path to another package; swap in its first anchored proof")
    parser.add_argument("--rpc", default=DEFAULT_RPC)
    parser.add_argument("--tx-fixture", default=None)
    parser.add_argument("--skip-chain", action="store_true")
    args = parser.parse_args()

    with open(args.report, encoding="utf-8") as handle:
        report = json.load(handle)
    with open(args.pem, "rb") as handle:
        pem = handle.read()

    events = report.get("events") or []
    proofs = report.get("inclusionProofs") or []
    if args.event_index is None:
        event_index = first_anchored_index(events, proofs)
        index_source = "first_anchored"
    else:
        event_index = args.event_index
        index_source = "flag"
        if event_index < 0 or event_index >= len(events):
            raise SystemExit(f"event index {event_index} out of range")

    event_id = events[event_index].get("id")
    print(f"event_index={event_index}")
    print(f"event_index_source={index_source}")
    print(f"event_id={event_id}")

    if args.tamper_event_url is not None:
        report["events"][event_index]["url"] = args.tamper_event_url
    if args.reorder_sibling:
        proof = proof_for_event(report.get("events") or [], report.get("inclusionProofs") or [], event_index)
        if not proof or not isinstance(proof.get("path"), list):
            print("VERIFY_FAIL")
            print("reason=reorder_path_missing")
            return 1
        proof["path"] = list(reversed(proof["path"]))
    if args.foreign_proof:
        with open(args.foreign_proof, encoding="utf-8") as handle:
            other = json.load(handle)
        foreign = first_anchored_proof(other.get("inclusionProofs") or [])
        if not foreign:
            print("VERIFY_FAIL")
            print("reason=foreign_proof_missing")
            return 1
        replace_proof_for_index(report, event_index, foreign)

    signature_status, hash_ok, computed = verify_signature(report, pem)
    event, proof = events[event_index], proof_for_event(
        report.get("events") or [], report.get("inclusionProofs") or [], event_index
    )
    if proof is None:
        raise SystemExit("no inclusion proof for selected event")
    inclusion_ok = verify_inclusion(hash_event(event), proof)

    chain_status = "skipped"
    chain_detail = "not checked"
    if args.skip_chain:
        chain_detail = "skipped by flag"
    else:
        fixture = None
        if args.tx_fixture:
            with open(args.tx_fixture, encoding="utf-8") as handle:
                fixture = json.load(handle)
        chain_status, chain_detail = check_chain(proof, args.rpc, fixture)

    print(f"signature={signature_status}")
    print(f"sha256={'ok' if hash_ok else 'mismatch'}")
    print(f"sha256_computed={computed}")
    print(f"inclusion={'ok' if inclusion_ok else 'fail'}")
    print(f"chain={chain_status}")
    print(f"chain_detail={chain_detail}")
    print(f"merkle_root={proof.get('merkleRoot')}")
    print(f"leaf_index={proof.get('leafIndex')}")
    print(f"tree_size={proof.get('treeSize')}")

    chain_ok = chain_status == "ok" or (args.skip_chain and chain_status == "skipped")
    if signature_status == "ok" and hash_ok and inclusion_ok and chain_ok:
        print("VERIFY_OK")
        return 0

    reasons = []
    if signature_status != "ok":
        reasons.append("signature_invalid")
    if not hash_ok:
        reasons.append("sha256_mismatch")
    if not inclusion_ok:
        reasons.append("inclusion_fail")
    if not chain_ok:
        reasons.append(f"chain_{chain_status}")
    print("VERIFY_FAIL")
    print(f"reason={','.join(reasons)}")
    return 1


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