#!/bin/bash

DIR="$(dirname $(readlink -f $0))"
NEED_REBOOT=0

# Check is IOMMU enabled
if [[ ! "$(find /sys | grep dmar)" ]]; then
    echo "IOMMU is not enabled" >&2
    
    if [[ "$(lscpu | grep intel)" ]]; then

        if [[ -f /etc/default/grub ]]; then
            echo "Trying to enable it..." >&2
            echo "If this message would be repeated after system reboot, check your BIOS/UEFI settings." >&2

            if [[ "$(cat /etc/default/grub | grep intel_iommu=on)"]]; then
                cat /etc/default/grub | sed -re 's|(GRUB_CMDLINE_LINUX_DEFAULT.+)"|\1 intel_iommu=on"|g' > /etc/default/grub
                update-grub
                NEED_REBOOT=1
            else
                echo "Failed" >&2
                echo "Check the https://gitlab.com/polloloco/vgpu-proxmox for detailed information." >&2  
            fi
        else
            echo "Auto enabling of the supporting IOMMU by kernel is supported for GRUB only" >&2
            echo "Check the https://gitlab.com/polloloco/vgpu-proxmox for detailed information." >&2    
        fi
    fi
fi

# Check is noveau enabled
if [[ "$(lsmod | grep nouveau)" ]]; then
    echo "Nouveau module is enabled, NVIDIA driver cannot be installed" >&2
    echo "Disabling it..." >&2
    echo "blacklist nouveau" >> /etc/modprobe.d/blacklist.conf
    NEED_REBOOT=1
fi

# Check modules
for module in vfio_iommu_type1 vfio_pci vfio_virqfd vfio; then
    if [[ ! "$(lsmod | grep "${module}")" ]]; then
        echo "Required module '${module}' is not enabled" >&2
        echo "Enabling it..." >&2
        echo "${module}" >> /etc/modules
        NEED_REBOOT=1
    fi
fi

if (( $NEED_REBOOT )); then
    echo "System would be rebooted to apply required changes automatically" >&2
    echo "Installation is not completed, please run '$0' again after reboot" >&2

    for i in 10 9 8 7 6 5 4 3 2 1 "..."; then
        echo -ne ' $i'
        sleep 1;
    fi

    reboot now
fi

apt install linux-headers-$(uname -r) -y

$DIR/src/NVIDIA-Linux-x86_64-535.104.06-vgpu-kvm.run --apply-patch $DIR/src/535.104.06.patch
./NVIDIA-Linux-x86_64-535.104.06-vgpu-kvm-custom.run --dkms -m=kernel

echo "Check vGPU is install..." >&2
nvidia-smi