#!/bin/sh
set -eux

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

    echo "Test passed"
    exit 0
}

architecture="$(uname -m)"

if [ "${architecture}" != "x86_64" ] && [ "${architecture}" != "s390x" ]; then
  echo "Skipping test as CPU hotplugging not supported on ${architecture}"
fi

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

waitVMAgent() (
  set +x
  # shellcheck disable=SC3043
  local vmName="$1"
  for i in $(seq 90) # Wait up to 90s.
  do
    if incus info "${vmName}" | grep -qF 127.0.0.1; then
      return 0 # Success.
    fi

    sleep 1
  done

  echo "VM ${vmName} agent not running after ${i}s"
  return 1 # Failed.
)

# Configure Incus
incus network create incusbr0
incus profile device add default eth0 nic network=incusbr0

poolName="vmpool$$"
poolDriver=dir

echo "==> Create storage pool using driver ${poolDriver}"
incus storage create "${poolName}" "${poolDriver}"

echo "==> Create ephemeral VM and boot"
incus launch images:ubuntu/22.04 v1 --vm -s "${poolName}" --ephemeral
waitVMAgent v1
incus info v1

# Get number of CPUs
# shellcheck disable=SC2010
cpuCount="$(ls /sys/devices/system/cpu | grep -Ec 'cpu[[:digit:]]+')"

# VMs should have only 1 CPU per default
[ "$(incus exec v1 -- ls /sys/devices/system/cpu | grep -Ec 'cpu[[:digit:]]+')" -eq "1" ]

# Set valid CPU limits (low to high)
for i in $(seq 2 "${cpuCount}"); do
  incus config set v1 limits.cpu="${i}"
  [ "$(incus exec v1 -- ls /sys/devices/system/cpu | grep -Ec 'cpu[[:digit:]]+')" -eq "${i}" ]
done

# Try setting more CPUs than available
! incus config set v1 limits.cpu="$(( cpuCount + 1 ))" || false

# Set valid CPU limits (high to low)
for i in $(seq "${cpuCount}" -1 1); do
  incus config set v1 limits.cpu="${i}"
  [ "$(incus exec v1 -- ls /sys/devices/system/cpu | grep -Ec 'cpu[[:digit:]]+')" -eq "${i}" ]
done

# Try doing pinning while VM is running (shouldn't work)
! incus config set v1 limits.cpu=1,2 || false

# Set max CPU count
incus config set v1 limits.cpu="${cpuCount}"
[ "$(incus exec v1 -- ls /sys/devices/system/cpu | grep -Ec 'cpu[[:digit:]]+')" -eq "${cpuCount}" ]

# Unset CPU limit
incus config unset v1 limits.cpu

# Unsetting the limit should leave the VM with 1 CPU
[ "$(incus exec v1 -- ls /sys/devices/system/cpu | grep -Ec 'cpu[[:digit:]]+')" -eq "1" ]

echo "==> Stopping and deleting ephemeral VM"
# Stop VM and check its deleted.
incus stop -f v1
! incus info v1 || false

incus profile device remove default eth0

echo "==> Deleting storage pool"
incus storage delete "${poolName}"

echo "==> Deleting storage pool"
incus network delete incusbr0

FAIL=0
