#!/usr/bin/env python3
"""Measure BTCRecover candidate generation with synthetic passwords only.

This lab never loads a wallet, seed, key, hash, or real password. It invokes
BTCRecover's ``--listpass`` mode against fixed fictional inputs and records
candidate counts plus whether the fictional target stayed inside the search.

The ``--dsw`` flag is used only to keep this synthetic, no-wallet lab
non-interactive. Do not copy it into a real recovery command.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path


BASE_PASSWORD = "CobaltOrbit2024!"


@dataclass(frozen=True)
class Scenario:
    scenario_id: str
    list_type: str
    contents: str
    target: str
    extra_args: tuple[str, ...] = ()
    description: str = ""


SCENARIOS = (
    Scenario(
        "fragments_unstructured",
        "tokenlist",
        "Cobalt\ncobalt\nOrbit\norbit\n2024\n2025\n!\n?\n",
        BASE_PASSWORD,
        ("--max-tokens", "4"),
        "Eight separate fragments; related variants can be combined as nonsense.",
    ),
    Scenario(
        "fragments_grouped",
        "tokenlist",
        "Cobalt cobalt\nOrbit orbit\n2024 2025\n! ?\n",
        BASE_PASSWORD,
        ("--max-tokens", "4"),
        "Same-family alternatives share a line and are mutually exclusive.",
    ),
    Scenario(
        "fragments_required",
        "tokenlist",
        "+ Cobalt cobalt\n+ Orbit orbit\n+ 2024 2025\n+ ! ?\n",
        BASE_PASSWORD,
        ("--max-tokens", "4"),
        "All four remembered fragment families are required, but order is open.",
    ),
    Scenario(
        "fragments_required_and_positioned",
        "tokenlist",
        "+ ^Cobalt ^cobalt\n+ ^2^Orbit ^2^orbit\n+ ^3^2024 ^3^2025\n+ ^4^! ^4^?\n",
        BASE_PASSWORD,
        ("--max-tokens", "4"),
        "All fragment families are required and fixed to remembered positions.",
    ),
    Scenario(
        "whole_guess_exact",
        "passwordlist",
        BASE_PASSWORD + "\n",
        BASE_PASSWORD,
        description="One complete remembered guess, tried verbatim.",
    ),
    Scenario(
        "one_case_change",
        "passwordlist",
        BASE_PASSWORD + "\n",
        "cobaltOrbit2024!",
        ("--typos", "1", "--typos-case"),
        "One capitalization uncertainty, searched with the matching typo class.",
    ),
    Scenario(
        "one_adjacent_swap",
        "passwordlist",
        BASE_PASSWORD + "\n",
        "CoabltOrbit2024!",
        ("--typos", "1", "--typos-swap"),
        "One adjacent transposition, searched with the matching typo class.",
    ),
    Scenario(
        "one_deleted_character",
        "passwordlist",
        BASE_PASSWORD + "\n",
        "CobaltOrbt2024!",
        ("--typos", "1", "--typos-delete"),
        "One missing character, searched with the matching typo class.",
    ),
    Scenario(
        "one_repeated_character",
        "passwordlist",
        BASE_PASSWORD + "\n",
        "CobalttOrbit2024!",
        ("--typos", "1", "--typos-repeat"),
        "One doubled character, searched with the matching typo class.",
    ),
    Scenario(
        "five_typo_classes_one_change",
        "passwordlist",
        BASE_PASSWORD + "\n",
        "cobaltOrbit2024!",
        (
            "--typos",
            "1",
            "--typos-capslock",
            "--typos-swap",
            "--typos-repeat",
            "--typos-delete",
            "--typos-case",
        ),
        "Five typo classes enabled together for one possible change.",
    ),
    Scenario(
        "five_typo_classes_two_changes",
        "passwordlist",
        BASE_PASSWORD + "\n",
        "coabltOrbit2024!",
        (
            "--typos",
            "2",
            "--typos-capslock",
            "--typos-swap",
            "--typos-repeat",
            "--typos-delete",
            "--typos-case",
        ),
        "The same five classes widened from one to two possible changes.",
    ),
)


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def repo_commit(repo_root: Path) -> str:
    completed = subprocess.run(
        ["git", "-C", str(repo_root), "rev-parse", "HEAD"],
        check=True,
        capture_output=True,
        text=True,
        timeout=30,
    )
    return completed.stdout.strip()


def git_value(repo_root: Path, *args: str) -> str:
    completed = subprocess.run(
        ["git", "-C", str(repo_root), *args],
        check=True,
        capture_output=True,
        text=True,
        timeout=30,
    )
    return completed.stdout.strip()


def tool_version(btcrecover_path: Path) -> str:
    completed = subprocess.run(
        [sys.executable, str(btcrecover_path), "--version"],
        check=True,
        capture_output=True,
        text=True,
        timeout=30,
    )
    combined = "\n".join((completed.stdout, completed.stderr))
    match = re.search(r"btcrecover\s+([^\s]+)", combined, re.IGNORECASE)
    return match.group(1) if match else combined.strip()


def run_scenario(btcrecover_path: Path, scenario: Scenario, work_dir: Path) -> dict:
    list_path = work_dir / f"{scenario.scenario_id}.txt"
    list_path.write_text(scenario.contents, encoding="utf-8", newline="\n")
    flag = "--tokenlist" if scenario.list_type == "tokenlist" else "--passwordlist"
    command = [
        sys.executable,
        str(btcrecover_path),
        "--listpass",
        flag,
        str(list_path),
        "--dsw",
        *scenario.extra_args,
    ]
    completed = subprocess.run(
        command,
        check=True,
        capture_output=True,
        text=True,
        encoding="utf-8",
        timeout=120,
    )
    summary_match = re.search(r"(\d+) password combinations", completed.stderr)
    reported_count = int(summary_match.group(1)) if summary_match else None
    stdout_lines = [
        line
        for line in completed.stdout.splitlines()
        if line
        and not line.startswith("Duplicate Check Level:")
        and not line.startswith("Notice: Loading File:")
        and not line.startswith("Notice: Finished File:")
    ]
    if reported_count is None or len(stdout_lines) != reported_count:
        raise RuntimeError(
            f"{scenario.scenario_id}: BTCRecover reported {reported_count}, "
            f"but filtered stdout contained {len(stdout_lines)} candidate lines"
        )
    # BTCRecover may print duplicate-check and passwordlist-loading notices to
    # stdout around --listpass candidates. Remove those known notices, then
    # independently check the candidate count and uniqueness.
    candidates = stdout_lines
    unique_candidates = set(candidates)
    return {
        "scenario_id": scenario.scenario_id,
        "description": scenario.description,
        "list_type": scenario.list_type,
        "candidate_count": len(candidates),
        "unique_candidate_count": len(unique_candidates),
        "target_present": scenario.target in unique_candidates,
        "input_sha256": hashlib.sha256(scenario.contents.encode("utf-8")).hexdigest(),
        "arguments": ["--listpass", flag, "SYNTHETIC_INPUT", "--dsw", *scenario.extra_args],
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--btcrecover-root",
        required=True,
        type=Path,
        help="Path to a checked-out 3rdIteration/btcrecover repository.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path(__file__).with_name("candidate-space-results.json"),
        help="Result JSON path.",
    )
    args = parser.parse_args()

    repo_root = args.btcrecover_root.resolve()
    btcrecover_path = repo_root / "btcrecover.py"
    if not btcrecover_path.is_file():
        parser.error(f"btcrecover.py not found under {repo_root}")

    with tempfile.TemporaryDirectory(prefix="cora-a02-synthetic-") as temp_name:
        temp_dir = Path(temp_name)
        runs = [run_scenario(btcrecover_path, scenario, temp_dir) for scenario in SCENARIOS]

    by_id = {run["scenario_id"]: run for run in runs}
    unstructured = by_id["fragments_unstructured"]["candidate_count"]
    positioned = by_id["fragments_required_and_positioned"]["candidate_count"]
    one_change = by_id["five_typo_classes_one_change"]["candidate_count"]
    two_changes = by_id["five_typo_classes_two_changes"]["candidate_count"]
    results = {
        "study": "Cora Aegis A-02 synthetic BTCRecover candidate-space lab",
        "scope": "Candidate generation only; no wallet, seed, key, hash, funds, or real password.",
        "python_version": sys.version.split()[0],
        "btcrecover_version": tool_version(btcrecover_path),
        "btcrecover_commit": repo_commit(repo_root),
        "btcrecover_tree": git_value(repo_root, "rev-parse", "HEAD^{tree}"),
        "btcrecover_origin": git_value(repo_root, "remote", "get-url", "origin"),
        "btcrecover_worktree_dirty": bool(
            git_value(repo_root, "status", "--porcelain", "--untracked-files=no")
        ),
        "btcrecover_py_sha256": sha256_file(btcrecover_path),
        "lab_script_sha256": sha256_file(Path(__file__).resolve()),
        "scenario_count": len(runs),
        "all_targets_present": all(run["target_present"] for run in runs),
        "all_candidates_unique_within_each_run": all(
            run["candidate_count"] == run["unique_candidate_count"] for run in runs
        ),
        "derived": {
            "structured_reduction_percent": round(
                (1 - positioned / unstructured) * 100, 4
            ),
            "unstructured_to_positioned_ratio": round(unstructured / positioned, 2),
            "two_typo_to_one_typo_ratio": round(two_changes / one_change, 2),
        },
        "runs": runs,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(results, indent=2, sort_keys=True) + "\n", encoding="utf-8", newline="\n"
    )
    print(json.dumps(results, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())