#!/bin/bash

# Function to check if a command exists
command_exists() {
    type "$1" &> /dev/null ;
}

# Check if the script is run as root
if [ "$(id -u)" != "0" ]; then
   echo "This script must be run as root. Please use sudo." 1>&2
   exit 1
fi

# Update and install necessary packages
echo "Updating package list and installing necessary tools..."
apt-get update && apt-get upgrade -y
if command_exists lspci; then
    echo "lspci is already installed."
else
    echo "Installing lspci..."
    apt-get install pciutils -y
fi
if command_exists grep; then
    echo "grep is already installed."
else
    echo "Installing grep..."
    apt-get install grep -y
fi

# Identify the GPU device you want to passthrough
echo "Listing all PCI devices to identify your GPU:"
lspci
read -p "Enter the PCI ID of your GPU (e.g., 00:01.0): " gpu_id

# Verify GPU identification
echo "Verifying GPU with lspci -k..."
lspci -k | grep -i "$gpu_id"
if [ $? -ne 0 ]; then
    echo "GPU not found or incorrect ID entered." >&2
    exit 1
fi

# Configure Proxmox for GPU passthrough
echo "Configuring Proxmox to enable GPU passthrough..."
# Edit the VM configuration file, e.g., /etc/pve/qemu-server/YOUR_VM_ID.conf
read -p "Enter the ID of the VM you want to configure (e.g., 100 for vm100): " vm_id
echo "Editing VM configuration file for GPU passthrough..."
sed -i "/^[^#]*args=/ s/$/ --device $gpu_id/" /etc/pve/qemu-server/$vm_id.conf

# Optionally, add the GPU to a specific PCI slot (useful if you have multiple GPUs)
echo "Do you want to bind this GPU to a specific PCI slot in Proxmox? This might require creating a new VMDX file."
read -p "Type 'yes' or 'no': " response
if [[ "$response" == "yes" ]]; then
    read -p "Enter the desired PCI slot (e.g., 00:02.0): " pci_slot
    echo "Adding GPU to specified PCI slot..."
    # You might need to create a new VMDX file or modify an existing one, this step depends on your Proxmox setup and version.
fi

# Blacklist the GPU driver modules
echo "Do you want to blacklist the GPU driver modules? This will prevent them from loading at boot."
read -p "Type 'yes' or 'no': " response
if [[ "$response" == "yes" ]]; then
    echo "Blacklisting GPU driver modules..."
    # Create a blacklist file for the relevant drivers
    echo "${gpu_id}*" > /etc/modprobe.d/blacklist-gpu.conf
fi

# Restart the VM for changes to take effect
echo "Changes require restarting the VM. Please power off your VM:"
read -p "Enter the ID of the VM you want to restart (e.g., 100): " vm_id
virsh shutdown $vm_id
read -p "Power on the VM now? Type 'yes': " response
if [[ "$response" == "yes" ]]; then
    virsh start $vm_id
fi

echo "GPU passthrough setup is complete. Please check your VM for GPU functionality."