#!/bin/bash
# launch_qemu_vgpu.sh - Launch QEMU with NVIDIA vGPU mediated device
# Usage: ./launch_qemu_vgpu.sh [install|run] [MDEV_UUID] [GUEST_IMAGE] [OPTIONS]

set -euo pipefail

MODE="${1:-run}"
if [[ "$MODE" == "install" ]]; then
    shift
elif [[ "$MODE" == "run" ]]; then
    shift
else
    # Default to run if first arg doesn't match
    MODE="run"
fi

MDEV_UUID="${1:-}"
GUEST_IMAGE="${2:-/var/lib/libvirt/images/ubuntu-22.04.qcow2}"
shift 2 || true
EXTRA_QEMU_ARGS=("$@")

ISO_IMAGE="/var/lib/libvirt/images/ubuntu-24.04.5-live-server-amd64.iso"

# Find mdev device if UUID not specified
if [[ -z "$MDEV_UUID" ]]; then
    echo "Scanning for available mdev devices..."
    mapfile -t MDEVS < <(find /sys/bus/mdev/devices -maxdepth 1 -name "????????-????-????-????-????????????" -printf "%f\n" 2>/dev/null)
    if [[ ${#MDEVS[@]} -eq 0 ]]; then
        echo "ERROR: No mdev devices found. Run create_mdev.sh first."
        exit 1
    elif [[ ${#MDEVS[@]} -eq 1 ]]; then
        MDEV_UUID="${MDEVS[0]}"
        echo "Auto-selected mdev: $MDEV_UUID"
    else
        echo "Multiple mdev devices found:"
        for i in "${!MDEVS[@]}"; do
            MDEV="${MDEVS[$i]}"
            TYPE=$(cat "/sys/bus/mdev/devices/$MDEV/mdev_type/name" 2>/dev/null || echo "unknown")
            echo "  [$i] $MDEV ($TYPE)"
        done
        read -p "Select mdev [0-$((${#MDEVS[@]}-1))]: " idx
        MDEV_UUID="${MDEVS[$idx]}"
    fi
fi

MDEV_PATH="/sys/bus/mdev/devices/$MDEV_UUID"

if [[ ! -d "$MDEV_PATH" ]]; then
    echo "ERROR: mdev device $MDEV_UUID not found at $MDEV_PATH"
    exit 1
fi

# Verify guest image
if [[ ! -f "$GUEST_IMAGE" ]]; then
    echo "ERROR: Guest image not found: $GUEST_IMAGE"
    echo "Create one with: qemu-img create -f qcow2 $GUEST_IMAGE 50G"
    exit 1
fi

# Check KVM
if [[ ! -r /dev/kvm ]]; then
    echo "ERROR: KVM not available (/dev/kvm not readable)"
    exit 1
fi

# Check VFIO
if [[ ! -e /dev/vfio/vfio ]]; then
    echo "ERROR: VFIO not available"
    exit 1
fi

# Get mdev info
VGPU_TYPE=$(cat "$MDEV_PATH/mdev_type/name" 2>/dev/null || echo "unknown")
PCI_ADDR=$(readlink -f "$MDEV_PATH/../../.." | xargs basename)

echo "=========================================="
echo "Launching QEMU with NVIDIA vGPU ($MODE mode)"
echo "=========================================="
echo "mdev UUID:  $MDEV_UUID"
echo "vGPU Type:  $VGPU_TYPE"
echo "PCI Addr:   $PCI_ADDR"
echo "Guest Image: $GUEST_IMAGE"
echo "Extra Args:  ${EXTRA_QEMU_ARGS[*]}"
echo "=========================================="

# CPU configuration for vGPU
CPU_OPTS=(
    -cpu host,hv_time,hv_relaxed,hv_vapic,hv_spinlocks=0x1fff,hv_vendor_id=NVIDIA42
    -smp 8,sockets=1,cores=8,threads=1
    -m 16G,slots=4,maxmem=64G
    -mem-prealloc
)

# Memory backing for hugepages (better performance)
HUGEPAGE_OPTS=(
    -object memory-backend-file,id=mem,size=16G,mem-path=/dev/hugepages,share=on
    -numa node,memdev=mem
)

# VFIO device configuration
VFIO_OPTS=(
    -device vfio-pci,sysfsdev="$MDEV_PATH",display=off
)

# Display - use qxl-vga with GTK for local visible window
DISPLAY_OPTS=(
    -device qxl-vga
    -display gtk
)

# Network
NET_OPTS=(
    -netdev user,id=net0,hostfwd=tcp::2222-:22
    -device virtio-net-pci,netdev=net0
)

# Storage
if [[ "$MODE" == "install" ]]; then
    if [[ ! -f "$ISO_IMAGE" ]]; then
        echo "ERROR: ISO image not found: $ISO_IMAGE"
        exit 1
    fi
    STORAGE_OPTS=(
        -drive file="$GUEST_IMAGE",format=qcow2,if=virtio,cache=writeback,aio=threads
        -drive file="$ISO_IMAGE",format=raw,if=virtio,media=cdrom
    )
    echo "Installation mode: ISO attached ($ISO_IMAGE)"
else
    STORAGE_OPTS=(
        -drive file="$GUEST_IMAGE",format=qcow2,if=virtio,cache=writeback,aio=threads
    )
    echo "Run mode: No ISO attached"
fi

# USB tablet for better mouse integration
USB_OPTS=(
    -usb
    -device usb-tablet
)

# RNG
RNG_OPTS=(
    -object rng-random,id=rng0,filename=/dev/urandom
    -device virtio-rng-pci,rng=rng0
)

# Serial log for debugging
SERIAL_OPTS=(
    -serial file:/tmp/qemu-serial.log
)

# QEMU command
QEMU_CMD=(
    qemu-system-x86_64
    -enable-kvm
    -machine q35,accel=kvm,kernel_irqchip=split
    -bios /usr/share/ovmf/OVMF.fd
    "${CPU_OPTS[@]}"
    "${HUGEPAGE_OPTS[@]}"
    "${VFIO_OPTS[@]}"
    "${DISPLAY_OPTS[@]}"
    "${NET_OPTS[@]}"
    "${STORAGE_OPTS[@]}"
    "${USB_OPTS[@]}"
    "${RNG_OPTS[@]}"
    "${SERIAL_OPTS[@]}"
    "${EXTRA_QEMU_ARGS[@]}"
    -name "vGPU-Guest-$VGPU_TYPE"
    -pidfile "/tmp/qemu-vgpu-$MDEV_UUID.pid"
    -daemonize
)

echo "Starting QEMU..."
echo "Command: ${QEMU_CMD[*]}"
echo ""

"${QEMU_CMD[@]}"

sleep 3

if [[ -f "/tmp/qemu-vgpu-$MDEV_UUID.pid" ]]; then
    PID=$(cat "/tmp/qemu-vgpu-$MDEV_UUID.pid")
    if kill -0 "$PID" 2>/dev/null; then
        echo "SUCCESS: QEMU started with PID $PID"
        echo ""
        echo "GTK window should be visible on your display"
        echo "Connect via SSH: ssh -p 2222 user@localhost"
        echo ""
        echo "To stop: kill $PID"
        echo "Serial log: /tmp/qemu-serial.log"
    else
        echo "ERROR: QEMU failed to start"
        exit 1
    fi
else
    echo "ERROR: QEMU PID file not created"
    exit 1
fi