#!/usr/bin/env python3
"""
Export 33 x 33 x 6 = 6534 Chain 138 soak wallet addresses from a BIP39 mnemonic.

Uses `cast wallet address` (Foundry) per index — parallel workers for speed.
Never writes the mnemonic; read from env PMM_SOAK_GRID_MNEMONIC only.

  PMM_SOAK_GRID_MNEMONIC='...' python3 scripts/deployment/pmm-soak-export-wallet-grid.py --out config/pmm-soak-wallet-grid.json
  # Smoke / CI (first N indices only):
  PMM_SOAK_GRID_MNEMONIC='...' python3 scripts/deployment/pmm-soak-export-wallet-grid.py --out /tmp/grid-smoke.json --count 10
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed

LPBCA_COUNT = 33
BRANCH_COUNT = 33
CLASS_COUNT = 6
LINEAR_COUNT = LPBCA_COUNT * BRANCH_COUNT * CLASS_COUNT  # 6534


def linear_to_coords(linear: int) -> tuple[int, int, int]:
    lpbca = linear // (BRANCH_COUNT * CLASS_COUNT)
    rem = linear % (BRANCH_COUNT * CLASS_COUNT)
    branch = rem // CLASS_COUNT
    klass = rem % CLASS_COUNT
    return lpbca, branch, klass


def derive_address(linear_index: int) -> tuple[int, str]:
    mn = os.environ.get("PMM_SOAK_GRID_MNEMONIC", "").strip()
    if not mn:
        raise RuntimeError("PMM_SOAK_GRID_MNEMONIC is not set")
    path = f"m/44'/60'/0'/0/{linear_index}"
    r = subprocess.run(
        [
            "cast",
            "wallet",
            "address",
            "--mnemonic",
            mn,
            "--mnemonic-derivation-path",
            path,
        ],
        capture_output=True,
        text=True,
        check=False,
    )
    if r.returncode != 0:
        raise RuntimeError(f"cast failed for index {linear_index}: {r.stderr.strip()}")
    addr = r.stdout.strip()
    if not addr.startswith("0x") or len(addr) != 42:
        raise RuntimeError(f"bad address for index {linear_index}: {addr!r}")
    return linear_index, addr


def main() -> int:
    ap = argparse.ArgumentParser(description="Export PMM soak 33x33x6 wallet grid addresses.")
    ap.add_argument(
        "--out",
        required=True,
        help="Output JSON path (use a gitignored file for real grids).",
    )
    ap.add_argument("--jobs", type=int, default=8, help="Parallel cast workers (default 8).")
    ap.add_argument(
        "--count",
        type=int,
        default=None,
        metavar="N",
        help=f"Export only linear indices 0..N-1 (default: full grid {LINEAR_COUNT}). For smoke tests.",
    )
    args = ap.parse_args()

    if not os.environ.get("PMM_SOAK_GRID_MNEMONIC", "").strip():
        print("ERROR: set PMM_SOAK_GRID_MNEMONIC in the environment.", file=sys.stderr)
        return 1

    n = LINEAR_COUNT if args.count is None else int(args.count)
    if n < 1 or n > LINEAR_COUNT:
        print(f"ERROR: --count must be 1..{LINEAR_COUNT}", file=sys.stderr)
        return 1

    addrs: list[str | None] = [None] * n
    jobs = max(1, min(args.jobs, 32))

    with ProcessPoolExecutor(max_workers=jobs) as ex:
        futs = {ex.submit(derive_address, i): i for i in range(n)}
        for fut in as_completed(futs):
            i, addr = fut.result()
            addrs[i] = addr

    wallets = []
    for li in range(n):
        lpbca, branch, klass = linear_to_coords(li)
        wallets.append(
            {
                "lpbca": lpbca,
                "branch": branch,
                "class": klass,
                "linearIndex": li,
                "address": addrs[li],
            }
        )

    doc: dict = {
        "version": 1,
        "dimensions": {
            "lpbcaCount": LPBCA_COUNT,
            "branchCount": BRANCH_COUNT,
            "classCount": CLASS_COUNT,
            "linearCount": LINEAR_COUNT,
            "linearCountExported": n,
        },
        "derivation": {
            "pathTemplate": "m/44'/60'/0'/0/{linearIndex}",
            "linearIndexFormula": "linear = lpbca * 198 + branch * 6 + class",
        },
        "wallets": wallets,
    }
    if n < LINEAR_COUNT:
        doc["partialExport"] = True

    out_path = args.out
    out_dir = os.path.dirname(os.path.abspath(out_path))
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(doc, f, indent=0)
        f.write("\n")

    print(f"Wrote {n} wallet(s) to {out_path}", file=sys.stderr)
    return 0


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