#!/usr/bin/env python3
"""Environment audit: hardware, Python stack, backends, cached open-weight checkpoints, cluster nodes.

Writes docs/ENVIRONMENT_AUDIT.json (and prints a short report). Cluster probing is optional
(``--cluster``) and uses the user's existing SSH aliases; nodes that do not answer are reported as
unreachable rather than assumed online. No secrets are read or written.
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from wp_ai_0001.instrumentation.provenance import hardware, library_versions, utc_now  # noqa: E402

NODES = ["M3U96a", "M3U96b", "M2U64", "M4M64a", "M4M64b", "M4M64c", "M5M36", "M4M36", "M4BP48", "m4mm", "m4mn"]
REMOTE_CMD = (
    'echo "$(sysctl -n machdep.cpu.brand_string)|$(( $(sysctl -n hw.memsize)/1073741824 ))|'
    '$(df -k ~ | tail -1 | awk \'{print int($4/1048576)}\')|$(python3 --version 2>&1 | cut -d" " -f2)|'
    '$(python3 -c "import mlx_lm;print(mlx_lm.__version__)" 2>/dev/null || echo none)|'
    '$(python3 -c "import torch;print(torch.__version__)" 2>/dev/null || echo none)|'
    '$(ls ~/.cache/huggingface/hub 2>/dev/null | grep -c models--)"'
)


def hf_cache() -> list[dict]:
    hub = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")) / "hub"
    out = []
    if hub.exists():
        for d in sorted(hub.glob("models--*")):
            size = sum(f.stat().st_size for f in d.rglob("*") if f.is_file())
            out.append({"repo": d.name[len("models--"):].replace("--", "/"), "size_gb": round(size / 2**30, 2)})
    return out


def backends() -> dict:
    b: dict = {}
    try:
        import mlx.core as mx

        b["mlx"] = {"version": mx.__version__, "metal": bool(mx.metal.is_available())}
    except Exception as e:  # noqa: BLE001
        b["mlx"] = {"error": repr(e)}
    try:
        import torch

        b["torch"] = {"version": torch.__version__, "mps": bool(torch.backends.mps.is_available()), "cuda": bool(torch.cuda.is_available())}
    except Exception as e:  # noqa: BLE001
        b["torch"] = {"error": repr(e)}
    return b


def probe_nodes() -> dict:
    res = {}
    for n in NODES:
        try:
            out = subprocess.run(["ssh", "-o", "ConnectTimeout=6", "-o", "BatchMode=yes", n, REMOTE_CMD], capture_output=True, text=True, timeout=30)
            line = out.stdout.strip().splitlines()[-1] if out.stdout.strip() else ""
            if out.returncode != 0 or "|" not in line:
                res[n] = {"reachable": False, "error": (out.stderr or out.stdout).strip()[:120]}
                continue
            chip, mem, free, py, mlxlm, torch_v, hf = line.split("|")
            res[n] = {"reachable": True, "chip": chip, "memory_gb": int(mem), "free_disk_gb": int(free), "python": py, "mlx_lm": mlxlm, "torch": torch_v, "hf_models_cached": int(hf)}
        except Exception as e:  # noqa: BLE001
            res[n] = {"reachable": False, "error": repr(e)[:120]}
    return res


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--cluster", action="store_true")
    args = ap.parse_args()
    report = {"utc": utc_now(), "hardware": hardware(), "libraries": library_versions(), "backends": backends(), "hf_cache": hf_cache()}
    if args.cluster:
        report["cluster"] = probe_nodes()
    (ROOT / "docs" / "ENVIRONMENT_AUDIT.json").write_text(json.dumps(report, indent=2, default=str))
    print(json.dumps({k: report[k] for k in ("hardware", "backends")}, indent=2, default=str))
    print("cached checkpoints:", [m["repo"] for m in report["hf_cache"]])
    if args.cluster:
        for n, r in report["cluster"].items():
            print(f"  {n}: {'OK ' + r['chip'] + ' ' + str(r['memory_gb']) + 'GB mlx_lm=' + r['mlx_lm'] if r['reachable'] else 'UNREACHABLE'}")
    print("written docs/ENVIRONMENT_AUDIT.json")


if __name__ == "__main__":
    main()