"""Record the machine and its tool versions as one JSON line in the lab notebook.

Purpose: capture the baseline every later measurement in this course is compared
         against: which machine, which operating system and kernel, which accelerator
         and how much memory it may address, which tool versions, and on what date.
Platform: all (probes that do not apply to a track are left out of the record; nothing
          here needs an accelerator to be present, and nothing is changed)
Minimum memory: 8 GB
Assumes: Python 3.9 or newer, ideally the Part 1 environment (so torch and mlx versions
         are recorded); the notebook file already exists (Part 1 created labbook.md).
         Every shell tool it asks is optional and is recorded as absent if missing.

Usage: python record-baseline.py --labbook labbook.md [--track auto] [--note "..."]
       python record-baseline.py --print     (show the record, write nothing)
"""
import argparse
import json
import os
import platform
import shutil
import subprocess
import sys
from datetime import date
from pathlib import Path

TRACKS = ("auto", "spark", "strix", "mac", "nvidia")
GIB = 1024 ** 3
TTM_PAGE_BYTES = 4096


def run(cmd, timeout=30):
    """Run a command; return its trimmed stdout, or None if absent or failing."""
    exe = cmd[0] if os.path.isabs(cmd[0]) else shutil.which(cmd[0])
    if exe is None or not os.path.exists(exe):
        return None
    try:
        out = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, check=False)
    except (OSError, subprocess.SubprocessError):
        return None
    if out.returncode != 0:
        return None
    return out.stdout.strip() or None


def first_line(text, contains=""):
    for line in (text or "").splitlines():
        if contains in line:
            return line.strip()
    return None


def read_text(path):
    try:
        return Path(path).read_text(encoding="utf-8").strip()
    except OSError:
        return None


def os_release():
    fields = {}
    for line in (read_text("/etc/os-release") or "").splitlines():
        key, _, value = line.partition("=")
        fields[key] = value.strip('"')
    return fields.get("PRETTY_NAME")


def is_wsl():
    return "microsoft" in (read_text("/proc/version") or "").lower()


def nvidia_smi():
    if shutil.which("nvidia-smi"):
        return "nvidia-smi"
    if os.path.exists("/usr/lib/wsl/lib/nvidia-smi"):
        return "/usr/lib/wsl/lib/nvidia-smi"
    return None


def guess_track():
    if sys.platform == "darwin":
        return "mac"
    if nvidia_smi():
        # The GB10 is the only aarch64 machine in this course with an NVIDIA GPU.
        return "spark" if platform.machine() == "aarch64" else "nvidia"
    if list(Path("/sys/class/drm").glob("card*/device/mem_info_gtt_total")):
        return "strix"
    return "unknown"


def memtotal_bytes():
    if sys.platform == "darwin":
        out = run(["sysctl", "-n", "hw.memsize"])
        return int(out) if out and out.isdigit() else None
    for line in (read_text("/proc/meminfo") or "").splitlines():
        if line.startswith("MemTotal:"):
            return int(line.split()[1]) * 1024
    return None


def machine():
    total = memtotal_bytes()
    info = {
        "arch": platform.machine(),
        "system_memory_gib": round(total / GIB, 2) if total else None,
        "kernel": platform.release(),
    }
    if sys.platform == "darwin":
        info["os"] = "macOS " + (run(["sw_vers", "-productVersion"]) or "unknown")
        info["chip"] = run(["sysctl", "-n", "machdep.cpu.brand_string"])
    else:
        info["os"] = os_release()
        info["wsl2"] = is_wsl()
    return info


