#!/bin/bash -e
# Helper for creating QEMU images with GPU pass-through support
#
# Author: Christian Kastner <ckk@kvr.at>
# License: MIT

# Option defaults
authorized_keys_file=
mirror="http://deb.debian.org/debian"
release="unstable"
image_size=20G
stack=

function usage() {
    cat >&2 <<-EOF

	Create a QEMU VM image for suitable for GPU pass-through.

	STACK is required. Currently support values are: [rocm, cuda]

	IMAGE, if not specified, defaults to STACK-RELEASE.img.

	This script will build an EFI-bootable image IMAGE, with packages
	nvidia-kernel-dkms and openssh-server preinstalled, and (optionally) other
	minor customizations.

	The image will have a regular (non-system) user 'user', who will be in all
	the necessary system groups.

	When using gpuisol-qemu-run with this image, you will be able to share a host
	directory with the guest.

	Synopsis:
	  $0 -h

	  $0 [-a AKFILE] [-m MIRROR] [-r RELEASE] STACK [IMAGE]

	Options:
	  -h          Show this help
	  -a AKFILE   File to copy to /root/.ssh/authorized_keys and to
	              /home/rocm/.ssh/authorized_keys in the guest. 
	  -m MIRROR   Download packages from here. This will be also used within
	              the image. Ideally, this points to an APT cache on the host
	              machine, such as apt-cacher-ng or approx.
	  -s SIZE     Image size (default: '$image_size')
	  -r RELEASE  Release (default: '$release')
	              If RELEASE is 'experimental', APT sources for both 'unstable'
	              and 'experimental' will be added to the image.
	              If RELEASE contains a dash, APT sources for the "basename"
	              will be added, e.g.: 'trixie-backports' will include
	              sources for 'trixie' and 'trixie-backports'.

	Examples:

	  # Configure the system for GPU pass-through with AMD GPUs

	  \$ gpuisol-qemu-setup -u <user> rocm

	  # Creates rocm-unstable.img based on unstable

	  \$ $0 rocm

	  # Same as above, but name the output unstable-autopkgtest-amd64.img

	  \$ $0 rocm unstable-autopkgtest-amd64.img

	  # Create rocm-unstable.img using the mirror at 10.1.2.3:9999, which
	  # ideally is a local APT cache

	  \$ $0 -m http://10.1.2.3:9999/debian rocm

	  # Creates rocm-trixie.img

	  \$ $0 -r trixie rocm

	  # Copies this file to {/root,/home/rocm}/.ssh/authorized_keys

	  \$ $0 -a ~/.ssh/id_rsa.pub rocm

	EOF
}

while getopts "a:hm:r:s:" OPTNAME; do
    case "$OPTNAME" in
    a) authorized_keys_file="$OPTARG" ;;
    h)
        usage
        exit 0
        ;;
    m) mirror="$OPTARG" ;;
    r) release="$OPTARG" ;;
    s) image_size="$OPTARG" ;;
    ?)
        usage
        exit 1
        ;;
    esac
done
shift $((OPTIND - 1))

if [ "$#" -eq 0 ]; then
    echo "Missing GPU stack argument: [rocm, cuda]" >&2
    exit 1
fi
stack="$1"
if [ "$stack" != "rocm" ] && [ "$stack" != "cuda" ]; then
    echo "Unsupported GPU stack argument value: $stack" >&2
    exit 1
fi

imagename="${2:-$stack-$release}"

# trixie             -> trixie
# trixie-backports   -> trixie
# experimental       -> unstable
release_base=$(echo "$release" | sed -r 's/-.+$//')
[ "$release_base" != "experimental" ] || release_base="unstable"

script="/usr/share/debootstrap/scripts/$release_base"
if grep -q archive.ubuntu.com "$script"; then
    distro="ubuntu"
else
    distro="debian"
fi

