#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# Load vGPU stack in converged mode
# Usage: sudo ./load_converged.sh [dry_run]

set -euo pipefail

OPEN_IOV_ROOT="${OPEN_IOV_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}"
NVIDIA_OPEN_KERNEL="$OPEN_IOV_ROOT/nvidia-open-gpu-kernel-modules/kernel-open"
VGPU_KERNEL="$OPEN_IOV_ROOT/kernel"

DRY_RUN="${1:-false}"

echo "=== Open-IOV Converged Mode Module Load ==="
echo "Open-IOV Root: $OPEN_IOV_ROOT"
echo "Dry-run mode: $DRY_RUN"
echo ""

# Function to check if module is loaded
is_loaded() {
    lsmod | grep -q "^$1 "
}

# Function to load module with error checking
load_module() {
    local module_path=$1
    local module_name=$(basename $module_path .ko)
    local params="${2:-}"
    
    if is_loaded "$module_name"; then
        echo "Module $module_name already loaded"
        return 0
    fi
    
    echo "Loading $module_name..."
    if sudo insmod "$module_path" $params; then
        echo "  ✓ $module_name loaded successfully"
        return 0
    else
        echo "  ✗ Failed to load $module_name"
        return 1
    fi
}

# 1. Load mdev core
echo "Step 1: Loading mdev core module..."
sudo modprobe mdev
echo "  ✓ mdev loaded"
echo ""

# 2. Load host nvidia.ko with vGPU support
echo "Step 2: Loading host nvidia.ko (patched open-gpu-kernel-modules)..."
load_module "$NVIDIA_OPEN_KERNEL/nvidia.ko" "NVreg_RegisterPCIDriver=1 NVreg_RmMsg=:"
echo ""

# 3. Load dependent modules
echo "Step 3: Loading dependent modules..."
load_module "$NVIDIA_OPEN_KERNEL/nvidia-modeset.ko"
load_module "$NVIDIA_OPEN_KERNEL/nvidia-uvm.ko"
load_module "$NVIDIA_OPEN_KERNEL/nvidia-drm.ko"
load_module "$NVIDIA_OPEN_KERNEL/nvidia-peermem.ko"
echo ""

# 4. Load nvidia-vgpu-vfio.ko
echo "Step 4: Loading nvidia-vgpu-vfio.ko..."
if [ "$DRY_RUN" = "dry_run" ] || [ "$DRY_RUN" = "true" ]; then
    echo "  Running in DRY-RUN mode first..."
    load_module "$VGPU_KERNEL/nvidia-vgpu-vfio.ko" "dyndbg='+p' dry_run=1"
    echo "  Dry-run successful, unloading..."
    sudo rmmod nvidia-vgpu-vfio
    echo "  Loading for real..."
    load_module "$VGPU_KERNEL/nvidia-vgpu-vfio.ko" "dyndbg='+p'"
else
    load_module "$VGPU_KERNEL/nvidia-vgpu-vfio.ko" "dyndbg='+p'"
fi
echo ""

# 5. Verify
echo "Step 5: Verifying vGPU registration..."
sleep 2
if [ -d "/sys/bus/pci/devices/0000:01:00.0/mdev_supported_types" ]; then
    echo "  ✓ mdev_supported_types directory exists"
    ls -la /sys/bus/pci/devices/0000:01:00.0/mdev_supported_types/
else
    echo "  ✗ mdev_supported_types NOT found"
    echo "  Check dmesg for errors: dmesg -T | grep -i vgpu"
    exit 1
fi

echo ""
echo "=== Converged Mode Load Complete ==="
echo ""
echo "Next steps:"
echo "  - Create mdev device: echo <uuid> | sudo tee /sys/bus/pci/devices/0000:01:00.0/mdev_supported_types/nvidia-vgpu-turing/create"
echo "  - Run QEMU with vGPU passthrough"
echo "  - Run tests: $OPEN_IOV_ROOT/tools/test/run_integration_tests.sh"