#!/bin/bash

## Last update 2018/11/20

## Testing pourposes
#set -x

## What do I need
## To output the names of the devices
## to select
##   To check where they are
##   Who they are attached to
##   What is using them
## To unload the kernel modules automaticly
## to modify the config file

## So this is a fucking nightmare
## The script will only output what you have to copy to the config file. something like:
# IOMMU_GPU=06:00.0
# IOMMU_GPU_AUDIO=06:00.1
# IOMMU_USB=07:00.3
# VIRSH_GPU=pci_0000_06_00_0
# VIRSH_GPU_AUDIO=pci_0000_06_00_1
# VIRSH_USB=pci_0000_07_00_3
# videoid="10de 1184"
# audioid="10de 0e0a"
# usbid="1022 145f"
# videobusid="0000:06:00.0"
# audiobusid="0000:06:00.1"
# usbbusid="0000:07:00.3"

find_iommu(){
  for iommu_group in \
    $(find /sys/kernel/iommu_groups/ -maxdepth 1 -mindepth 1 -type d); do \
    echo "IOMMU group $(basename "$iommu_group")"; \
    for device in $(ls -1 "$iommu_group"/devices/); do \
      echo -n $'\t'; lspci -nns "$device"; \
    done; \
  done
  ## Sample output
  ## 06:00.0 VGA compatible controller [0300]: NVIDIA Corporation GK104 [GeForce GTX 770] [10de:1184] (rev a1)
	## 06:00.1 Audio device [0403]: NVIDIA Corporation GK104 HDMI Audio Controller [10de:0e0a] (rev a1)
}


get_the_gpu(){
  ## WIP
  find_iommu | grep $SELECTED_GPU | grep $1 | awk '{print $1}'
  find_iommu | grep $SELECTED_GPU | grep "VGA compatible controller" | grep $1 | awk -F'[][]' '{print $4}'
}

## Ask for the GPU
## For now it only works with NVIDIA so set the variable to NVIDIA.
read -p "Which GPU do you have? AMD/NVIDIA: " SELECTED_GPU
case $SELECTED_GPU in
  [Aa][Mm][Dd]) SELECTED_GPU=AMD ;;
  [Nn][Vv][Ii][Dd][Ii][Aa]) SELECTED_GPU=NVIDIA ;;
  *) echo "I do not recognise that GPU"; exit ;;
esac
echo ""
if [ -z "$(find_iommu | grep $SELECTED_GPU)" ]
then
  echo "Seems like you dont have a $SELECTED_GPU card"
else
  echo "You have a $SELECTED_GPU card"
fi
#SELECTED_GPU=NVIDIA


## Find the GPUs in the system and make the user select them.
## Inputting the GPU number [770] should be enough.
#find_iommu | grep $SELECTED_GPU | grep "VGA compatible controller" | awk -F'[][]' '{print $4}'
echo "Available GPUs"
echo "  $(find_iommu | grep $SELECTED_GPU | grep "VGA compatible controller" | awk -F'[][]' '{print $4}')"
read -p "Which GPU do you wish to use? $( for i in $( find_iommu | grep $SELECTED_GPU | grep "VGA compatible controller" | awk -F'[][]' '{print $4}' | grep -o '[0-9]\+'); do echo "$i"; done): " GPU_TO_BE_USED
## If you do not input anything then it will asume there is only one GPU.
#if [ -z "$GPU_TO_BE_USED" ];
#then
#  get_the_gpu $( find_iommu | grep $SELECTED_GPU | grep "VGA compatible controller" | awk -F'[][]' '{print $4}' | grep -o '[0-9]\+' )
#else
#  get_the_gpu $GPU_TO_BE_USED
#fi

## Set the location for the config file
## I have no idea how to do it so it's diabled and the config file will be saved in the folder the script is executed
#echo ""
#read -p "Select the path for the config file to be saved [/tmp/config]: " $CONFIG_LOCATION
#if [ -z "$CONFIG_LOCATION" ]
#then
#  CONFIG_LOCATION=/tmp/config
#elif $CONFIG_LOCATION

## Set the locations for the images and stuff
#echo ""
#read -p "Select the path where you have ovmf, virtio and else [/home/$(whoami)/vm/]: " PATH_TO_OVMF



IOMMU_GPU=$( find_iommu | grep $SELECTED_GPU | grep "VGA compatible controller" | awk '{print $1}' )
IOMMU_GPU_AUDIO=$( find_iommu | grep $SELECTED_GPU | grep "Audio device" | awk '{print $1}' )
VIRSH_GPU="pci_0000_$( echo $IOMMU_GPU | awk -F ":" '{print $1}' )_$( echo $IOMMU_GPU | awk -F ":" '{print $2}' | cut -f1 -d"." )_$( echo $IOMMU_GPU | awk -F "." '{print $2}' )"
VIRSH_GPU_AUDIO="pci_0000_$( echo $IOMMU_GPU_AUDIO | awk -F ":" '{print $1}' )_$( echo $IOMMU_GPU_AUDIO | awk -F ":" '{print $2}' | cut -f1 -d"." )_$( echo $IOMMU_GPU_AUDIO | awk -F "." '{print $2}' )"
videoid="$( find_iommu | grep $SELECTED_GPU | grep "VGA compatible controller" | awk -F'[][]' '{print $6}' | awk -F ":" '{print $1 " " $2}' )"
audioid="$( find_iommu | grep $SELECTED_GPU | grep "Audio device" | awk -F'[][]' '{print $4}' | awk -F ":" '{print $1 " " $2}' )"
videobusid=0000:$( find_iommu | grep $SELECTED_GPU | grep "VGA compatible controller" | awk '{print $1}' )
audiobusid=0000:$( find_iommu | grep $SELECTED_GPU | grep "Audio device" | awk '{print $1}' )


echo ""
echo "You have to copy the following into the config file:"
echo ""
echo "IOMMU_GPU=$IOMMU_GPU"
echo "IOMMU_GPU_AUDIO=$IOMMU_GPU_AUDIO"
echo "VIRSH_GPU=$VIRSH_GPU"
echo "VIRSH_GPU_AUDIO=$VIRSH_GPU_AUDIO"
echo "videoid=$videoid"
echo "audioid=$audioid"
echo "videobusid=$videobusid"
echo "audiobusid=$audiobusid"
