# vim: set ft=make :

# Setup and configure virtualization and vfio
setup-virtualization ACTION="":
    #!/usr/bin/bash
    source /usr/lib/ujust/ujust.sh
    # Check if we are running on a Steam Deck
    if /usr/libexec/hardware/valve-hardware; then
      echo "Virtualization is not supported on Steam Deck"
      exit 0
    else
      if [ "$(systemctl is-enabled libvirtd.service)" == "disabled" ]; then
        echo "${b}libvirtd${n} service is ${red}disabled${n}!"
        echo "${green}enabling${n} and starting libvirtd"
        echo "If virt-manager says libvirtd.sock is not available after a big update, re-run this command."
        sudo systemctl enable --now libvirtd 2> /dev/null
        echo "Press ESC if you want to exit and do not need to do anything"
      fi
    fi
    OPTION={{ ACTION }}
    if [ "$OPTION" == "help" ]; then
      echo "Usage: ujust configure-grub <option>"
      echo "  <option>: Specify the quick option to skip the prompt"
      echo "  Use 'enable' to select Enable Virtualization"
      echo "  Use 'disable' to select Disable Virtualization"
      echo "  Use 'group' to select Add $USER to libvirt group"
      echo "  Use 'vfio-on' to select Enable VFIO drivers"
      echo "  Use 'vfio-off' to select Disable VFIO drivers"
      echo "  Use 'shm' to select Autocreate Looking-Glass shm"
      exit 0
    elif [ "$OPTION" == "" ]; then
      echo "${bold}Virtualization Setup${normal}"
      echo "NOTE: Enabling Virtualization will layer virt-manager and qemu"
      echo "      this will slow down system updates by a lot."
      OPTION=$(Choose \
        "Enable Virtualization" \
        "Disable Virtualization" \
        "Add $USER to libvirt group" \
        "Enable VFIO drivers" \
        "Disable VFIO drivers" \
        "Autocreate Looking-Glass shm" \
      )
    fi
    if [[ "${OPTION,,}" =~ ^enable[[:space:]]virt ]]; then
      virt_test=$(rpm-ostree status | grep -A 4 "●" | grep "virt-manager")
      if [[ -z ${virt_test} ]]; then
        echo "Installing QEMU and virt-manager..."
        rpm-ostree install -y virt-manager edk2-ovmf qemu
        rpm-ostree kargs \
        --append-if-missing="kvm.ignore_msrs=1" \
        --append-if-missing="kvm.report_ignored_msrs=0"
        echo 'Please re-run "ujust setup-virtualization" after the reboot to enable libvirtd service'
      fi
    elif [[ "${OPTION,,}" =~ ^disable[[:space:]]virt ]]; then
      virt_test=$(rpm-ostree status | grep -A 4 "●" | grep "virt-manager")
      if [[ ${virt_test} ]]; then
        if [ "$(systemctl is-enabled libvirtd.service)" == "enabled" ]; then
          echo "${red}Disabling${n} libvirtd before removal"
          sudo systemctl disable --now libvirtd 2> /dev/null
        fi
        echo "Removing QEMU and virt-manager..."
        rpm-ostree remove -y virt-manager edk2-ovmf qemu
        rpm-ostree kargs \
        --delete-if-present="kvm.ignore_msrs=1" \
        --delete-if-present="kvm.report_ignored_msrs=0"
        echo 'Please re-run "ujust enable-virtualization" after the reboot to finish setup'
      fi
    elif [[ "${OPTION,,}" =~ (^enable[[:space:]]vfio|vfio-on) ]]; then
      echo "Enabling VFIO..."
      VIRT_TEST=$(rpm-ostree kargs)
      CPU_VENDOR=$(grep "vendor_id" "/proc/cpuinfo" | uniq | awk -F": " '{ print $2 }')
      VENDOR_KARG="unset"
      if [[ ${VIRT_TEST} == *kvm.report_ignored_msrs* ]]; then
        echo 'add_drivers+=" vfio vfio_iommu_type1 vfio-pci "' | sudo tee /etc/dracut.conf.d/vfio.conf
        sudo touch /etc/bazzite/initramfs/rebuild
        if [[ ${CPU_VENDOR} == "AuthenticAMD" ]]; then
          VENDOR_KARG="amd_iommu=on"
        elif [[ ${CPU_VENDOR} == "GenuineIntel" ]]; then
          VENDOR_KARG="intel_iommu=on"  
        fi
        if [[ ${VENDOR_KARG} == "unset" ]]; then
          echo "Failed to get CPU vendor, exiting..."
          exit 1
        else
          rpm-ostree kargs \
            --append-if-missing="${VENDOR_KARG}" \
            --append-if-missing="iommu=pt" \
            --append-if-missing="rd.driver.pre=vfio_pci" \
            --append-if-missing="vfio_pci.disable_vga=1"
          echo "VFIO will be enabled on next boot, make sure you enable IOMMU, VT-d or AMD-v in your BIOS!"
          echo "Please understand that since this is such a niche use case, support will be very limited!"
          echo "To add your unused/second GPU device ids to the vfio driver by running"
          echo 'rpm-ostree kargs --append-if-missing="vfio-pci.ids=xxxx:yyyy,xxxx:yyzz"'
          echo "NOTE: Your second GPU will not be usable by the host after you do this!"
        fi
      fi
    elif [[ "${OPTION,,}" =~ (^disable[[:space:]]vfio|vfio-off) ]]; then
      echo ""
      echo "Make sure you have ${b}disabled autostart of all VMs using VFIO${n} before continuing!"
      CONFIRM=$(Choose Cancel Continue)
      if [ "$CONFIRM" == "Continue" ]; then
        echo "Disabling VFIO..."
        VFIO_IDS="$(rpm-ostree kargs | sed -E 's/.+(vfio_pci.ids=.+\s)/\1/' | awk '{ print $1 }' | grep vfio_pci.ids)"
        VFIO_IDS_KARG=""
        if [ -n "$VFIO_IDS" ]; then
          echo "Found VFIO ids in kargs, adding the below line to removal list"
          echo "$VFIO_IDS"
          VFIO_IDS_KARG="--delete-if-present=\"$VFIO_IDS\""
        fi
        echo "Removing dracut modules"
        sudo rm /etc/dracut.conf.d/vfio.conf
        echo "Issuing initramfs rebuild for next boot"
        sudo touch /etc/bazzite/initramfs/rebuild
        rpm-ostree kargs \
        --delete-if-present="iommu=pt" \
        --delete-if-present="iommu=on" \
        --delete-if-present="amd_iommu=on" \
        --delete-if-present="intel_iommu=on" \
        --delete-if-present="rd.driver.pre=vfio_pci" \
        --delete-if-present="vfio_pci.disable_vga=1" \
        --delete-if-present="vfio_pci.disable_vga=0" \
        $VFIO_IDS_KARG
      fi
    elif [[ "${OPTION,,}" =~ shm ]]; then
      echo "Creating tmpfile definition for shm file in /etc/tmpfiles.d/"
      sudo bash -c "tee << LOOKING_GLASS_TMP > /etc/tmpfiles.d/10-looking-glass.conf
      # Type Path               Mode UID  GID Age Argument
      f /dev/shm/looking-glass 0660 1000 qemu -
      LOOKING_GLASS_TMP"
      echo "Adding SELinux context record for /dev/shm/looking-glass"
      sudo semanage fcontext -a -t svirt_tmpfs_t /dev/shm/looking-glass
    elif [[ "${OPTION,,}" =~ group ]]; then
      sudo usermod -aG libvirt $USER
    fi
