#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# setup_guest_driver.sh - Setup NVIDIA driver in QEMU guest
# Run this INSIDE the guest VM after installation

set -euo pipefail

echo "=========================================="
echo "Setting up NVIDIA driver in guest"
echo "=========================================="

# Check if running as root
if [[ $EUID -ne 0 ]]; then
    echo "ERROR: This script must be run as root"
    exit 1
fi

# Check kernel version
KERNEL_VER=$(uname -r)
echo "Kernel: $KERNEL_VER"

# Install dependencies
echo "Installing build dependencies..."
apt-get update
apt-get install -y \
    build-essential \
    dkms \
    linux-headers-"$KERNEL_VER" \
    linux-headers-generic \
    pciutils \
    clinfo

# Check for NVIDIA GPU in guest
echo "Checking for NVIDIA GPU in guest..."
lspci -d 10de: | grep -E "VGA|3D controller" || {
    echo "WARNING: No NVIDIA GPU detected in guest"
    echo "Make sure QEMU is started with -device vfio-pci,sysfsdev=..."
}

# Option 1: Install from Ubuntu repository (easiest)
echo ""
echo "Choose installation method:"
echo "  1) Ubuntu repository (nvidia-driver-550-open)"
echo "  2) NVIDIA runfile (latest)"
echo "  3) Build from source (open-gpu-kernel-modules)"
read -p "Select [1-3]: " choice

case $choice in
    1)
        echo "Installing from Ubuntu repository..."
        apt-get install -y nvidia-driver-550-open
        ;;
    2)
        echo "Downloading NVIDIA runfile..."
        RUNFILE="NVIDIA-Linux-x86_64-570.144.run"
        wget -q "https://us.download.nvidia.com/XFree86/Linux-x86_64/570.144/$RUNFILE"
        chmod +x "$RUNFILE"
        ./"$RUNFILE" --silent --dkms --no-opengl-files
        ;;
    3)
        echo "Building from open-gpu-kernel-modules source..."
        git clone https://github.com/NVIDIA/open-gpu-kernel-modules.git /tmp/open-gpu-kernel-modules
        cd /tmp/open-gpu-kernel-modules
        git checkout 570.144
        make -C kernel-open/modules SYSSRC=/usr/src/linux-headers-"$KERNEL_VER" -j$(nproc)
        make -C kernel-open/modules SYSSRC=/usr/src/linux-headers-"$KERNEL_VER" modules_install
        depmod -a
        ;;
    *)
        echo "Invalid choice"
        exit 1
        ;;
esac

# Install firmware
echo "Installing GSP firmware..."
mkdir -p /lib/firmware/nvidia/tu102/gsp
# Firmware files should be copied from host or downloaded
# For TU104 (Turing), we need:
# - bootloader-535.113.01.bin
# - booter_load-535.113.01.bin
# - gsp-535.113.01.bin
echo "Copy firmware files to /lib/firmware/nvidia/tu102/gsp/"
echo "  (bootloader, booter_load, gsp binaries)"

# Load driver
echo "Loading NVIDIA driver..."
modprobe nvidia

# Verify
echo ""
echo "=========================================="
echo "Verification"
echo "=========================================="
nvidia-smi
echo ""
dmesg -T | grep -i nvidia | tail -20

echo ""
echo "Setup complete! Recommended: reboot the guest"