#!/bin/bash
# Helper for creating autopkgtest-capable podman images with GPU support
#
# Author: Christian Kastner <ckk@kvr.at>
# License: MIT
set -eu

# Option defaults
mirror="http://deb.debian.org/debian"
release="unstable"

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

	Create a container image suitable for use with GPUs.

	STACK is required. Currently supported values are: [rocm, cuda]. Note
	that this will NOT automatically install the stack in the image, that
	is up to you. This just sets p some basic functionality.

	IMAGE, if not specified, defaults to <STACK>/debian:RELEASE.

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

	Synopsis:
	  $0 -h

	  $0 [-m MIRROR] [-r RELEASE] [-u USER] STACK [IMAGE]

	Options:
	  -h          Show this help
	  -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.
	  -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-in-container use with AMD GPUs

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

	  # Creates the image rocm/debian:unstable

	  \$ $0 rocm

	  # Same image as above, but simply tagged 'myimage'

	  \$ $0 rocm myimage

	  # Creates rocm/debian:trixie, using a fast local mirror

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

	EOF
}

while getopts "hm:r:u:" OPTNAME; do
    case $OPTNAME in
    h)
        usage
        exit 0
        ;;
    m) mirror="$OPTARG" ;;
    r) release="$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

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
imagename="${2:-$stack/$distro:$release}"

# /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

if [ ! -e /usr/bin/autopkgtest-build-podman ]; then
    echo "autopkgtest not installed on this system. Aborting." >&2
    exit 1
elif podman image exists "$imagename"; then
    echo "Image $imagename already exists, refusing to overwrite." >&2
    exit 1
fi

# Render group doesn't exist in base images
renderGID="$(getent group render | cut -d: -f3)"
postcmd="groupadd -g $renderGID render"
postcmd+="; useradd -m -G render,video user"
# This looks pointless, but is necessitated by how autopkgtest runs tests
# internally (through su, which behaves differently in containers)
postcmd+="; usermod -a -G render,video root"

# See https://bugs.debian.org/1032487
# This can be removed with autopkgtest >= 5.30
umask 0022

autopkgtest-build-podman \
    --tag "$imagename" \
    --mirror "$mirror" \
    --release "$release" \
    --post-command="$postcmd"
