#!/usr/bin/env python
"""Phase 0 environment audit: records hardware, toolchain and framework capabilities.

Writes hardware/device_inventory.json (local machine) and prints a summary.
Never prints secrets. Remote cluster nodes are audited by scripts/cluster_audit.sh
(SSH, read-only commands) and merged into the same JSON under "remote_nodes".
"""
from __future__ import annotations
import argparse, json, os, platform, shutil, subprocess, sys, time
from pathlib import Path


def sh(cmd: list[str]) -> str:
    try:
        return subprocess.run(cmd, capture_output=True, text=True, timeout=20).stdout.strip()
    except Exception as e:  # noqa: BLE001
        return f"<unavailable: {type(e).__name__}>"


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--out", default="hardware/device_inventory.json")
    ap.add_argument("--smoke", action="store_true", help="run a tiny forward/backward on every available backend")
    args = ap.parse_args()

    inv: dict = {
        "audited_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "hostname": platform.node(),
        "os": f"{platform.system()} {platform.release()} ({platform.mac_ver()[0] or platform.version()})",
        "machine": platform.machine(),
        "python": sys.version.split()[0],
        "python_executable": sys.executable,
        "cpu_brand": sh(["sysctl", "-n", "machdep.cpu.brand_string"]) if sys.platform == "darwin" else platform.processor(),
        "cpu_logical_cores": os.cpu_count(),
        "ram_bytes": int(sh(["sysctl", "-n", "hw.memsize"]) or 0) if sys.platform == "darwin" else None,
        "disk_free_bytes": shutil.disk_usage(os.getcwd()).free,
        "toolchains": {t: shutil.which(t) for t in ["git", "latexmk", "pdflatex", "lualatex", "pandoc", "nvidia-smi", "uv", "cmake", "clang"]},
        "frameworks": {},
        "backends": {},
    }
    if sys.platform == "darwin":
        sp = sh(["system_profiler", "SPDisplaysDataType"])
        gpu_cores = None
        for line in sp.splitlines():
            if "Total Number of Cores" in line:
                gpu_cores = int(line.split(":")[-1].strip())
        inv["gpu_cores"] = gpu_cores
        inv["metal"] = any("Metal" in l for l in sp.splitlines())
    try:
        import torch
        inv["frameworks"]["torch"] = torch.__version__
        inv["backends"]["cpu"] = True
        inv["backends"]["mps"] = bool(torch.backends.mps.is_available())
        inv["backends"]["cuda"] = bool(torch.cuda.is_available())
        if torch.cuda.is_available():
            inv["cuda_devices"] = [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
    except Exception as e:  # noqa: BLE001
        inv["frameworks"]["torch"] = f"<missing: {e}>"
    try:
        import mlx.core as mx
        inv["frameworks"]["mlx"] = mx.__version__
    except Exception:  # noqa: BLE001
        inv["frameworks"]["mlx"] = None
    for pkg in ["numpy", "tokenizers", "datasets", "safetensors", "scipy", "pandas", "matplotlib", "pytest", "hypothesis", "yaml"]:
        try:
            m = __import__(pkg)
            inv["frameworks"][pkg] = getattr(m, "__version__", "present")
        except Exception:  # noqa: BLE001
            inv["frameworks"][pkg] = None
    inv["network"] = {"note": "outbound HTTPS assumed available for dataset download; verified at data preparation time"}
    inv["secrets_policy"] = "no secret values are recorded in this inventory"

    if args.smoke:
        import torch
        res = {}
        for dev in [d for d, ok in inv["backends"].items() if ok]:
            try:
                t0 = time.perf_counter()
                x = torch.randn(64, 256, device=dev, requires_grad=True)
                w = torch.randn(256, 256, device=dev, requires_grad=True)
                y = (x @ w).relu().sum()
                y.backward()
                if dev == "mps":
                    torch.mps.synchronize()
                res[dev] = {"ok": bool(torch.isfinite(w.grad).all().item()), "seconds": round(time.perf_counter() - t0, 4)}
            except Exception as e:  # noqa: BLE001
                res[dev] = {"ok": False, "error": repr(e)}
        inv["smoke_forward_backward"] = res

    out = Path(args.out)
    existing = {}
    if out.exists():
        try:
            existing = json.loads(out.read_text())
        except Exception:  # noqa: BLE001
            existing = {}
    if "remote_nodes" in existing:
        inv["remote_nodes"] = existing["remote_nodes"]
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(inv, indent=2, ensure_ascii=False))
    print(json.dumps({k: v for k, v in inv.items() if k != "remote_nodes"}, indent=2, ensure_ascii=False))
    return 0


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