#!/bin/bash
# Setup script for NVIDIA vGPU mdev device
# This script creates an mdev device for QEMU testing

set -euo pipefail

GPU_PCI_ADDR="${GPU_PCI_ADDR:-0000:01:00.0}"
VGPU_TYPE="${VGPU_TYPE:-nvidia-vgpu-turing}"
MDEV_UUID="${MDEV_UUID:-$(cat /proc/sys/kernel/random/uuid)}"

MDEV_CREATE_PATH="/sys/bus/pci/devices/$GPU_PCI_ADDR/mdev_supported_types/$VGPU_TYPE/create"

echo "=== NVIDIA vGPU mdev Device Setup ==="
echo "GPU PCI: $GPU_PCI_ADDR"
echo "VGPU Type: $VGPU_TYPE"
echo "MDEV UUID: $MDEV_UUID"
echo ""

# Check if GPU device exists
if [[ ! -d "/sys/bus/pci/devices/$GPU_PCI_ADDR" ]]; then
    echo "ERROR: GPU device $GPU_PCI_ADDR not found"
    echo "Available NVIDIA GPUs:"
    lspci -nn -d 10de: | grep -E "VGA|3D"
    exit 1
fi

# Check if mdev_supported_types exists (requires NVIDIA vGPU host driver)
if [[ ! -d "/sys/bus/pci/devices/$GPU_PCI_ADDR/mdev_supported_types" ]]; then
    echo "ERROR: mdev_supported_types not found for $GPU_PCI_ADDR"
    echo ""
    echo "This requires the NVIDIA vGPU host driver (licensed product for data center GPUs)."
    echo "The regular NVIDIA driver does not provide vGPU mdev support."
    echo ""
    echo "For testing on consumer GPUs (TU104, etc.), use dry-run mode:"
    echo "  sudo insmod /path/to/nvidia-vgpu-vfio.ko dry_run=1"
    echo ""
    echo "In dry-run mode, the module loads but cannot create real mdev devices."
    exit 1
fi

# Check if requested vGPU type exists
if [[ ! -d "/sys/bus/pci/devices/$GPU_PCI_ADDR/mdev_supported_types/$VGPU_TYPE" ]]; then
    echo "ERROR: vGPU type $VGPU_TYPE not supported"
    echo "Available types:"
    ls "/sys/bus/pci/devices/$GPU_PCI_ADDR/mdev_supported_types/"
    exit 1
fi

# Create mdev device
echo "Creating mdev device..."
echo "$MDEV_UUID" | sudo tee "$MDEV_CREATE_PATH" > /dev/null

if [[ -d "/sys/bus/mdev/devices/$MDEV_UUID" ]]; then
    echo "SUCCESS: mdev device created at /sys/bus/mdev/devices/$MDEV_UUID"
    echo ""
    echo "Device info:"
    cat "/sys/bus/mdev/devices/$MDEV_UUID/description" 2>/dev/null || echo "  (no description)"
    echo ""
    echo "To use with QEMU:"
    echo "  export MDEV_UUID=$MDEV_UUID"
    echo "  /path/to/launch_vgpu_qemu.sh"
else
    echo "ERROR: Failed to create mdev device"
    exit 1
fi