#!/bin/bash

# Get the scripts directory
SCRIPTDIR=$(dirname "$(realpath "$0")")
cd "$SCRIPTDIR"

# Get the config paths
source "$SCRIPTDIR/lib/paths.sh"

# Get the CPU Vendor
CPU_VENDOR=$(cat /proc/cpuinfo | grep vendor | head -1 | cut -f 2 | cut -d " " -f 2)
CMDLINE="iommu=pt"

# Adjust our kernel_args based on cpu vendor
if [ "$CPU_VENDOR" == "GenuineIntel" ];
then
    CMDLINE="$CMDLINE intel_iommu=on"
elif [ "$CPU_VENDOR" == "AuthenticAMD" ];
then
    CMDLINE="$CMDLINE amd_iommu=on"
fi

# Clear the screen
clear

# Show the user a warning before we start
printf "Welcome to QuickPassthrough!
Written for systems using initramfs.

The setup done by this script is quite complex and is prone to human error or hardware incompatibilities.
It is HIGHLY RECOMMENDED to make a backup/snapshot of your system using something like timeshift or snapper before starting.

Once everything is configured, your 2nd graphic card will hopefully be dedicated for use inside a virtual machine.

Even though that this script is intended to make this setup easier, it is recommended that you
read through a guide and learn what is actually being done so that you can get familiar the process.
A full documentation for debian/ubuntu systems can be found here: https://github.com/hikariknight/vfio-setup-docs/wiki

Press ENTER to continue once you have made a backup of your system.
"
read -r

# Separator
printf "
############################################################
"

echo "This script assumes a few things:
* You have already enabled IOMMU, VT-d and/or AMD-v inside your UEFI/BIOS advanced settings.
* You have already added \"$CMDLINE\" to your
  kernel boot arguments and booted your system with these kernel arguments active.
* You are comfortable with navigating and changing settings in your UEFI/BIOS.
* You know how edit your bootloader configuration and kernel arguments.
* Your Linux distribution is an EFI installation (important to get VFIO working).

NOTE: If your computer no longer fully shut down after enabling IOMMU, then there is possibly a bug
      with your motherboard and a piece of hardware in your system, it only prevents you from using
      the system in a headless mode with working shutdown and is otherwise just an annoying
      quirk with IOMMU on some boards.

This is a list of prerequisites you will be needing before starting with VFIO:
* 2 very different GPUs (iGPU/APU included), the easiest combination is to have 2 from different vendors (amd/intel/nvidia)
  if both cards share the same device id (ex: both are identified as 1022:145c), then passthrough will
  not be possible unless you swap out one of the cards.
* A \"ghost display\" dummy plug for your second graphic card (or having it hooked to a separate input on your monitor).
* If you are planning to use the inegrated GPU on your CPU, make sure your monitor is connected to it and working/enabled before continuing.
* Preferably a motherboard verified to work with IOMMU and with good IOMMU groups.
  https://reddit.com/r/vfio is a good resource for this info.
  (If you are unsure, you will find out while using this script)

* It is also highly recommended to have access to your VM through VNC (RDP will not work) as once you pass through
  a graphic card, it will most likely not possible to interact with it through spice and the normal qemu display window!

Press ENTER to start creating your config from scratch.
NOTE: continuing will delete the contents of \"$SCRIPTDIR/config\"
"
read -r

# Separator
printf "
############################################################
"

if [ -d "$SCRIPTDIR/config" ];
then
    rm -r "$SCRIPTDIR/config"
fi

# Make the directories
mkdir -p "$SCRIPTDIR/$MODPROBE"
mkdir -p "$SCRIPTDIR/$QUICKEMU"

# Write the cmdline file
echo "$CMDLINE" > "$SCRIPTDIR/config/kernel_args"

# Copy system configs that exists into our config folder so we can safely edit them
if [ -f "/etc/modules" ];
then
    # This copies /etc/modules without the vfio module lines
    grep -v "vfio" "/etc/modules" > "$SCRIPTDIR/$ETCMODULES"
fi

if [ -f "/etc/default/grub" ];
then
    # Create the default folder
    mkdir -p "$SCRIPTDIR/$DEFAULT"

    # Copy grub config
    cp "/etc/default/grub" "$SCRIPTDIR/$DEFAULT/grub"
fi

if [ -f "/etc/initramfs-tools/modules" ];
then
    # Create the initramfs folder
    mkdir -p "$SCRIPTDIR/$INITRAMFS"

    # This copies /etc/initramfs-tools/modules without the vfio modules
    grep -v "vfio" "/etc/initramfs-tools/modules" > "$SCRIPTDIR/$INITRAMFS/modules"
fi

if [ -f "/etc/mkinitcpio.conf" ];
then
    # Copy mkinitcpio.conf to our config folder
    cp "/etc/mkinitcpio.conf" "$SCRIPTDIR/$MKINITCPIO"
fi

if [ -f "/etc/dracut.conf" ];
then
    # Create the dracut folder
    mkdir -p "$SCRIPTDIR/$DRACUT"
    # Create a dracut.conf.d file
    touch "$SCRIPTDIR/$DRACUT/10-vfio.conf" 
fi

# Run ls-iommu so we can verify that IOMMU properly working
LS_IOMMU=$("$SCRIPTDIR/utils/ls-iommu")

# Detect if IOMMU is disabled (output will start with "IOMMU Group *:" if disabled)
if [[ $LS_IOMMU =~ ^IOMMU[[:space:]]Group[[:space:]]+\*: ]];
then
    # Tell user to enable IOMMU then try again
    echo "IOMMU IS NOT ENABLED!
Please enable IOMMU, VT-d or AMD-v inside your UEFI/BIOS and add \"$CMDLINE\"
to your kernel boot arguments and reboot your system, then re-run this script!
"
    exit 1

else
    # Show the output of ls-iommu for manual verification
    echo "$LS_IOMMU"
fi

echo ""

# Have user visually verify that IOMMU is working
read -r -p "Is there more than 1 group in the output above? [y/N]: " YESNO
case "${YESNO}" in
    [Yy]*)
        echo ""
    ;;
    [Nn]*)
        echo "Please enable IOMMU, VT-d or AMD-v inside your UEFI/BIOS and add \"$CMDLINE\"
to your kernel boot arguments and reboot your system, then re-run this script!
"
        exit 1
    ;;
    *)
        echo "Please enable IOMMU, VT-d or AMD-v inside your UEFI/BIOS and add \"$CMDLINE\"
to your kernel boot arguments and reboot your system, then re-run this script!
"
        exit 1
    ;;
esac

exec "$SCRIPTDIR/lib/get_GPU.sh"