#!/bin/bash
# Guest Driver Validation Suite
# Tests the guest NVIDIA driver with vGPU

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OPEN_IOV_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

echo "=== Guest Driver Validation Suite ==="
echo "Open-IOV root: $OPEN_IOV_ROOT"

# Check if we're root
if [ "$EUID" -ne 0 ]; then
    echo "ERROR: This script must be run as root"
    exit 1
fi

# Check if mdev exists
MDEV_DEVICES=$(ls /sys/bus/mdev/devices/ 2>/dev/null | wc -l)
if [ "$MDEV_DEVICES" -eq 0 ]; then
    echo "ERROR: No mdev devices found. Run create_mdev.sh first."
    exit 1
fi

echo "Found $MDEV_DEVICES mdev device(s)"

# Launch QEMU with vGPU
echo "Launching QEMU with vGPU..."
"$SCRIPT_DIR/launch_qemu_vgpu.sh" &
QEMU_PID=$!

# Wait for QEMU to start
sleep 10

# Check if QEMU is running
if ! kill -0 $QEMU_PID 2>/dev/null; then
    echo "ERROR: QEMU failed to start"
    exit 1
fi

echo "QEMU started with PID $QEMU_PID"
echo "Guest should be booting..."

# Wait for guest to boot and SSH to be available
echo "Waiting for guest SSH (timeout 120s)..."
for i in {1..24}; do
    if ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null root@localhost -p 2222 "echo 'SSH ready'" 2>/dev/null; then
        echo "Guest SSH is ready"
        break
    fi
    sleep 5
done

# Run guest driver setup
echo "Setting up guest NVIDIA driver..."
ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null root@localhost -p 2222 \
    "bash -s" < "$SCRIPT_DIR/setup_guest_driver.sh"

# Test nvidia-smi
echo "Testing nvidia-smi..."
ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null root@localhost -p 2222 \
    "nvidia-smi"

# Test CUDA
echo "Testing CUDA vectorAdd..."
ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null root@localhost -p 2222 \
    "cd /usr/local/cuda/samples/1_Utilities/vectorAdd && make && ./vectorAdd"

# Cleanup
echo "Shutting down QEMU..."
kill $QEMU_PID 2>/dev/null || true
wait $QEMU_PID 2>/dev/null || true

echo "=== Guest Driver Validation Complete ==="