import sys
import os
import shutil

def check_gpu():
    print("\n--- GPU Check (via nvidia-smi) ---")
    if shutil.which("nvidia-smi"):
        ret = os.system("nvidia-smi")
        if ret == 0:
            print("✅ NVIDIA Driver detected.")
        else:
            print("⚠️ nvidia-smi found but returned error.")
    else:
        print("❌ nvidia-smi not found. CUDA might not be in PATH.")

def check_llama_cuda():
    print("\n--- Llama.cpp CUDA Check ---")
    try:
        from llama_cpp import Llama
        print("✅ llama-cpp-python is installed.")
        print(f"Llama.cpp package location: {sys.modules['llama_cpp'].__file__}")
    except ImportError:
        print("❌ llama-cpp-python is NOT installed.")
        sys.exit(1)

if __name__ == "__main__":
    print(f"Python Version: {sys.version}")
    check_gpu()
    check_llama_cuda()
    print("\nEnvironment check complete.")