# /etc/apt/sources.list within the VM
if [ "$distro" = "ubuntu" ]; then
    AUTOPKGTEST_APT_SOURCES=$(
        cat <<-EOF
		deb     $mirror $release          main restricted universe multiverse
		deb-src $mirror $release          main restricted universe multiverse
		deb     $mirror $release-updates  main restricted universe multiverse
		deb-src $mirror $release-updates  main restricted universe multiverse
		deb     $mirror $release-security main restricted universe multiverse
		deb-src $mirror $release-security main restricted universe multiverse
EOF
    )
else
    components="main contrib non-free-firmware"
    [ "$stack" != "cuda" ] || components="$components non-free"

    if [ "$release" != "$release_base" ]; then
        AUTOPKGTEST_APT_SOURCES=$(
            cat <<-EOF
		deb     $mirror $release_base $components
		deb-src $mirror $release_base $components

	deb     $mirror $release $components
	deb-src $mirror $release $components
EOF
        )
    else
        AUTOPKGTEST_APT_SOURCES=$(
            cat <<-EOF
		deb     $mirror $release $components
		deb-src $mirror $release $components
EOF
        )
    fi
fi
export AUTOPKGTEST_APT_SOURCES

scratch_dir="$(mktemp -d || exit 1)"
cleanup() {
    rm -rf "$scratch_dir"
}
trap cleanup EXIT

# Image customization script
#
# Unfortunately the modscript cannot be stuffed in a variable; it must be an
# executable file
modscript_file="$scratch_dir/modscript"
touch "$modscript_file"
chmod +x "$modscript_file"
if [ "$stack" = "rocm" ]; then
    extra_packages="firmware-amd-graphics"
elif [ "$stack" = "cuda" ]; then
    extra_packages="linux-headers-amd64 nvidia-kernel-dkms nvidia-driver-bin nvidia-driver-libs libcuda1"
fi

# The script generated below will be executed by autopkgtest-build-qemu. The
# one and only argument to this script is the path to where the guest's root fs
# is temporarily mounted on the host.
cat >"$modscript_file" <<EOF
# Workaround for https://bugs.debian.org/1111870
if [ -z "\$1" ]; then
    VMROOT="/"
else
    VMROOT="\$1"
fi

export DEBIAN_FRONTEND=noninteractive

# This user is created by autopkgtest-build-qemu
chroot "\$VMROOT" usermod -a -G audio,render,video user

chroot "\$VMROOT" apt-get install --quiet -y --no-install-recommends \
	openssh-server $extra_packages pciutils

if [ -n "$authorized_keys_file" ]; then
    mkdir --mode=0700 "\$VMROOT/root/.ssh"
    install -m 0600 "$authorized_keys_file" "\$VMROOT/root/.ssh/authorized_keys"

    mkdir --mode=0700 "\$VMROOT/home/user/.ssh"
    install -m 0600 "$authorized_keys_file" "\$VMROOT/home/user/.ssh/authorized_keys"
    chroot "\$VMROOT" chown -R "user": "/home/user/.ssh"
fi

# Mount point for a shared folder, if the VM is started with one
/bin/echo "9p\n9pnet\n9pnet_virtio" >> "\$VMROOT/etc/initramfs-tools/modules"
mkdir -m 755 "\$VMROOT/shared"
echo "qemu-host-dir /shared 9p trans=virtio,version=9p2000.L,auto,nofail 0 0" >> "\$VMROOT/etc/fstab"

echo "Setting hostname"
echo "$release-guest" > "\$VMROOT/etc/hosts"
echo "$release-guest" > "\$VMROOT/etc/hostname"
EOF

# Finally: build the image
# This invocation will create files inside the chroot that are expected to be
# world-readable, e.g. /etc/apt/sources.list
umask 0022
mmdebstrap-autopkgtest-build-qemu \
    --boot=efi \
    --architecture=amd64 \
    --mirror "$mirror" \
    --size="$image_size" \
    --script "$modscript_file" \
    "$release_base" "$scratch_dir/$imagename.raw"
qemu-img convert -f raw -O qcow2 "$scratch_dir/$imagename.raw" "$imagename.img"
