#!/bin/bash

# GPU Passthrough Setup for AMD Strix Halo iGPU to Windows Container
# This script enables IOMMU and configures VFIO-PCI for GPU passthrough

set -e

echo "=== AMD Strix Halo GPU Passthrough Setup ==="
echo ""

# Check if running as root
if [ "$EUID" -ne 0 ]; then 
    echo "Please run as root (sudo)"
    exit 1
fi

# GPU PCI IDs
GPU_PCI="0000:c5:00.0"
GPU_AUDIO_PCI="0000:c5:00.1"
GPU_VENDOR_ID="1002"
GPU_DEVICE_ID="1586"

# Check current IOMMU status
echo "Current IOMMU status:"
cat /proc/cmdline | grep -o "amd_iommu=[^ ]*" || echo "IOMMU not configured in kernel parameters"
echo ""

# Step 1: Enable IOMMU in GRUB
echo "Step 1: Checking GRUB configuration..."
GRUB_FILE="/etc/default/grub"

if ! grep -q "amd_iommu=on" "$GRUB_FILE"; then
    echo "Adding IOMMU parameters to GRUB..."
    cp "$GRUB_FILE" "$GRUB_FILE.backup.$(date +%Y%m%d-%H%M%S)"
    
    # Update GRUB_CMDLINE_LINUX_DEFAULT
    sed -i 's/GRUB_CMDLINE_LINUX_DEFAULT="\(.*\)"/GRUB_CMDLINE_LINUX_DEFAULT="\1 amd_iommu=on iommu=pt vfio-pci.ids=1002:1586,1002:1640"/' "$GRUB_FILE"
    
    # Clean up double spaces
    sed -i 's/GRUB_CMDLINE_LINUX_DEFAULT="  /GRUB_CMDLINE_LINUX_DEFAULT="/' "$GRUB_FILE"
    
    echo "Updating GRUB..."
    update-grub
    
    echo ""
    echo "✓ GRUB updated. IOMMU will be enabled after reboot."
else
    echo "✓ IOMMU already configured in GRUB"
fi

# Step 2: Configure VFIO modules
echo ""
echo "Step 2: Configuring VFIO modules..."
VFIO_CONF="/etc/modprobe.d/vfio.conf"

cat > "$VFIO_CONF" << EOF
# Bind AMD Strix Halo GPU to VFIO for passthrough
options vfio-pci ids=1002:1586,1002:1640
softdep amdgpu pre: vfio-pci
softdep snd_hda_intel pre: vfio-pci
EOF

echo "✓ VFIO configuration created"

# Step 3: Update initramfs
echo ""
echo "Step 3: Adding VFIO modules to initramfs..."
MODULES_FILE="/etc/initramfs-tools/modules"

if ! grep -q "vfio_pci" "$MODULES_FILE"; then
    cat >> "$MODULES_FILE" << EOF

# VFIO modules for GPU passthrough
vfio
vfio_iommu_type1
vfio_pci
vfio_virqfd
EOF
    echo "✓ VFIO modules added to initramfs"
else
    echo "✓ VFIO modules already in initramfs"
fi

echo ""
echo "Updating initramfs..."
update-initramfs -u

# Step 4: Summary and next steps
echo ""
echo "=== Setup Complete ==="
echo ""
echo "Configuration applied:"
echo "  - IOMMU enabled in GRUB (amd_iommu=on iommu=pt)"
echo "  - GPU (1002:1586) bound to vfio-pci driver"
echo "  - Audio (1002:1640) bound to vfio-pci driver"
echo "  - Initramfs updated with VFIO modules"
echo ""
echo "⚠️  REBOOT REQUIRED to apply changes by running: 'sudo update-grub' and 'sudo reboot'"
echo ""
echo "After reboot:"
echo "  1. Verify IOMMU is enabled: dmesg | grep -i iommu"
echo "  2. Check GPU binding: lspci -nnk -d 1002:1586"
echo "  3. Start Windows container: docker-compose up -d"
echo "  4. Install AMD drivers in Windows"
echo ""
echo "To reboot now, run: reboot"