#!/usr/bin/bash

set -e
set -u
set -o pipefail

TASK=${1:-default}

# Helper function so that we (under "set -e") don't error out when
# qvm-check returns nonzero for a domain that isn't running.
function shutdown_if_running () {
    set +e
    qvm-check --running "$1" && qvm-shutdown --wait "$1"
    set -e
}

# To allow the template of an AppVM to be changed, the following two
# conditions must be met:
# 1. The AppVM must be powered off
# 2. The AppVM must not be a DispVM template that used as the default DispVM
#    for an AppVM, nor the system default DispVM.
if [[ $TASK == "prepare" ]]; then
    # sd-app, we simply shutdown the machine as we want to preserve the data
    if qvm-check sd-app --quiet; then
        BASE_TEMPLATE=$(qvm-prefs sd-app template)
        if [[ ! $BASE_TEMPLATE =~ "small-bookworm" ]]; then
            shutdown_if_running "sd-app"
        fi
    fi

    # For sd-viewer and sd-devices-dvm, DispVM templates. We can delete both
    # VMs since they contain no persistent data. The installer will re-create them
    # as part of the provisioning process.
    # We set the default DispVM to empty string to ensure nothing is opened in an
    # insecure (unmanaged or not yet updated) or networked vm, until the
    # provisioning process runs again and sets that value to sd-viewer
    if qvm-check --quiet sd-viewer; then
        BASE_TEMPLATE=$(qvm-prefs sd-viewer template)
        if [[ ! $BASE_TEMPLATE =~ "large-bookworm" ]]; then
            qubes-prefs default_dispvm ''
            shutdown_if_running "sd-viewer"
            qvm-remove -f sd-viewer
        fi
    fi

    if qvm-check --quiet sd-devices; then
        BASE_TEMPLATE=$(qvm-prefs sd-devices-dvm template)
        if [[ ! $BASE_TEMPLATE =~ "large-bookworm" ]]; then
            shutdown_if_running "sd-printers"
            shutdown_if_running "sd-devices"
            shutdown_if_running "sd-devices-dvm"
            qvm-remove -f sd-printers
            qvm-remove -f sd-devices
            qvm-remove -f sd-devices-dvm
        fi
    fi

    if qvm-check --quiet sd-proxy-dvm; then
        BASE_TEMPLATE=$(qvm-prefs sd-proxy-dvm template)
        if [[ ! $BASE_TEMPLATE =~ "large-bookworm" ]]; then
            shutdown_if_running "sd-proxy"
        fi
    fi

    # For sd-gpg, we simply shutdown the machine
    if qvm-check --quiet sd-gpg; then
        BASE_TEMPLATE=$(qvm-prefs sd-gpg template)
        if [[ ! $BASE_TEMPLATE =~ "small-bookworm" ]]; then
            shutdown_if_running "sd-gpg"
        fi
    fi

    # Shut down sd-log last, since other VMs will autostart it by sending logs
    if qvm-check --quiet sd-log; then
        BASE_TEMPLATE=$(qvm-prefs sd-log template)
        if [[ ! $BASE_TEMPLATE =~ "small-bookworm" ]]; then
            shutdown_if_running "sd-log"
        fi
    fi
elif [[ $TASK == "remove" ]]; then
    # For each template, ensure the TemplateVM exists, that it is shut down
    # before deleting it.
    # TODO: clean this up, we don't have separate templates anymore and nobody
    # will be upgrading from the original setup
    for template in sd-small-bullseye-template sd-large-bullseye-template
    do
        if qvm-check "${template}" --quiet; then
            shutdown_if_running "${template}"
            qvm-remove -f "${template}"
        fi
    done
else
    echo "Please specify prepare or remove"
    exit 1
fi
