#!/bin/bash
# create_mdev.sh - Create NVIDIA vGPU mediated device for QEMU
# Usage: ./create_mdev.sh [PCI_ADDRESS] [VGPU_TYPE]

set -euo pipefail

PCI_ADDR="${1:-}"
VGPU_TYPE="${2:-${VGPU_MDEV_TYPE:-nvidia-nvidia-vgpu-turing}}"

# Supported vGPU types (matching the mdev type names the driver registers)
VALID_TYPES=("nvidia-nvidia-vgpu-ampere" "nvidia-nvidia-vgpu-hopper" "nvidia-nvidia-vgpu-turing" "nvidia-nvidia-vgpu-turing-1g" "nvidia-nvidia-vgpu-turing-2g" "nvidia-nvidia-vgpu-pascal")

# Find NVIDIA GPU if not specified
if [[ -z "$PCI_ADDR" ]]; then
    echo "Scanning for NVIDIA GPUs..."
    mapfile -t GPUS < <(lspci -D -d 10de: | grep -E "VGA|3D controller" | awk '{print $1}')
    if [[ ${#GPUS[@]} -eq 0 ]]; then
        echo "ERROR: No NVIDIA GPU found"
        exit 1
    elif [[ ${#GPUS[@]} -eq 1 ]]; then
        PCI_ADDR="${GPUS[0]}"
        echo "Auto-selected GPU: $PCI_ADDR"
    else
        echo "Multiple GPUs found:"
        for i in "${!GPUS[@]}"; do
            echo "  [$i] ${GPUS[$i]}"
        done
        read -p "Select GPU [0-$((${#GPUS[@]}-1))]: " idx
        PCI_ADDR="${GPUS[$idx]}"
    fi
fi

# Validate PCI address format
if [[ ! "$PCI_ADDR" =~ ^[0-9a-fA-F]{4}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9a-fA-F]$ ]]; then
    echo "ERROR: Invalid PCI address format. Expected: dddd:bb:dd.f"
    exit 1
fi

# Check if GPU exists
if [[ ! -e "/sys/bus/pci/devices/$PCI_ADDR" ]]; then
    echo "ERROR: PCI device $PCI_ADDR not found"
    exit 1
fi

# Validate vGPU type
if [[ ! " ${VALID_TYPES[*]} " =~ " ${VGPU_TYPE} " ]]; then
    echo "ERROR: Invalid vGPU type. Valid types: ${VALID_TYPES[*]}"
    exit 1
fi

MDEV_TYPE_PATH="/sys/bus/pci/devices/$PCI_ADDR/mdev_supported_types/$VGPU_TYPE"

if [[ ! -d "$MDEV_TYPE_PATH" ]]; then
    echo "ERROR: vGPU type $VGPU_TYPE not supported on $PCI_ADDR"
    echo "Available types:"
    ls "/sys/bus/pci/devices/$PCI_ADDR/mdev_supported_types/" 2>/dev/null || echo "  (none - check driver loaded)"
    exit 1
fi

# Check available instances
AVAILABLE=$(cat "$MDEV_TYPE_PATH/available_instances" 2>/dev/null || echo "0")
if [[ "$AVAILABLE" -le 0 ]]; then
    echo "ERROR: No available instances for $VGPU_TYPE on $PCI_ADDR"
    exit 1
fi

echo "Creating $VGPU_TYPE on $PCI_ADDR..."

# Generate UUID
UUID=$(uuidgen)
echo "$UUID" | sudo tee "$MDEV_TYPE_PATH/create" > /dev/null

# Verify creation
MDEV_PATH="/sys/bus/mdev/devices/$UUID"
if [[ -d "$MDEV_PATH" ]]; then
    echo "SUCCESS: Created mdev device $UUID"
    echo "  Path: $MDEV_PATH"
    echo "  PCI:  $PCI_ADDR"
    echo "  Type: $VGPU_TYPE"
    echo ""
    echo "For QEMU, use:"
    echo "  -device vfio-pci,sysfsdev=$MDEV_PATH,display=off"
else
    echo "ERROR: Failed to create mdev device"
    exit 1
fi