#!/bin/bash

# Args from libvirt
DOMAIN="$1"
OPERATION="$2"
SUBOPERATION="$3"
EXTRA_ARG="$4"

# Allowed CPU cores for the host
ALLOWED_CPUS="8-15,24-31"
ALL_CPUS="0-31"

# GPU
GPU_ID="1002:73bf"
GPU_ADDRESS="$(lspci -nn | grep "$GPU_ID" | awk '{print $1}')"
GPU_BUS="$(echo "$GPU_ADDRESS" | awk -F':' '{print $1}')"
GPU_FUNCTION="$(echo "$GPU_ADDRESS" | awk -F'.' '{print $2}')"

# Needed for notifications
export DISPLAY=":0",	
export WAYLAND_DISPLAY="wayland-0",	
export XDG_RUNTIME_DIR="/run/user/1000",	
export DBUS_SESSION_BUS_ADDRESS="unix:path=/run/user/1000/bus"

# Function for notifications 
notify(){
	su "$(id -nu 1000)" -c "notify-send \"$1\""
	paplay --server /run/user/1000/pulse/native /usr/share/sounds/freedesktop/stereo/dialog-warning.oga
}

set_bar(){
	# Unbind from amdgpu to set bar sizes
	echo 0000:07:00.0 > /sys/bus/pci/drivers/amdgpu/unbind
	echo 14 > /sys/bus/pci/devices/0000:07:00.0/resource0_resize
	echo 3 > /sys/bus/pci/devices/0000:07:00.0/resource2_resize
}

cpu_pin(){
	systemctl set-property --runtime -- system.slice AllowedCPUs="$1"
	systemctl set-property --runtime -- user.slice AllowedCPUs="$1"
	systemctl set-property --runtime -- init.scope AllowedCPUs="$1"
}

case "$OPERATION" in
	# Before a VM is started
	"prepare")
		if grep -v 'pci' "/etc/libvirt/qemu/${DOMAIN}.xml" | grep -q "bus='0x$GPU_BUS'.*function='0x$GPU_FUNCTION'"; then
			# Prevent VM from starting if GPU is in use
			if fuser -s /dev/dri/card1 || fuser -s /dev/dri/renderD129; then
				notify "GPU is in use"
				# Notify of running programs
				progs="$(sudo fuser /dev/dri/renderD129 2>/dev/null)"
				for prog in $progs; do
					notify "$(ps -p "$prog" -o comm=)"
				done
				exit 1
			fi
			set_bar
		fi
		# Prevent system from sleeping while a VM is active
		systemctl start libvirt-nosleep@"$DOMAIN"
		;;
	"started")
		notify "Domain $DOMAIN started."
		if grep -v 'pci' "/etc/libvirt/qemu/${DOMAIN}.xml" | grep -q "bus='0x$GPU_BUS'.*function='0x$GPU_FUNCTION'"; then
			# Ungrab evdev
			/usr/local/bin/evdev-toggle
			# Pin CPU cores
			cpu_pin "$ALLOWED_CPUS"
		fi
		;;
	"stopped")
		;;
	"release")
		notify "Domain $DOMAIN shut down."
		if grep -v 'pci' "/etc/libvirt/qemu/${DOMAIN}.xml" | grep -q "bus='0x$GPU_BUS'.*function='0x$GPU_FUNCTION'"; then
			# Set monitor back to DP-1
			ddcutil --sn 21470B000061 --brief setvcp 60 x0f
			# Kill looking glass
			pkill looking-glass
			# Unpin CPU cores
			cpu_pin "$ALL_CPUS"
			# Fix sway taking card back
			udevadm trigger --verbose --type=devices --action=remove --subsystem-match=drm --property-match="MINOR=1"
		fi
		# Allow host to sleep
		systemctl stop libvirt-nosleep@"$DOMAIN"
		;;
esac

# Update waybar
pkill -RTMIN+2 waybar
