#!/usr/bin/env bash
# ============================================================
# fix_gpu_host.sh — Run on the HOST (not container) to fix
# CUDA on RTX 3090 with NVIDIA open kernel module.
#
# Problem: The open kernel module (nvidia-open) blocks CUDA
# on GeForce RTX 3090. nvidia-smi works but cuInit() fails
# with error 999 (UNKNOWN). This is a known driver issue.
#
# Solution: Switch from nvidia-open to the proprietary driver.
# ============================================================
set -euo pipefail

echo "============================================"
echo " GPU Fix — RTX 3090 + NVIDIA Driver"
echo "============================================"

# Check current driver
echo ""
echo "Current driver:"
nvidia-smi --query-gpu=name,driver_version --format=csv 2>/dev/null || echo "nvidia-smi not found"

# Check kernel module type
echo ""
echo "Kernel module type:"
if lsmod | grep -q nvidia; then
    if modinfo nvidia 2>/dev/null | grep -q "open"; then
        echo "  OPEN kernel module detected — THIS IS THE PROBLEM"
        echo "  GeForce GPUs need the proprietary module for full CUDA"
    else
        echo "  Proprietary kernel module detected"
    fi
else
    echo "  No nvidia module loaded"
fi

echo ""
echo "=== Fix options ==="
echo ""
echo "Option 1: Switch to proprietary driver (recommended)"
echo "  sudo apt-get remove --purge nvidia-dkms-580-open nvidia-driver-580-open"
echo "  sudo apt-get install nvidia-driver-580 nvidia-dkms-580"
echo "  sudo reboot"
echo ""
echo "Option 2: For RunPod specifically"
echo "  Stop this pod and restart with template:"
echo "    'RunPod Pytorch 2.4+' or 'NVIDIA CUDA 12.4'"
echo "  Ensure the template does NOT say 'open kernel'"
echo ""
echo "Option 3: Quick fix (try without reboot)"
echo "  sudo rmmod nvidia-uvm nvidia-drm nvidia-modeset nvidia"
echo "  sudo modprobe nvidia"
echo "  sudo modprobe nvidia-uvm"
echo "  sudo nvidia-persistenced --user root"
echo ""

# Auto-detect and apply if running as root on host
if [[ "$(id -u)" -eq 0 ]] && [[ -f /proc/driver/nvidia/version ]]; then
    KERNEL_TYPE=$(cat /proc/driver/nvidia/version | grep -i "open" && echo "open" || echo "proprietary")
    if [[ "$KERNEL_TYPE" == "open" ]]; then
        echo "⚠️  Open kernel module detected. Would you like to attempt fix? (y/n)"
        read -r answer
        if [[ "$answer" == "y" ]]; then
            echo "Attempting fix..."
            rmmod nvidia-uvm nvidia-drm nvidia-modeset nvidia 2>/dev/null || true
            modprobe nvidia 2>/dev/null || true
            modprobe nvidia-uvm 2>/dev/null || true
            nvidia-persistenced --user root 2>/dev/null || true
            echo "Fix applied. Test with: python3 -c 'import torch; print(torch.cuda.is_available())'"
        fi
    else
        echo "✅ Proprietary driver detected — CUDA should work"
    fi
fi