def accelerator():
    """What the machine says about its accelerator and the memory it may address."""
    info = {}
    smi = nvidia_smi()
    if smi:
        out = run([smi, "--query-gpu=name,driver_version,memory.total", "--format=csv,noheader"])
        if out:
            info["nvidia_smi"] = out.splitlines()
    nvcc = shutil.which("nvcc") or ("/usr/local/cuda/bin/nvcc" if os.path.exists("/usr/local/cuda/bin/nvcc") else None)
    if nvcc:
        info["nvcc"] = first_line(run([nvcc, "--version"]), "release")
    vk = run(["vulkaninfo", "--summary"])
    if vk:
        devices = [ln.split("=", 1)[1].strip() for ln in vk.splitlines() if "deviceName" in ln and "=" in ln]
        drivers = [ln.split("=", 1)[1].strip() for ln in vk.splitlines() if "driverName" in ln and "=" in ln]
        info["vulkan"] = {"devices": devices, "drivers": drivers}
    rocm = run(["rocminfo"])
    if rocm:
        agents = [ln.split()[1] for ln in rocm.splitlines() if ln.split()[:1] == ["Name:"] and len(ln.split()) > 1 and ln.split()[1].startswith("gfx")]
        info["rocminfo_gfx"] = agents
    amd_smi = run(["amd-smi", "version"])
    if amd_smi:
        info["amd_smi_version"] = first_line(amd_smi)
    for card in sorted(Path("/sys/class/drm").glob("card*/device")):
        vram, gtt = read_text(card / "mem_info_vram_total"), read_text(card / "mem_info_gtt_total")
        if vram and gtt and vram.isdigit() and gtt.isdigit():
            info.setdefault("amdgpu", {})[card.parent.name] = {
                "vram_gib": round(int(vram) / GIB, 2), "gtt_gib": round(int(gtt) / GIB, 2)}
    pages = read_text("/sys/module/ttm/parameters/pages_limit")
    if pages and pages.isdigit() and "amdgpu" in info:
        info["ttm_pages_limit"] = int(pages)
        info["ttm_limit_gib"] = round(int(pages) * TTM_PAGE_BYTES / GIB, 2)
    if sys.platform == "darwin":
        disp = run(["system_profiler", "SPDisplaysDataType"], timeout=60)
        info["mac_gpu"] = first_line(disp, "Chipset Model")
        info["mac_gpu_cores"] = first_line(disp, "Total Number of Cores")
        wired = run(["sysctl", "-n", "iogpu.wired_limit_mb"])
        info["iogpu_wired_limit_mb"] = int(wired) if wired and wired.isdigit() else wired
        try:
            import mlx.core as mx  # noqa: PLC0415 - optional, Track M only
            dev = mx.device_info()
            for key in ("max_recommended_working_set_size", "memory_size"):
                if key in dev:
                    info["mlx_" + key] = dev[key]
        except (ImportError, AttributeError, RuntimeError):
            pass
    return info


def tool_versions():
    versions = {}
    for name, cmd, pick in (
        ("uv", ["uv", "--version"], ""),
        ("git", ["git", "--version"], ""),
        ("docker", ["docker", "--version"], ""),
        ("docker_compose", ["docker", "compose", "version", "--short"], ""),
        ("podman", ["podman", "--version"], ""),
        ("cmake", ["cmake", "--version"], ""),
    ):
        out = run(cmd)
        if out:
            versions[name] = first_line(out, pick)
    runtimes = run(["docker", "info", "--format", "{{range $name, $rt := .Runtimes}}{{$name}} {{end}}"])
    if runtimes:
        versions["docker_runtimes"] = runtimes.split()
    hf = shutil.which("hf")
    if hf:
        out = run([hf, "version", "--format", "json"])
        try:
            versions["hf"] = json.loads(out)["version"] if out else None
        except (ValueError, KeyError):
            versions["hf"] = None
    for module in ("torch", "mlx.core", "transformers", "huggingface_hub"):
        try:
            mod = __import__(module, fromlist=["__version__"])
        except ImportError:
            continue
        versions[module] = str(getattr(mod, "__version__", "unknown"))
        if module == "torch":
            versions["torch_cuda"] = getattr(mod.version, "cuda", None)
            versions["torch_hip"] = getattr(mod.version, "hip", None)
    return versions


def storage():
    models = Path.home() / "models"
    hf_home = os.environ.get("HF_HOME")
    usage = shutil.disk_usage(models if models.exists() else Path.home())
    return {
        "hf_home": hf_home,
        "hf_home_is_library": hf_home == str(models / "hf"),
        "hf_token_stored": bool(hf_home) and (Path(hf_home) / "token").exists(),
        "library_readme": (models / "README.md").exists(),
        "free_disk_gib": round(usage.free / GIB, 1),
    }


def build_record(track, note):
    return {
        "lab": "part-05/prepare-your-machine",
        "date": date.today().isoformat(),
        "track": track,
        "note": note,
        "machine": machine(),
        "python": platform.python_version(),
        "python_executable": sys.executable,
        "accelerator": accelerator(),
        "tools": tool_versions(),
        "storage": storage(),
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--labbook", default=None,
                        help="append one JSON line to this existing file (for example labbook.md)")
    parser.add_argument("--track", default="auto", choices=TRACKS,
                        help="platform track; 'auto' guesses from what is installed")
    parser.add_argument("--note", default="", help="free text kept with the record")
    parser.add_argument("--print", dest="show", action="store_true",
                        help="print the record and write nothing")
    args = parser.parse_args()

    if args.labbook and not args.show and not Path(args.labbook).is_file():
        sys.exit(f"record-baseline: {args.labbook} does not exist. Run this from ~/llm-course, "
                 "where Part 1 created labbook.md, or pass its full path.")

    track = guess_track() if args.track == "auto" else args.track
    record = build_record(track, args.note)
    print(json.dumps(record, indent=2, sort_keys=True))

    if args.show or not args.labbook:
        if not args.labbook:
            print("\nNothing written: pass --labbook labbook.md to record this.")
        return

    with Path(args.labbook).open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, sort_keys=True) + "\n")
    print(f"\nrecorded in {args.labbook}")


if __name__ == "__main__":
    main()