"""
Run this FIRST before setup.bat
It tells you which llama-cpp-python build to install.
"""

import subprocess, sys, platform

print("\n=== SunoDoc GPU Check ===\n")
print(f"OS       : {platform.system()} {platform.release()}")
print(f"Python   : {sys.version.split()[0]}")

# --- Check NVIDIA GPU via nvidia-smi ---
try:
    result = subprocess.run(
        ["nvidia-smi", "--query-gpu=name,memory.total,driver_version",
         "--format=csv,noheader"],
        capture_output=True, text=True, timeout=10
    )
    if result.returncode == 0:
        gpu_info = result.stdout.strip()
        print(f"\n✅ NVIDIA GPU found!")
        print(f"   {gpu_info}")

        # Check CUDA version
        cuda_result = subprocess.run(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader,nounits"],
            capture_output=True, text=True
        )
        ver_result = subprocess.run(
            ["nvcc", "--version"], capture_output=True, text=True
        )
        if ver_result.returncode == 0:
            cuda_line = [l for l in ver_result.stdout.split('\n') if 'release' in l]
            if cuda_line:
                print(f"   CUDA: {cuda_line[0].strip()}")

        print("\n📋 USE THIS install command in setup.bat:")
        print("   pip install llama-cpp-python --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu121")
        print("\n   (If CUDA version is 12.4+, change cu121 → cu124)")
        print("   SET GPU=CUDA")
    else:
        raise FileNotFoundError
except (FileNotFoundError, subprocess.TimeoutExpired):
    print("\n⚠️  No NVIDIA GPU detected (or drivers not installed)")
    print("   Running on CPU — model will be slower (~30-60s per response)")
    print("\n📋 USE THIS install command in setup.bat:")
    print("   pip install llama-cpp-python")
    print("   SET GPU=CPU")

# --- Check RAM ---
try:
    import psutil
    ram = psutil.virtual_memory().total / (1024**3)
    print(f"\n💾 RAM: {ram:.1f} GB")
    if ram < 6:
        print("   ⚠️  Warning: Need ~6GB free RAM. Close other apps before running.")
    elif ram < 8:
        print("   ✅ Enough RAM (tight, close other apps)")
    else:
        print("   ✅ RAM is fine")
except ImportError:
    print("\n💾 RAM: Install psutil to check (pip install psutil)")

print("\n=== Done. Now run setup.bat ===\n")