#!/bin/bash
# test_vgpu_stack.sh - Comprehensive vGPU stack test

set -euo pipefail

BASE_DIR="$OPEN_IOV_ROOT"
LIB_DIR="$BASE_DIR/userspace/lib"
EXAMPLE="$BASE_DIR/userspace/examples/example_vgpu"

echo "=========================================="
echo "vGPU Stack Test Suite"
echo "=========================================="

# Test 1: Module loading
echo "Test 1: Module loading"
for mod in nvidia open_iov_vfio gsp_vgpu; do
    if lsmod | grep -q "^$mod "; then
        echo "  PASS: $mod loaded"
    else
        echo "  FAIL: $mod not loaded"
        exit 1
    fi
done

# Test 2: Device nodes
echo "Test 2: Device nodes"
if [[ -c /dev/gsp-vgpu ]]; then
    echo "  PASS: /dev/gsp-vgpu exists"
else
    echo "  FAIL: /dev/gsp-vgpu missing"
    exit 1
fi

# Test 3: Library loading
echo "Test 3: Library loading"
if LD_LIBRARY_PATH="$LIB_DIR" ldd "$EXAMPLE" 2>&1 | grep -q "libiov.so"; then
    echo "  PASS: libiov.so loads correctly"
else
    echo "  FAIL: libiov.so load failed"
    exit 1
fi

# Test 4: Example app help
echo "Test 4: Example app help"
if LD_LIBRARY_PATH="$LIB_DIR" "$EXAMPLE" --help 2>&1 | grep -q "Usage"; then
    echo "  PASS: example_vgpu runs"
else
    echo "  FAIL: example_vgpu failed"
    exit 1
fi

# Test 5: mdev creation (if GPU available)
echo "Test 5: mdev creation"
GPU_PCI=$(lspci -D -d 10de: | grep -E "VGA|3D controller" | head -1 | awk '{print $1}')
if [[ -n "$GPU_PCI" ]]; then
    echo "  Found GPU: $GPU_PCI"
    MDEV_TYPE_PATH="/sys/bus/pci/devices/$GPU_PCI/mdev_supported_types/nvidia-vgpu-turing"
    if [[ -d "$MDEV_TYPE_PATH" ]]; then
        AVAIL=$(cat "$MDEV_TYPE_PATH/available_instances" 2>/dev/null || echo 0)
        if [[ "$AVAIL" -gt 0 ]]; then
            UUID=$(uuidgen)
            echo "$UUID" | sudo tee "$MDEV_TYPE_PATH/create" > /dev/null
            sleep 1
            if [[ -d "/sys/bus/mdev/devices/$UUID" ]]; then
                echo "  PASS: mdev created: $UUID"
                # Cleanup
                echo 1 | sudo tee "/sys/bus/mdev/devices/$UUID/remove" > /dev/null
            else
                echo "  FAIL: mdev creation failed"
                exit 1
            fi
        else
            echo "  SKIP: No available instances"
        fi
    else
        echo "  SKIP: nvidia-vgpu-turing type not available"
    fi
else
    echo "  SKIP: No NVIDIA GPU found"
fi

# Test 6: Kernel log check
echo "Test 6: Kernel log check"
if dmesg -T | grep -q "GSP vGPU-IO driver loaded"; then
    echo "  PASS: gsp-vgpu driver loaded successfully"
else
    echo "  WARN: gsp-vgpu load message not found in dmesg"
fi

if dmesg -T | grep -q "NVIDIA vGPU VFIO driver loaded"; then
    echo "  PASS: open-iov-vfio driver loaded successfully"
else
    echo "  WARN: open-iov-vfio load message not found in dmesg"
fi

# Test 7: Firmware verification
echo "Test 7: Firmware verification"
FW="$BASE_DIR/firmware/570-GSP/gsp-570.144.fwimage.bin"
if [[ -f "$FW" ]]; then
    SIZE=$(stat -c%s "$FW")
    EXPECTED=29622272  # ~28.5MB
    if [[ "$SIZE" -eq "$EXPECTED" ]]; then
        echo "  PASS: Firmware size correct ($SIZE bytes)"
    else
        echo "  WARN: Firmware size mismatch: $SIZE (expected $EXPECTED)"
    fi
    
    # Check ELF header
    if file "$FW" | grep -q "ELF 64-bit LSB executable, UCB RISC-V"; then
        echo "  PASS: Firmware is valid RISC-V ELF"
    else
        echo "  FAIL: Firmware is not valid RISC-V ELF"
        exit 1
    fi
else
    echo "  FAIL: Firmware not found"
    exit 1
fi

echo ""
echo "=========================================="
echo "All Tests Passed!"
echo "=========================================="
echo ""
echo "Stack is ready for QEMU guest launch."
echo "Run: $BASE_DIR/scripts/create_mdev.sh"
echo "Then: $BASE_DIR/scripts/launch_qemu_vgpu.sh"