#!/usr/bin/env python3
"""
Apply Elemental Imperium matrix labels to config/pmm-soak-wallet-grid.json.

Each wallet gains:
  cellId      — immutable coordinate label (EI-L##-B##-C#)
  networkCode — from registry default, overlay rangeRules, then cellOverrides
  asn         — same layering; null until known
  label       — {networkCode}.{cellId}[.AS{n}]

Run from repo root:
  python3 scripts/lib/apply_elemental_imperium_wallet_labels.py

Options:
  --network-code CODE   Base networkCode for all wallets (before overlay).
  --asn N               Base ASN for all wallets (before overlay).
  --overlay PATH        JSON overlay (see config/elemental-imperium-wallet-overlay.example.json)
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

_REPO_LIB = Path(__file__).resolve().parent
if str(_REPO_LIB) not in sys.path:
    sys.path.insert(0, str(_REPO_LIB))

from elemental_imperium_wallet_common import (  # noqa: E402
    build_label,
    cell_id,
    resolve_network_asn,
)


def load_registry(repo_root: Path) -> tuple[str, dict]:
    reg_path = repo_root / "config" / "elemental-imperium-network-registry.json"
    reg = json.loads(reg_path.read_text(encoding="utf-8"))
    default = reg.get("defaultNetworkCode", "unassigned")
    return default, reg


def load_overlay(path: Path) -> dict:
    o = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(o, dict):
        raise SystemExit("Overlay must be a JSON object")
    return o


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--network-code",
        dest="network_code",
        default=None,
        help="Base networkCode for every wallet (before overlay)",
    )
    parser.add_argument(
        "--asn",
        type=int,
        default=None,
        help="Base ASN for every wallet (before overlay)",
    )
    parser.add_argument(
        "--overlay",
        type=Path,
        default=None,
        help="Path to overlay JSON (rangeRules + cellOverrides)",
    )
    args = parser.parse_args()

    repo_root = Path(__file__).resolve().parents[2]
    grid_path = repo_root / "config" / "pmm-soak-wallet-grid.json"
    default_net, registry = load_registry(repo_root)
    base_network = args.network_code or default_net
    base_asn = args.asn
    overlay: dict | None = None
    overlay_path_str: str | None = None
    if args.overlay:
        overlay = load_overlay(args.overlay.resolve())
        overlay_path_str = str(args.overlay)

    data = json.loads(grid_path.read_text(encoding="utf-8"))
    wallets = data.get("wallets")
    if not isinstance(wallets, list):
        raise SystemExit("Invalid grid: missing wallets array")

    spec = registry.get("labelSpec", {})
    labeling = {
        "specId": spec.get("id", "ei-wallet-label-v1"),
        "registryPath": "config/elemental-imperium-network-registry.json",
        "defaultNetworkCode": default_net,
        "baseNetworkCode": base_network,
        "baseAsn": base_asn,
        "overlayPath": overlay_path_str,
        "cellIdFormat": spec.get("cellIdFormat", "EI-L{lpbca:02d}-B{branch:02d}-C{class}"),
        "labelFormats": {
            "default": spec.get("labelFormatDefault", "{networkCode}.{cellId}"),
            "withAsn": spec.get("labelFormatWithAsn", "{networkCode}.{cellId}.AS{asn}"),
        },
    }

    for w in wallets:
        lpbca = int(w["lpbca"])
        branch = int(w["branch"])
        class_ = int(w["class"])
        cid = cell_id(lpbca, branch, class_)
        w["cellId"] = cid
        net, asn = resolve_network_asn(w, cid, base_network, base_asn, overlay)
        w["networkCode"] = net
        w["asn"] = asn
        w["label"] = build_label(net, cid, asn)

    ver = max(int(data.get("version", 1)), 2)

    out = {
        "version": ver,
        "dimensions": data["dimensions"],
        "derivation": data["derivation"],
        "labeling": labeling,
        "wallets": wallets,
    }

    grid_path.write_text(json.dumps(out, indent=2, sort_keys=False) + "\n", encoding="utf-8")
    print(f"Wrote {len(wallets)} wallets to {grid_path}")


if __name__ == "__main__":
    main()