#!/usr/bin/env python3
"""Reproduce the A-01 Receipt / Signing / Recovery drill on Bitcoin Core regtest.

Safety properties:
- Hard-coded regtest network; no real bitcoin is used.
- P2P networking, listening, peer discovery, and DNS are disabled.
- A fresh temporary data directory is removed after the node stops.
- Addresses, transaction IDs, descriptors, wallet backups, and RPC credentials are
  never printed. Only sanitized measurements are written to stdout.

Example:
    python custody-proof-lab.py --bin-dir /opt/bitcoin-31.1/bin
"""

from __future__ import annotations

import argparse
import json
import os
import re
import socket
import subprocess
import sys
import tempfile
import time
from decimal import Decimal
from pathlib import Path
from typing import Any


EXPECTED_VERSION = "v31.1.0"
TEST_AMOUNT = Decimal("0.01000000")
RESTORED_SPEND = Decimal("0.00200000")


class LabError(RuntimeError):
    """A sanitized lab failure that does not expose command arguments or output."""

    def __init__(self, stage: str, message: str, rpc_code: int | str | None = None):
        super().__init__(message)
        self.stage = stage
        self.rpc_code = rpc_code


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--bin-dir",
        required=True,
        type=Path,
        help="Directory containing the verified Bitcoin Core 31.1 binaries.",
    )
    return parser.parse_args()


def executable(bin_dir: Path, name: str) -> Path:
    suffix = ".exe" if os.name == "nt" else ""
    candidate = (bin_dir / f"{name}{suffix}").resolve()
    if not candidate.is_file():
        raise LabError("preflight", f"Required executable is missing: {name}{suffix}")
    return candidate


def free_loopback_port() -> int:
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        return int(sock.getsockname()[1])


def sum_balance_buckets(balances: dict[str, Any]) -> Decimal:
    total = Decimal("0")
    for bucket_name in ("mine", "watchonly"):
        bucket = balances.get(bucket_name, {})
        total += Decimal(str(bucket.get("trusted", 0)))
        total += Decimal(str(bucket.get("untrusted_pending", 0)))
    return total


