#!/usr/bin/env python3
"""
On-chain readiness audit for EI matrix wallets (config/pmm-soak-wallet-grid.json).

Queries ERC-20 balanceOf for each address on one or both chains:
  - Ethereum mainnet cWUSDC (default from env CWUSDC_MAINNET)
  - Chain 138 cUSDC (default canonical CompliantUSDC)

Use for strength profiling: segment by class/lpbca via --report-by-class, find gaps vs thresholds.

Environment (optional defaults for thresholds):
  EI_MATRIX_AUDIT_MIN_MAINNET_RAW, EI_MATRIX_AUDIT_MIN_138_RAW, EI_MATRIX_AUDIT_WORKERS

Examples:
  python3 scripts/lib/ei_matrix_onchain_readiness_audit.py --mainnet-only --min-mainnet-raw 1
  python3 scripts/lib/ei_matrix_onchain_readiness_audit.py --both \\
    --shard-size 400 --min-mainnet-raw 12000000 --min-138-raw 0 --workers 3 \\
    --report-by-class --json-out reports/status/ei-matrix-readiness-audit-latest.json
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

# balanceOf(address) selector
BALANCE_OF = bytes.fromhex("70a08231")
ADDR_PAD = 12 * b"\x00"


def encode_balance_of_call(addr: str) -> str:
    a = addr.lower().removeprefix("0x")
    if len(a) != 40:
        raise ValueError(f"bad address {addr}")
    data = BALANCE_OF + ADDR_PAD + bytes.fromhex(a)
    return "0x" + data.hex()


def rpc_eth_call(to: str, data: str, rpc_url: str, timeout: float = 30.0) -> str:
    body = json.dumps(
        {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "eth_call",
            "params": [{"to": to, "data": data}, "latest"],
        }
    ).encode()
    req = urllib.request.Request(rpc_url, data=body, headers={"Content-Type": "application/json"}, method="POST")
    with urllib.request.urlopen(req, timeout=timeout) as r:
        j = json.loads(r.read().decode())
    if "error" in j:
        raise RuntimeError(str(j["error"]))
    return j.get("result") or "0x0"


def hex_to_int(h: str) -> int:
    h = h.strip()
    if not h or h == "0x":
        return 0
    return int(h, 16)


def collect_rows_for_slice(
    slice_items: list[tuple[int, dict]],
    *,
    do_main: bool,
    do_138: bool,
    mainnet_rpc: str,
    chain138_rpc: str,
    mainnet_token: str,
    chain138_cusdc: str,
    workers: int,
) -> list[dict]:
    def fetch_one(item: tuple[int, dict]) -> tuple[int, dict, int, int]:
        idx, w = item
        addr = w["address"]
        mbal, bbal = 0, 0
        if do_main:
            calldata = encode_balance_of_call(addr)
            res = rpc_eth_call(mainnet_token.lower(), calldata, mainnet_rpc)
            mbal = hex_to_int(res)
        if do_138:
            calldata = encode_balance_of_call(addr)
            res = rpc_eth_call(chain138_cusdc.lower(), calldata, chain138_rpc)
            bbal = hex_to_int(res)
        return idx, w, mbal, bbal

    rows: list[dict] = []
    with ThreadPoolExecutor(max_workers=max(1, workers)) as ex:
        futs = [ex.submit(fetch_one, it) for it in slice_items]
        for fut in as_completed(futs):
            idx, w, mbal, bbal = fut.result()
            cls = int(w.get("class", 0))
            row = {
                "linearIndex": idx,
                "address": w["address"],
                "cellId": w.get("cellId"),
                "class": cls,
                "mainnetCwusdcRaw": mbal if do_main else None,
                "chain138CusdcRaw": bbal if do_138 else None,
            }
            rows.append(row)
    return rows


def write_indices(path: Path, indices: list[int]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(str(i) for i in indices) + ("\n" if indices else ""), encoding="utf-8")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--grid", default="config/pmm-soak-wallet-grid.json")
    ap.add_argument("--offset", type=int, default=0)
    ap.add_argument("--limit", type=int, default=0, help="0 = all from offset to grid end")
    ap.add_argument(
        "--shard-size",
        type=int,
        default=int(os.environ.get("EI_MATRIX_AUDIT_SHARD_SIZE", "0")),
        help="If >0, query in sequential shards of this size (eases RPC load). 0 = single batch.",
    )
    ap.add_argument("--workers", type=int, default=int(os.environ.get("EI_MATRIX_AUDIT_WORKERS", "4")))
    ap.add_argument("--mainnet-only", action="store_true")
    ap.add_argument("--chain138-only", action="store_true")
    ap.add_argument("--both", action="store_true")
    ap.add_argument("--mainnet-rpc", default=os.environ.get("ETHEREUM_MAINNET_RPC") or os.environ.get("RPC_URL_1") or "")
    ap.add_argument("--chain138-rpc", default=os.environ.get("RPC_URL_138") or os.environ.get("CHAIN138_PUBLIC_RPC_URL") or "")
    ap.add_argument("--mainnet-token", default=os.environ.get("CWUSDC_MAINNET", "0x2de5F116bFcE3d0f922d9C8351e0c5Fc24b9284a"))
    ap.add_argument(
        "--chain138-cusdc",
        default=os.environ.get("CUSDC_CHAIN138", "0xf22258f57794CC8E06237084b353Ab30fFfa640b"),
    )
    ap.add_argument(
        "--min-mainnet-raw",
        type=int,
        default=int(os.environ.get("EI_MATRIX_AUDIT_MIN_MAINNET_RAW", "0")),
        help="fail wallets strictly below this (mainnet); env EI_MATRIX_AUDIT_MIN_MAINNET_RAW",
    )
    ap.add_argument(
        "--min-138-raw",
        type=int,
        default=int(os.environ.get("EI_MATRIX_AUDIT_MIN_138_RAW", "0")),
        help="fail wallets strictly below this (138); env EI_MATRIX_AUDIT_MIN_138_RAW",
    )
    ap.add_argument("--report-by-class", action="store_true", help="aggregate counts by matrix class 0..5")
    ap.add_argument("--json-out", default="", help="write full per-wallet rows + summary")
    ap.add_argument(
        "--gaps-mainnet-out",
        default="",
        help="write newline-separated linear indices below mainnet minimum (only if mainnet queried)",
    )
    ap.add_argument(
        "--gaps-138-out",
        default="",
        help="write newline-separated linear indices below 138 minimum (only if 138 queried)",
    )
    ap.add_argument("--max-list", type=int, default=200, help="max gap indices to print on stderr")
    args = ap.parse_args()

    repo = Path(__file__).resolve().parents[2]
    grid_path = repo / args.grid if not os.path.isabs(args.grid) else Path(args.grid)
    data = json.loads(grid_path.read_text(encoding="utf-8"))
    wallets: list[dict] = data["wallets"]
    n = len(wallets)
    scan_end = n if args.limit <= 0 else min(n, args.offset + args.limit)
    scan_start = args.offset
    if scan_start < 0 or scan_start > n:
        print("Invalid --offset", file=sys.stderr)
        return 2
    if scan_end < scan_start:
        print("Invalid --limit / range", file=sys.stderr)
        return 2

    do_main = args.mainnet_only or args.both
    do_138 = args.chain138_only or args.both
    if not do_main and not do_138:
        print("Specify --mainnet-only, --chain138-only, or --both", file=sys.stderr)
        return 2
    if do_main and not args.mainnet_rpc:
        print("Need --mainnet-rpc or ETHEREUM_MAINNET_RPC / RPC_URL_1", file=sys.stderr)
        return 2
    if do_138 and not args.chain138_rpc:
        print("Need --chain138-rpc or RPC_URL_138", file=sys.stderr)
        return 2

    shard = max(0, args.shard_size)
    rows: list[dict] = []
    if shard <= 0:
        slice_items = list(enumerate(wallets[scan_start:scan_end], start=scan_start))
        rows = collect_rows_for_slice(
            slice_items,
            do_main=do_main,
            do_138=do_138,
            mainnet_rpc=args.mainnet_rpc,
            chain138_rpc=args.chain138_rpc,
            mainnet_token=args.mainnet_token,
            chain138_cusdc=args.chain138_cusdc,
            workers=args.workers,
        )
    else:
        for start in range(scan_start, scan_end, shard):
            chunk_end = min(scan_end, start + shard)
            slice_items = list(enumerate(wallets[start:chunk_end], start=start))
            print(f"Shard {start}..{chunk_end} ({len(slice_items)} wallets)", file=sys.stderr)
            rows.extend(
                collect_rows_for_slice(
                    slice_items,
                    do_main=do_main,
                    do_138=do_138,
                    mainnet_rpc=args.mainnet_rpc,
                    chain138_rpc=args.chain138_rpc,
                    mainnet_token=args.mainnet_token,
                    chain138_cusdc=args.chain138_cusdc,
                    workers=args.workers,
                )
            )

    rows.sort(key=lambda r: r["linearIndex"])

    by_class: dict[int, dict] = {i: {"n": 0, "mainnet_below": 0, "138_below": 0} for i in range(6)}
    if args.report_by_class:
        for r in rows:
            cls = int(r.get("class", 0))
            if cls not in by_class:
                continue
            by_class[cls]["n"] += 1
            if do_main and r["mainnetCwusdcRaw"] < args.min_mainnet_raw:
                by_class[cls]["mainnet_below"] += 1
            if do_138 and r["chain138CusdcRaw"] < args.min_138_raw:
                by_class[cls]["138_below"] += 1

    gaps_main: list[int] = []
    gaps_138: list[int] = []
    for r in rows:
        if do_main and r["mainnetCwusdcRaw"] < args.min_mainnet_raw:
            gaps_main.append(r["linearIndex"])
        if do_138 and r["chain138CusdcRaw"] < args.min_138_raw:
            gaps_138.append(r["linearIndex"])

    summary = {
        "gridPath": str(grid_path),
        "slice": {"offset": scan_start, "endExclusive": scan_end, "count": len(rows)},
        "shardSize": shard if shard > 0 else None,
        "mainnet": {
            "token": args.mainnet_token if do_main else None,
            "rpc": args.mainnet_rpc[:48] + "…" if do_main and len(args.mainnet_rpc) > 48 else args.mainnet_rpc,
            "minRaw": args.min_mainnet_raw,
            "belowMin": len(gaps_main),
        },
        "chain138": {
            "token": args.chain138_cusdc if do_138 else None,
            "minRaw": args.min_138_raw,
            "belowMin": len(gaps_138),
        },
        "byClass": by_class if args.report_by_class else None,
    }

    print(json.dumps(summary, indent=2))
    if gaps_main:
        print(
            f"\nMainnet cWUSDC below min ({args.min_mainnet_raw}) — {len(gaps_main)} wallets "
            f"(first {args.max_list} indices):",
            file=sys.stderr,
        )
        print(", ".join(str(x) for x in gaps_main[: args.max_list]), file=sys.stderr)
    if gaps_138:
        print(
            f"\nChain 138 cUSDC below min ({args.min_138_raw}) — {len(gaps_138)} wallets "
            f"(first {args.max_list} indices):",
            file=sys.stderr,
        )
        print(", ".join(str(x) for x in gaps_138[: args.max_list]), file=sys.stderr)

    if args.json_out:
        outp = repo / args.json_out if not os.path.isabs(args.json_out) else Path(args.json_out)
        outp.parent.mkdir(parents=True, exist_ok=True)
        outp.write_text(json.dumps({"summary": summary, "rows": rows}, indent=2), encoding="utf-8")
        print(f"\nWrote {outp}", file=sys.stderr)

    if do_main and args.gaps_mainnet_out:
        gp = repo / args.gaps_mainnet_out if not os.path.isabs(args.gaps_mainnet_out) else Path(args.gaps_mainnet_out)
        write_indices(gp, gaps_main)
        print(f"Wrote mainnet gap indices ({len(gaps_main)}): {gp}", file=sys.stderr)
    if do_138 and args.gaps_138_out:
        gp = repo / args.gaps_138_out if not os.path.isabs(args.gaps_138_out) else Path(args.gaps_138_out)
        write_indices(gp, gaps_138)
        print(f"Wrote 138 gap indices ({len(gaps_138)}): {gp}", file=sys.stderr)

    fail = bool(gaps_main or gaps_138)
    return 1 if fail else 0


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