#!/usr/bin/env python3
# =============================================================================
#  Project   : localvm-research
#  File      : benchmarks/hardware_manifest.py
#  Purpose   : macOS hardware/software fingerprint embedded in every result file
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Created   : 2026-08-11
#  Modified  : 2026-08-11
#  Platform  : macOS / Apple Silicon (arm64)
#  License   : All rights reserved (research code)
# =============================================================================
"""Collect a reproducibility manifest for the current Mac.

Records chip model, P/E core counts, GPU core count, unified memory size,
SSD model, macOS version, and versions of the key software stack (Python,
MLX, PyTorch, NumPy), plus the current git commit. Every benchmark result
JSON must embed this manifest (CLAUDE.md §0.2, §10).

Usage:
    python3 benchmarks/hardware_manifest.py            # pretty-print JSON
    from hardware_manifest import collect_manifest     # programmatic use
"""

from __future__ import annotations

import json
import platform
import subprocess
import sys
from datetime import datetime, timezone


def _run(cmd: list[str]) -> str:
    try:
        return subprocess.run(cmd, capture_output=True, text=True, timeout=30).stdout.strip()
    except (OSError, subprocess.TimeoutExpired):
        return ""


def _sysctl(key: str) -> str:
    return _run(["sysctl", "-n", key])


def _sysctl_int(key: str) -> int | None:
    val = _sysctl(key)
    try:
        return int(val)
    except ValueError:
        return None


def _gpu_cores() -> int | None:
    """GPU core count via system_profiler (no sysctl key exposes it)."""
    out = _run(["system_profiler", "SPDisplaysDataType", "-json"])
    try:
        displays = json.loads(out)["SPDisplaysDataType"]
        for gpu in displays:
            cores = gpu.get("sppci_cores")
            if cores is not None:
                return int(cores)
    except (json.JSONDecodeError, KeyError, ValueError, TypeError):
        pass
    return None


def _ssd_info() -> dict:
    out = _run(["system_profiler", "SPNVMeDataType", "-json"])
    try:
        items = json.loads(out)["SPNVMeDataType"]
        for controller in items:
            for dev in controller.get("_items", []):
                return {
                    "model": dev.get("device_model", "").strip(),
                    "size": dev.get("size", ""),
                    "smart_status": dev.get("smart_status", ""),
                }
    except (json.JSONDecodeError, KeyError, TypeError):
        pass
    return {"model": None, "size": None, "smart_status": None}


def _pkg_version(module: str) -> str | None:
    try:
        from importlib.metadata import version

        return version(module)
    except Exception:
        return None


def _git_commit() -> dict:
    commit = _run(["git", "rev-parse", "HEAD"])
    dirty = bool(_run(["git", "status", "--porcelain"]))
    return {"commit": commit or None, "dirty_tree": dirty}


def _thermal_state() -> str | None:
    # 'thermal pressure' via thermal level sysctl where available
    lvl = _sysctl("machdep.xcpm.cpu_thermal_level")
    return lvl or None


def collect_manifest() -> dict:
    """Return the full hardware/software manifest as a dict."""
    mem_bytes = _sysctl_int("hw.memsize") or 0
    manifest = {
        "author": "Simon-Pierre Boucher",
        "contact": "contact@spboucher.ai",
        "project": "localvm-research",
        "collected_utc": datetime.now(timezone.utc).isoformat(),
        "chip": {
            "brand": _sysctl("machdep.cpu.brand_string"),
            "arch": platform.machine(),
            "cores_total": _sysctl_int("hw.ncpu"),
            "cores_performance": _sysctl_int("hw.perflevel0.physicalcpu"),
            "cores_efficiency": _sysctl_int("hw.perflevel1.physicalcpu"),
            "gpu_cores": _gpu_cores(),
        },
        "memory": {
            "unified_bytes": mem_bytes,
            "unified_gb": round(mem_bytes / 2**30, 1),
            "pagesize": _sysctl_int("hw.pagesize"),
        },
        "ssd": _ssd_info(),
        "os": {
            "product": _run(["sw_vers", "-productName"]),
            "version": _run(["sw_vers", "-productVersion"]),
            "build": _run(["sw_vers", "-buildVersion"]),
            "kernel": platform.release(),
        },
        "software": {
            "python": sys.version.split()[0],
            "mlx": _pkg_version("mlx"),
            "mlx_lm": _pkg_version("mlx-lm"),
            "torch": _pkg_version("torch"),
            "numpy": _pkg_version("numpy"),
        },
        "git": _git_commit(),
        "thermal_level_at_collect": _thermal_state(),
    }
    return manifest


if __name__ == "__main__":
    print(json.dumps(collect_manifest(), indent=2))