def run_lab(bin_dir: Path) -> dict[str, Any]:
    bitcoind = executable(bin_dir, "bitcoind")
    bitcoin_cli = executable(bin_dir, "bitcoin-cli")
    version = subprocess.run(
        [str(bitcoind), "--version"],
        check=True,
        text=True,
        capture_output=True,
    ).stdout.splitlines()[0]
    started = time.monotonic()
    stage = "node-start"

    with tempfile.TemporaryDirectory(prefix="cora-a01-regtest-") as temp_dir:
        datadir = Path(temp_dir)
        backup_path = datadir / "receiver-preflight-backup.dat"
        rpc_port = free_loopback_port()
        node = subprocess.Popen(
            [
                str(bitcoind),
                "-regtest",
                f"-datadir={datadir}",
                f"-rpcport={rpc_port}",
                "-rpcbind=127.0.0.1",
                "-rpcallowip=127.0.0.1",
                "-server=1",
                "-daemon=0",
                "-networkactive=0",
                "-listen=0",
                "-maxconnections=0",
                "-dnsseed=0",
                "-fixedseeds=0",
                "-discover=0",
                "-natpmp=0",
                "-dns=0",
                "-listenonion=0",
                "-dbcache=64",
                "-fallbackfee=0.00001000",
                "-printtoconsole=0",
            ],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        base_cli = [
            str(bitcoin_cli),
            "-regtest",
            f"-datadir={datadir}",
            f"-rpcport={rpc_port}",
        ]

        def cli(
            method: str,
            *arguments: object,
            wallet: str | None = None,
            parse_json: bool = True,
        ) -> Any:
            command = base_cli.copy()
            if wallet:
                command.append(f"-rpcwallet={wallet}")
            command.extend([method, *[str(argument) for argument in arguments]])
            result = subprocess.run(command, text=True, capture_output=True)
            if result.returncode != 0:
                match = re.search(r"error code:\s*(-?\d+)", result.stderr or "")
                code: int | str = int(match.group(1)) if match else "rpc-error"
                raise LabError(stage, f"RPC failed: {method}", code)
            text = result.stdout.strip()
            if not parse_json:
                return text
            return json.loads(text) if text else None

        try:
            deadline = time.monotonic() + 60
            while True:
                try:
                    cli("getblockchaininfo")
                    break
                except LabError:
                    if node.poll() is not None:
                        raise LabError(stage, "bitcoind exited before RPC became ready")
                    if time.monotonic() >= deadline:
                        raise LabError(stage, "RPC readiness timeout")
                    time.sleep(0.25)

            stage = "wallet-setup"
            cli("createwallet", "funder")
            cli("createwallet", "receiver")
            receiver_address = cli("getnewaddress", wallet="receiver", parse_json=False)
            cli("backupwallet", backup_path, wallet="receiver")

            # Import the receive address before funding it. This avoids a historical
            # rescan while still producing a wallet that can observe but not sign.
            cli("createwallet", "watchproof", "true", "true")
            watch_descriptor = cli("getdescriptorinfo", f"addr({receiver_address})")["descriptor"]
            import_result = cli(
                "importdescriptors",
                json.dumps(
                    [{"desc": watch_descriptor, "timestamp": "now"}],
                    separators=(",", ":"),
                ),
                wallet="watchproof",
            )
            if not import_result[0].get("success"):
                raise LabError(stage, "Watch-only descriptor import failed")

            stage = "funding"
            mining_address = cli("getnewaddress", wallet="funder", parse_json=False)
            cli("generatetoaddress", 101, mining_address)
            cli(
                "sendtoaddress",
                receiver_address,
                f"{TEST_AMOUNT:.8f}",
                wallet="funder",
                parse_json=False,
            )
            cli("generatetoaddress", 1, mining_address)

            stage = "receipt-proof"
            receipt_balance = sum_balance_buckets(cli("getbalances", wallet="receiver"))
            watch_balance = sum_balance_buckets(cli("getbalances", wallet="watchproof"))
            watch_info = cli("getwalletinfo", wallet="watchproof")
            watch_utxos = cli("listunspent", 1, wallet="watchproof")
            if len(watch_utxos) != 1:
                raise LabError(stage, "Unexpected watch-only UTXO count")

            stage = "signing-proof"
            destination = cli("getnewaddress", wallet="funder", parse_json=False)
            source = watch_utxos[0]
            unsigned_psbt = cli(
                "createpsbt",
                json.dumps(
                    [{"txid": source["txid"], "vout": source["vout"]}],
                    separators=(",", ":"),
                ),
                json.dumps([{destination: 0.009}], separators=(",", ":")),
                parse_json=False,
            )
            processed_psbt = cli(
                "walletprocesspsbt",
                unsigned_psbt,
                "true",
                "ALL",
                "true",
                wallet="watchproof",
            )
            direct_spend_error_code: int | str | None = None
            try:
                cli(
                    "sendtoaddress",
                    destination,
                    "0.00100000",
                    wallet="watchproof",
                )
            except LabError as expected_error:
                direct_spend_error_code = expected_error.rpc_code

            stage = "recovery-proof"
            cli("unloadwallet", "receiver")
            cli("restorewallet", "receiver-restored", backup_path)
            restored_balance = sum_balance_buckets(
                cli("getbalances", wallet="receiver-restored")
            )
            restored_txid = cli(
                "sendtoaddress",
                destination,
                f"{RESTORED_SPEND:.8f}",
                wallet="receiver-restored",
                parse_json=False,
            )
            cli("generatetoaddress", 1, mining_address)
            restored_confirmations = int(
                cli(
                    "gettransaction",
                    restored_txid,
                    wallet="receiver-restored",
                )["confirmations"]
            )

            stage = "final-isolation-check"
            network_info = cli("getnetworkinfo")
            checks = {
                "exact_core_version": EXPECTED_VERSION in version,
                "node_network_disabled": network_info.get("networkactive") is False,
                "node_has_zero_connections": int(network_info.get("connections", -1)) == 0,
                "receipt_balance_matches": receipt_balance == TEST_AMOUNT,
                "watch_only_sees_same_balance": watch_balance == receipt_balance,
                "watch_only_private_keys_disabled": watch_info.get("private_keys_enabled") is False,
                "watch_only_cannot_complete_psbt": processed_psbt.get("complete") is False,
                "watch_only_direct_spend_rejected": direct_spend_error_code == -4,
                "backup_restores_same_balance": restored_balance == receipt_balance,
                "restored_wallet_signs_confirmed_spend": restored_confirmations >= 1,
            }
            passed = sum(bool(value) for value in checks.values())
            return {
                "schema_version": 1,
                "status": "PASS" if passed == len(checks) else "FAIL",
                "bitcoin_core": version,
                "network": "regtest",
                "duration_seconds": round(time.monotonic() - started, 2),
                "isolation": {
                    "networkactive": network_info.get("networkactive"),
                    "connections": int(network_info.get("connections", -1)),
                    "listen": False,
                    "dns": False,
                },
                "receipt": {
                    "test_amount_btc": f"{TEST_AMOUNT:.8f}",
                    "confirmed_balance_btc": f"{receipt_balance:.8f}",
                },
                "signing": {
                    "watch_only_balance_btc": f"{watch_balance:.8f}",
                    "private_keys_enabled": watch_info.get("private_keys_enabled"),
                    "psbt_complete": bool(processed_psbt.get("complete")),
                    "direct_spend_error_code": direct_spend_error_code,
                    "listunspent_spendable_observation": bool(
                        watch_utxos[0].get("spendable")
                    ),
                },
                "recovery": {
                    "backup_taken_before_receive": True,
                    "restored_balance_btc": f"{restored_balance:.8f}",
                    "restored_spend_btc": f"{RESTORED_SPEND:.8f}",
                    "restored_spend_confirmations": restored_confirmations,
                },
                "assertions": {
                    "passed": passed,
                    "total": len(checks),
                    "all_passed": passed == len(checks),
                    "checks": checks,
                },
                "sanitization": {
                    "real_funds_used": False,
                    "addresses_included": False,
                    "transaction_ids_included": False,
                    "descriptors_or_keys_included": False,
                    "wallet_backups_included": False,
                    "rpc_credentials_included": False,
                },
            }
        finally:
            try:
                cli("stop")
            except Exception:
                node.terminate()
            try:
                node.wait(timeout=30)
            except subprocess.TimeoutExpired:
                node.kill()
                node.wait(timeout=10)


def main() -> int:
    args = parse_args()
    try:
        result = run_lab(args.bin_dir.resolve())
    except LabError as error:
        result = {
            "schema_version": 1,
            "status": "ERROR",
            "stage": error.stage,
            "error": str(error),
            "rpc_code": error.rpc_code,
            "sanitized": True,
        }
        print(json.dumps(result, indent=2, sort_keys=True))
        return 1
    except Exception as error:
        result = {
            "schema_version": 1,
            "status": "ERROR",
            "stage": "unhandled",
            "error_type": type(error).__name__,
            "sanitized": True,
        }
        print(json.dumps(result, indent=2, sort_keys=True))
        return 1

    print(json.dumps(result, indent=2, sort_keys=True))
    return 0 if result["status"] == "PASS" else 1


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