#!/bin/bash

# Unbind the AMD GPU from VFIO and restore it to the host (amdgpu driver)
# Run this AFTER stopping the Windows container to restore host GPU access

set -e

if [ "$EUID" -ne 0 ]; then 
    echo "Please run as root: sudo $0"
    exit 1
fi

echo "=== Unbinding AMD GPU from VFIO ==="
echo ""

GPU_PCI="0000:c5:00.0"
AUDIO_PCI="0000:c5:00.1"

# Stop Windows container first
echo "Checking if Windows container is running..."
if docker ps | grep -q windows2; then
    echo "Stopping Windows container..."
    docker stop windows2
    echo "✓ Container stopped"
fi

# Unbind from vfio-pci
echo ""
echo "Unbinding from vfio-pci..."
if [ -e /sys/bus/pci/devices/$GPU_PCI/driver ]; then
    echo "$GPU_PCI" > /sys/bus/pci/devices/$GPU_PCI/driver/unbind
    echo "✓ GPU unbound from vfio-pci"
fi

if [ -e /sys/bus/pci/devices/$AUDIO_PCI/driver ]; then
    echo "$AUDIO_PCI" > /sys/bus/pci/devices/$AUDIO_PCI/driver/unbind
    echo "✓ Audio unbound from vfio-pci"
fi

# Remove device IDs from vfio-pci
echo "1002 1586" > /sys/bus/pci/drivers/vfio-pci/remove_id 2>/dev/null || true
echo "1002 1640" > /sys/bus/pci/drivers/vfio-pci/remove_id 2>/dev/null || true

sleep 1

# Rebind to host drivers
echo ""
echo "Binding back to host drivers..."
echo "$GPU_PCI" > /sys/bus/pci/drivers_probe
echo "$AUDIO_PCI" > /sys/bus/pci/drivers_probe

sleep 2

# Verify
GPU_DRIVER=$(lspci -nnk -s c5:00.0 | grep "Kernel driver in use" | awk '{print $5}')
AUDIO_DRIVER=$(lspci -nnk -s c5:00.1 | grep "Kernel driver in use" | awk '{print $5}')

echo ""
echo "=== Status ==="
if [ "$GPU_DRIVER" = "amdgpu" ]; then
    echo "✓ GPU restored to amdgpu"
else
    echo "⚠ GPU bound to: ${GPU_DRIVER:-none}"
fi

if [ "$AUDIO_DRIVER" = "snd_hda_intel" ]; then
    echo "✓ Audio restored to snd_hda_intel"
else
    echo "⚠ Audio bound to: ${AUDIO_DRIVER:-none}"
fi

echo ""
if [ "$GPU_DRIVER" = "amdgpu" ]; then
    echo "✓ GPU restored to host!"
    echo ""
    echo "You may need to restart your display manager:"
    echo "  sudo systemctl restart gdm3    # for GNOME"
    echo "  sudo systemctl restart lightdm  # for XFCE/other"
else
    echo "⚠ GPU not fully restored. You may need to reboot."
fi