#!/bin/bash

function insert_MODULES() {
    # Get the header and enabled modules separately from the /etc/modules file
    local MODULES_HEADER
    local MODULES_ENABLED
    local VENDOR_RESET
    MODULES_HEADER=$(head -n $1 "$2" | grep -P "^#" | grep -v "# Added by quickpassthrough")
    MODULES_ENABLED=$(cat "$2" | grep -vP "^#" | grep -v "vendor-reset")
    VENDOR_RESET=0
    
    # If vendor-reset is present
    if grep -q "vendor-reset" "$2" ;
    then
        VENDOR_RESET=1
    fi

    # Write header
    echo "$MODULES_HEADER" > "$2"
    
    # If vendor-reset existed from before
    if [ $VENDOR_RESET == 1 ];
    then
        # Write vendor-reset as the first module!
        echo "vendor-reset" >> "$2"
    fi
    
    # Append vfio 
    printf "
# Added by quickpassthrough #
vfio
vfio_iommu_type1
vfio_pci
vfio_virqfd
#############################
" >> "$2"

    # Write the previously enabled modules under vfio in the load order
    echo "$MODULES_ENABLED" >> "$2"
}

function set_MODULES () {
    # Get the config paths
    source "$SCRIPTDIR/lib/paths.sh"
    
    # Insert modules in the correct locations as early as possible without
    # conflicting with vendor-reset module if it is enabled
    insert_MODULES 4 "$SCRIPTDIR/$ETCMODULES"
    insert_MODULES 11 "$SCRIPTDIR/$INITRAMFS/modules"

    # Assign the GPU device ids to a variable
    GPU_DEVID="$1"

    # Get the kernel_args file content
    CMDLINE=$(cat "$SCRIPTDIR/config/kernel_args")

    # Ask if we shall disable video output on this card
    echo "
Disabling video output in Linux for the card you want to use in a VM
will make it easier to successfully do the passthrough without issues."
    read -p "Do you want to force disable video output in linux on this card? [Y/n]: " DISABLE_VGA
    case "${DISABLE_VGA}" in
    [Yy]*)
        # Update kernel_args file
        echo "${CMDLINE} vfio_pci.ids=${GPU_DEVID} vfio_pci.disable_vga=1" > "$SCRIPTDIR/config/kernel_args"

        # Update GPU_DEVID
        GPU_DEVID="$GPU_DEVID disable_vga=1"
    ;;
    [Nn]*)
        echo ""
    ;;
    *)
        # Update kernel_args file
        echo "${CMDLINE} vfio_pci.ids=${GPU_DEVID} vfio_pci.disable_vga=1" > "$SCRIPTDIR/config/kernel_args"

        # Update GPU_DEVID
        GPU_DEVID="$GPU_DEVID disable_vga=1"
    ;;
    esac

    # Write the vfio modprobe config
    printf "## This is an autogenerated file that stubs your graphic card for use with vfio
## This file should be placed inside /etc/modprobe.d/
# Uncomment the line below to \"hardcode\" your graphic card to be bound to the vfio-pci driver.
# In most cases this should not be neccessary, it will also prevent you from turning off vfio in the bootloader.
#options vfio_pci ids=%s

# Make sure vfio_pci is loaded before these modules: nvidia, nouveau, amdgpu and radeon
softdep nvidia pre: vfio vfio_pci
softdep nouveau pre: vfio vfio_pci
softdep amdgpu pre: vfio vfio_pci
softdep radeon pre: vfio vfio_pci
"  "${GPU_DEVID}"  > "$SCRIPTDIR/$MODPROBE/vfio.conf"

    exec "$SCRIPTDIR/lib/get_USB_CTL.sh"
}


function main () {
    SCRIPTDIR=$(dirname "$(which $0)" | perl -pe "s/\/\.\.\/lib//" | perl -pe "s/\/lib$//")

    set_MODULES "$1"
}

main "$1"