#!/bin/sh
set -eu

cleanup() {
    echo ""
    if [ "${FAIL}" = "1" ]; then
        echo "Test failed"
        exit 1
    fi

    echo "Test passed"
    exit 0
}

FAIL=1
trap cleanup EXIT HUP INT TERM

# Install test dependencies
apt-get remove --purge cloud-init --yes
apt-get install --yes curl

# Install Incus
curl -sL https://pkgs.zabbly.com/get/incus-daily | sh

# Configure Incus
incus admin init --auto

if [ "${1}" = "nvidia" ]; then
    # Reduce power draw
    nvidia-smi -q -d POWER

    # Enable SR-IOV
    /usr/lib/nvidia/sriov-manage -e ALL

    # Confirm GPU is online
    nvidia-smi

    # Disable MIG if enabled
    nvidia-smi -mig 0
fi

if [ "${1}" = "intel" ]; then
    # Configure SR-IOV on Intel Xe
    echo 0 > /sys/devices/pci0000:00/0000:00:02.0/sriov_drivers_autoprobe
    echo 7 > /sys/devices/pci0000:00/0000:00:02.0/sriov_numvfs
fi

# Incus resource API
incus info --resources

# Launch test instances
for i in $(seq 1 4); do
    incus init images:ubuntu/24.04/cloud "v${i}" --vm -c security.secureboot=false < /dev/null
    if [ "${1}" = "nvidia" ]; then
        incus config device add "v${i}" vgpu gpu gputype=mdev pci=0000:07:00.0 mdev=nvidia-468
    elif [ "${1}" = "intel" ]; then
        incus config device add "v${i}" vgpu gpu gputype=sriov pci=0000:00:02.0
    fi
    incus start "v${i}"
done

# Wait for them to start and list
sleep 30
incus list

if [ "${1}" = "nvidia" ]; then
    # Validate NVIDIA vGPU
    incus exec v4 -- apt-get update
    incus exec v4 -- apt-get dist-upgrade --yes
    incus exec v4 -- apt-get install --no-install-recommends --yes build-essential wget pciutils linux-headers-virtual
    incus restart v4
    sleep 30
    incus exec v4 -- wget -6 http://lab.linuxcontainers.org/nvidia/v570/nvidia-guest.deb
    incus exec v4 -- apt-get install --yes /root/nvidia-guest.deb
    incus exec v4 -- nvidia-smi
fi

if [ "${1}" = "intel" ]; then
    # Validate Intel vGPU
    incus exec v3 -- apt-get update
    incus exec v3 -- apt install --yes pciutils
    incus exec v3 -- lspci | grep VGA.*Intel
fi

FAIL=0
