#!/usr/bin/env python3

import re, subprocess, sys
from pathlib import Path

NEXUS_ROOT  = Path.home() / "nexus-core"
NVIDIA_REQS = NEXUS_ROOT / "requirements-nvidia.txt"

CUDA_TO_WHEEL = [
    ((12, 8), "cu128"),
    ((12, 6), "cu126"),
    ((12, 4), "cu124"),
    ((12, 1), "cu121"),
    ((11, 8), "cu118"),
]

def detect_cuda():
    try:
        out = subprocess.run(["nvidia-smi"], capture_output=True, text=True, timeout=10).stdout
        m = re.search(r"CUDA Version:\s*(\d+)\.(\d+)", out)
        if m:
            return int(m.group(1)), int(m.group(2))
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    return None, None

def wheel_suffix(major, minor):
    for (req_major, req_minor), suffix in CUDA_TO_WHEEL:
        if (major, minor) >= (req_major, req_minor):
            return suffix
    return "cu118"

def main():
    print("Detecting NVIDIA GPU...")
    major, minor = detect_cuda()

    if major is None:
        print("Error: nvidia-smi not found or CUDA version unreadable.")
        print("Ensure NVIDIA drivers are installed and nvidia-smi is on your PATH.")
        sys.exit(1)

    print(f"CUDA {major}.{minor} detected.")
    suffix = wheel_suffix(major, minor)
    print(f"PyTorch wheel: {suffix}")

    NVIDIA_REQS.write_text(f"""\
# --- Force NVIDIA/CUDA Priority ---
--index-url https://download.pytorch.org/whl/{suffix}
--extra-index-url https://pypi.org/simple

-r requirements-base.txt

# GPU Compute Stack
torch
torchaudio
torchvision
""")

    print(f"\nWritten: {NVIDIA_REQS}")
    print("Run 'ncp backup' to push it to the router.")

if __name__ == "__main__":
    main()