#!/bin/bash
# ============================================================
#  fix-gpu-host.sh  —  Run this on the PROXMOX HOST
#  Fixes NVIDIA driver version mismatch so GPU encoding works
# ============================================================
set -euo pipefail

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'

log()  { echo -e "${GREEN}[✓]${NC} $1"; }
warn() { echo -e "${YELLOW}[!]${NC} $1"; }
fail() { echo -e "${RED}[✗]${NC} $1"; exit 1; }

echo ""
echo "======================================================"
echo "  NVIDIA Driver Fix — Proxmox Host"
echo "======================================================"
echo ""

# Must be root
[[ $EUID -ne 0 ]] && fail "Please run as root: sudo bash fix-gpu-host.sh"

TARGET="580.65.06-1"
CURRENT=$(modinfo nvidia 2>/dev/null | awk '/^version:/{print $2}' || echo "none")

log "Current kernel module version: ${CURRENT}"
log "Target version: 580.65.06"
echo ""

if [[ "$CURRENT" == "580.65.06" ]]; then
    log "Already at correct version. Nothing to do."
    echo ""
    echo "Now run fix-gpu-vm.sh inside the VM."
    exit 0
fi

warn "The host driver (${CURRENT}) does not match the VM userspace (580.126.09)."
warn "This script will install nvidia-open-580=580.65.06 and reboot the host."
echo ""
read -p "Continue? (yes/no): " CONFIRM
[[ "$CONFIRM" != "yes" ]] && { warn "Aborted."; exit 0; }

echo ""
log "Updating package lists..."
apt-get update -qq

log "Installing nvidia-open-580=580.65.06-1 ..."
apt-get install -y "nvidia-open-580=${TARGET}" || \
    fail "Installation failed. Check that the CUDA local repo is still configured."

log "Pinning package to prevent accidental upgrades..."
apt-mark hold nvidia-open-580

echo ""
log "Done! The host will reboot in 10 seconds."
warn "After reboot, SSH back in and run fix-gpu-vm.sh inside the VM."
echo ""
sleep 10
reboot