#!/bin/bash

# Log file for debugging
LOG_FILE="/var/log/libvirt/gpu-reset.log"

# Function to log messages
log() {
    echo "$(date): $1" >> "$LOG_FILE"
}

# Create log directory if it doesn't exist
mkdir -p "$(dirname "$LOG_FILE")" 2>/dev/null
touch "$LOG_FILE" 2>/dev/null

VM_NAME="$1"
HOOK_TYPE="$2"
STATE="$3"
REASON="$4"

# Only proceed for the release/end event
if [[ "$HOOK_TYPE" == "release" && "$STATE" == "end" ]]; then
    log "Processing release/end for VM: $VM_NAME"
    
    # Read VM XML from stdin and save it for processing
    VM_XML=$(cat)
    
    # Log a snippet of the XML for debugging
    log "XML snippet: $(echo "$VM_XML" | grep -n "address domain" | head -3)"
    
    # Extract all PCI addresses from the VM XML using regex (for single quotes)
    PCI_ADDRESSES=$(echo "$VM_XML" | grep -oP "address domain='0x([0-9a-f]+)' bus='0x([0-9a-f]+)' slot='0x([0-9a-f]+)' function='0x([0-9a-f]+)" | 
                  sed -E "s/address domain='0x([0-9a-f]+)' bus='0x([0-9a-f]+)' slot='0x([0-9a-f]+)' function='0x([0-9a-f]+)/\1:\2:\3.\4/")

    # try double quotes if single quotes didn't work
    if [ -z "$PCI_ADDRESSES" ]; then
        PCI_ADDRESSES=$(echo "$VM_XML" | grep -oP 'address domain="0x([0-9a-f]+)" bus="0x([0-9a-f]+)" slot="0x([0-9a-f]+)" function="0x([0-9a-f]+)' | 
                      sed -E 's/address domain="0x([0-9a-f]+)" bus="0x([0-9a-f]+)" slot="0x([0-9a-f]+)" function="0x([0-9a-f]+)/\1:\2:\3.\4/')
    fi

    log "Extracted PCI addresses: $PCI_ADDRESSES"
    
    # Process each PCI address
    while IFS= read -r PCI_ADDR; do

        log "Processing PCI device: $PCI_ADDR"
        
        # Check if device exists
        if [ ! -e "/sys/bus/pci/devices/$PCI_ADDR" ]; then
            log "Device $PCI_ADDR does not exist, skipping"
            continue
        fi
        
        # Get vendor and device ID
        VENDOR=$(cat "/sys/bus/pci/devices/$PCI_ADDR/vendor" 2>/dev/null)
        DEVICE=$(cat "/sys/bus/pci/devices/$PCI_ADDR/device" 2>/dev/null)
        CLASS=$(cat "/sys/bus/pci/devices/$PCI_ADDR/class" 2>/dev/null)

        log "Device $PCI_ADDR: Vendor=$VENDOR, Device=$DEVICE, Class=$CLASS"
        
        # Check if this is an NVIDIA GPU
        if echo "$VENDOR" | grep -q "0x10de" && echo "$CLASS" | grep -q "0x030000"; then            
            # Attempt to unbind from vfio-pci
            if [ -e "/sys/bus/pci/drivers/vfio-pci/$PCI_ADDR" ]; then
                log "Unbinding $PCI_ADDR from vfio-pci"
                echo "$PCI_ADDR" > /sys/bus/pci/drivers/vfio-pci/unbind
                sleep 1
            fi
            
            # Check if nvidia driver is available
            if [ -d "/sys/bus/pci/drivers/nvidia" ]; then
                log "Rebinding $PCI_ADDR to nvidia driver"
                echo "nvidia" > /sys/bus/pci/devices/$PCI_ADDR/driver_override
                echo "$PCI_ADDR" > /sys/bus/pci/drivers/nvidia/bind
            else
                log "Nvidia driver not available for rebinding"
            fi
        fi
    done <<< "$PCI_ADDRESSES"
fi

exit 0