#!/bin/bash
# shellcheck disable=SC1091 # Runtime-resolved path; intentionally not followed.
. "$(dirname "$(realpath "${BASH_SOURCE[0]}")")/../local-snap.sh" || return 1
set -eux

# Install LXD
install_lxd

# Configure LXD
lxc network create lxdbr0
lxc profile device add default eth0 nic network=lxdbr0

IMAGE="${TEST_IMG:-ubuntu-minimal-daily:26.04}"

# checkNestedInstanceDNS checks if the nested instance can resolve DNS names.
checkNestedInstanceDNS() {
    local instName=$1
    local nestedInstName=$2

    # Test DNS resolution.
    for _ in $(seq 1 5); do
        if lxc exec "${instName}" -- lxc exec "${nestedInstName}" -- getent hosts canonical.com; then
            echo "===> PASS: DNS resolution: ${nestedInstName} (on instance ${instName})"
            return 0
        fi
        sleep 1
    done

    echo "===> FAIL: DNS resolution: ${nestedInstName} (on instance ${instName})"
    return 1
}

test_nesting_container() {
    if [[ "${LXD_SNAP_CHANNEL}" =~ ^4\.0/ ]]; then
        echo "Skipping container nesting tests on ${LXD_SNAP_CHANNEL}, not supported"
        return 0
    fi

    echo "==> Running nesting test for container"

    local poolDriverList="btrfs dir lvm zfs"
    local poolDriver poolName instName
    local nestedPoolDriverList nestedPoolDriver nestedPoolName nestedInstName nestedSource
    local zfsVersion zfsPerm

    lxc profile create container-nesting
    lxc profile set container-nesting security.nesting=true

    if hasNeededAPIExtension devlxd_images_vm; then
        lxc profile set container-nesting security.devlxd.images=true
    fi

    for poolDriver in ${poolDriverList}; do
        poolName="pool-${poolDriver}"
        instName="c1"

        # Create storage pool.
        lxc storage create "${poolName}" "${poolDriver}"

        # Configure the list of storage drivers to be tested in the nested container.
        # Storage driver "dir" is always supported.
        nestedPoolDriverList="dir"

        if [ "${poolDriver}" = "btrfs" ]; then
            # BTRFS can be used within a nested container on top of BTRFS storage
            # pool, if the source is a directory.
            nestedPoolDriverList+=" btrfs"
        elif [ "${poolDriver}" = "zfs" ]; then
            # To use ZFS within a nested container, a ZFS version must be greater than
            # or equal to 2.2, and ZFS delegation must be enabled.
            zfsVersion="$(modinfo -F version zfs)"
            if dpkg --compare-versions "${zfsVersion}" ge "2.2"; then
                lxc storage set "${poolName}" volume.zfs.delegate=true
                nestedPoolDriverList+=" zfs"

                # XXX: Ensure that the ZFS device is accessible within the nested container.
                #      Otherwise, the LXD storage pool creation will fail.
                zfsPerm=$(stat -c '%a' /dev/zfs)
                if [ $((zfsPerm & 7)) -eq 0 ]; then
                    # It's udev's job to fix the perms but the rule for it ships in
                    # the zfsutils-linux package which might not be installed.
                    chmod 0666 /dev/zfs
                fi
            else
                echo "ZFS version ${zfsVersion} does not support ZFS delegation and cannot be used within nested containers."
            fi
        fi

        # Launch a new container and install LXD.
        lxc launch "${IMAGE}" "${instName}" --profile default --profile container-nesting --storage "${poolName}"
        waitInstanceReady "${instName}"

        # Ensure snapd is ready.
        lxc exec "${instName}" -- sh -c "$(declare -f waitSnapdSeed); waitSnapdSeed"

        # Install LXD in the container.
        lxc exec "${instName}" -- snap install lxd --channel="${LXD_SNAP_CHANNEL}"
        lxc exec "${instName}" -- lxd init --auto

        # Test each storage driver in nested container.
        for nestedPoolDriver in ${nestedPoolDriverList}; do
            nestedPoolName="pool-nested-${nestedPoolDriver}"
            nestedInstName="nested-c1"

            # Create nested storage pool.
            if [ "${nestedPoolDriver}" = "btrfs" ]; then
                # BTRFS pool must use a directory as a source. The path must not already exist
                # (LXD creates the subvolume itself), so only generate a unique name (inside the
                # instance, not the host) rather than actually creating the directory. Anchored
                # under /root rather than /tmp, since the latter may be a separate tmpfs mount.
                nestedSource="$(lxc exec "${instName}" -- mktemp -u -d -p /root)"
                lxc exec "${instName}" -- lxc storage create "${nestedPoolName}" "${nestedPoolDriver}" source="${nestedSource}"
            elif [ "${nestedPoolDriver}" = "zfs" ]; then
                # Create a custom volume and attach it to the instance.
                lxc storage volume create "${poolName}" vol1
                lxc config device add "${instName}" vol1 disk pool="${poolName}" source=vol1 path=/mnt/vol1

                # Create a ZFS storage pool in the nested container using the attached volume.
                lxc exec "${instName}" -- lxc storage create "${nestedPoolName}" "${nestedPoolDriver}" source="${poolName}/custom/default_vol1"
            else
                lxc exec "${instName}" -- lxc storage create "${nestedPoolName}" "${nestedPoolDriver}"
            fi

            # Launch nested container.
            lxc exec "${instName}" -- lxc launch "${IMAGE}" "${nestedInstName}" --storage "${nestedPoolName}"
            lxc exec "${instName}" -- sh -c "$(declare -f waitInstanceReady); waitInstanceReady ${nestedInstName}"

            # Test container connection.
            checkNestedInstanceDNS "${instName}" "${nestedInstName}"

            # Test graceful shutdown of the nested container.
            lxc exec "${instName}" -- timeout 30 lxc stop "${nestedInstName}"
            lxc exec "${instName}" -- timeout 30 lxc delete "${nestedInstName}"
        done

        # Cleanup.
        lxc delete --force "${instName}"
        if echo "${nestedPoolDriverList}" | grep -wF zfs >/dev/null; then
            lxc storage volume delete "${poolName}" vol1
        fi
        lxc storage delete "${poolName}"
    done

    lxc profile delete container-nesting
}

test_nesting_vm() {
    echo "==> Running nesting test for vm"

    local VM_IN_CTN=1
    local VMs=10
    local nestedVMs=3
    local storageDriver="zfs"
    local zfsVersion instCount

    if [[ "${LXD_SNAP_CHANNEL}" =~ ^4\.0/ ]]; then
        # VMs on 4.0 are more resource intensive and slower to boot
        export MAX_WAIT_SECONDS=240

        # VMs inside containers do not work
        VM_IN_CTN=0

        VMs=3
        nestedVMs=2

        zfsVersion="$(modinfo -F version zfs)"
        if dpkg --compare-versions "${zfsVersion}" gt "2.2"; then
            storageDriver="lvm"
            echo "${LXD_SNAP_CHANNEL} does not support ZFS ${zfsVersion}, falling back to ${storageDriver}"
        fi
    fi

    lxc project switch default
    lxc storage create default "${storageDriver}" size=30GiB

    function parallel() {
        seq "$1" | xargs -P "$1" -I "{}" "${@:2}"
    }

    function init() {
        vm="${2:-}"
        if [ -z "${vm}" ]
        then
            parallel "$1" lxc init "${IMAGE}" "t{}" -s default -n lxdbr0 --quiet
        else
            parallel "$1" lxc init "${IMAGE}" "t{}" "${vm}" -s default -n lxdbr0 --quiet
        fi
    }

    function conf() {
        # Supports setting multiple key=value pairs in a single call.
        parallel "$1" lxc config set "t{}" "${@:2}"
    }

    function device_add() {
        parallel "$1" lxc config device add "t{}" "$2" "$3" "$4"
    }

    function start() {
        instances=()
        for i in $(seq "$1"); do
            instances["$i"]="t$i"
        done

        echo "Start ${instances[*]}"
        lxc start "${instances[@]}"
    }

    function wait() {
        for i in $(seq "$1"); do
            waitInstanceReady "t${i}"
            setup_instance_gocoverage "t${i}"
        done
    }

    function cmd() {
        parallel "$1" lxc exec "t{}" -- bash -c "$2"
    }

    function delete() {
        instances=()
        for i in $(seq "$1"); do
            instances["$i"]="t$i"
            prepare_vm_for_hard_stop "t$i"
        done

        echo "Delete ${instances[*]}"
        lxc delete -f "${instances[@]}"
    }

    echo "==> Test ${VMs} VMs in parallel"
    init "${VMs}" --vm
    start "${VMs}"
    delete "${VMs}"

    echo "==> Test ${VMs} VMs in parallel for vsock ID collision"
    init "${VMs}" --vm
    conf "${VMs}" volatile.vsock_id=42
    start "${VMs}"
    delete "${VMs}"

    echo "==> Test ${nestedVMs} VMs each with one nested VM"
    init "${nestedVMs}" --vm
    if hasNeededAPIExtension devlxd_images_vm; then
        conf "${nestedVMs}" security.devlxd.images=true
    fi
    start "${nestedVMs}"
    wait "${nestedVMs}"

    cmd "${nestedVMs}" "$(declare -f waitSnapdSeed); waitSnapdSeed"
    cmd "${nestedVMs}" "snap install lxd --channel ${LXD_SNAP_CHANNEL}"
    cmd "${nestedVMs}" "lxd init --auto"
    cmd "${nestedVMs}" "lxc profile device set default root size=3584MiB"
    cmd "${nestedVMs}" "lxc launch ${IMAGE} nested --vm -c limits.memory=512MiB --quiet"
    delete "${nestedVMs}"

    if [ "${VM_IN_CTN}" = "1" ]; then
        echo "==> Test ${nestedVMs} containers each with one nested VM"
        init "${nestedVMs}"
        if hasNeededAPIExtension devlxd_images_vm; then
            conf "${nestedVMs}" security.devlxd.images=true
        fi
        device_add "${nestedVMs}" kvm unix-char source=/dev/kvm
        device_add "${nestedVMs}" vhost-net unix-char source=/dev/vhost-net
        device_add "${nestedVMs}" vhost-vsock unix-char source=/dev/vhost-vsock
        device_add "${nestedVMs}" vsock unix-char source=/dev/vsock
        start "${nestedVMs}"
        cmd "${nestedVMs}" "$(declare -f waitSnapdSeed); waitSnapdSeed"
        cmd "${nestedVMs}" "snap install lxd --channel ${LXD_SNAP_CHANNEL}"
        cmd "${nestedVMs}" "lxd init --auto"
        cmd "${nestedVMs}" "lxc profile device set default root size=3584MiB"
        cmd "${nestedVMs}" "lxc launch ${IMAGE} nested --vm -c limits.memory=512MiB --quiet"
    else
        echo "Skipping VM in container tests on ${LXD_SNAP_CHANNEL}"
    fi

    echo "==> Cleaning up"
    instCount="$(lxc list -f csv -c n t | wc -l)"
    [ "$instCount" -gt 0 ] && delete "${instCount}"

    lxc storage delete default
}

# Allow for running a specific set of tests.
if [ "$#" -gt 0 ]; then
    for test_name in "$@"; do
        case "${test_name}" in
            container|nesting-container|test_container|test_nesting_container)
                test_nesting_container
                ;;
            vm|nesting-vm|test_vm|test_nesting_vm)
                test_nesting_vm
                ;;
            *)
                echo "Unknown test: ${test_name}" >&2
                exit 1
                ;;
        esac
    done
else
    test_nesting_container
    test_nesting_vm
fi

echo "==> Cleaning up environment"
lxc profile device remove default eth0
lxc network delete lxdbr0

# shellcheck disable=SC2034
FAIL=0
