#!/bin/sh
# vim: set ts=4 sts=4 sw=4 et :

CRITICAL_VMS='
cloud0
dc0
dc2
file0
k8s-amd64-n0
k8s-ctrl0
logs0
matrix0
smtp1
web0
'
VM_HOSTS='
vmhost0.pyrocufflink.blue
vmhost1.pyrocufflink.blue
'

MESSAGE="$1"

consolidate_vms() {
    count=0
    for h in ${VM_HOSTS}; do
        [ -n "${h}" ] || continue
        uri=qemu+ssh://root@${h}/system
        c=$(($(virsh -c ${uri} list --name | wc -l) - 1))
        printf 'VM host %s has %d running VMs\n' "${h}" "${c}" \
            | log -p daemon.debug
        if [ ${c} -gt ${count} ]; then
            target=${h}
            count=${c}
        fi
    done
    [ -n "${target}" ] || return 1
    printf 'Consolidating VMs onto host %s\n' "${target}" \
        | log -p daemon.warning
    for h in ${VM_HOSTS}; do
        [ -n "${h}" ] || continue
        [ "${h}" != "${target}" ] || continue
        migrate_vms ${h} ${target}
        remote_shutdown ${h}
    done
}

handle_online() {
    notify_ntfy high white_check_mark,electric_plug
}

handle_onbatt() {
    notify_ntfy high warning,battery
}

handle_lowbatt() {
    notify_ntfy max warning,battery

    if is_only_ups; then
        shut_down_all
    else
        low_power_mode
    fi
}

handle_commok() {
    notify_ntfy high white_check_mark
}

handle_commbad() {
    notify_ntfy high warning
}

handle_replbatt() {
    notify_ntfy default warning,battery
}

handle_nocomm() {
    notify_ntfy high warning
}

handle_unknown() {
    printf 'Unknown notification type %s for UPS %s: %s' \
        "${NOTIFYTYPE}" \
        "${UPSNAME}" \
        "${MESSAGE}" \
        | log -p daemon.notice
}

is_only_ups() {
    for ups in $(upsc -l); do
        bat_chg=$(upsc "${ups}" battery.charge)
        bat_low=$(upsc "${ups}" battery.charge.low)
        if [ "${bat_chg}" -gt "${bat_low}" ]; then
            return 1
        fi
    done
    return 0
}

log() {
    logger --id=$$ -t "${0##*/}" "$@"
}

low_power_mode() {
    echo 'UPS battery low, going to low power mode' \
        | log -p daemon.warning

    remote_shutdown host-dca632474c83.pyrocufflink.red
    remote_shutdown burp1.pyrocufflink.blue
    stop_noncritical_vms
    consolidate_vms
}

migrate_vms() {
    srcuri=qemu+ssh://root@${1}/system
    tgturi=qemu+ssh://root@${2}/system
    for vm in $(virsh -c ${srcuri} list --name); do
        [ -n "${vm}" ] || continue
        printf 'Migrating %s from %s to %s\n' "${vm}" "${1}" "${2}" \
            | log -p daemon.info
        virsh -c ${srcuri} migrate ${vm} ${tgturi} \
            --live \
            --persistent \
            --undefinesource
    done
}

notify_ntfy() {
    priority="${1:-default}"
    tags="${2}"

    curl -fsSL https://ntfy.pyrocufflink.blue/pixel5 \
        -o /dev/null \
        -H 'Content-Type: text/plain' \
        -H "Priority: ${priority}" \
        -H "Tags: ${tags}" \
        -d "${MESSAGE}"
}

remote_shutdown() {
    printf 'Shutting down %s\n' "${1}" \
        | log -p daemon.warning

    # XXX Should probably create a nut user on remote machines
    ssh \
        -o BatchMode=yes \
        -o StrictHostKeyChecking=no \
        -o ConnectTimeout=5 \
        -l root \
        "${1}" \
        systemctl poweroff
}

shut_down_all() {
    echo 'Battery low on last UPS, shutting down' \
        | log -p daemon.alert

    for h in ${VM_HOSTS}; do
        remote_shutdown ${h}
    done
}

stop_noncritical_vms() {
    echo 'Stopping all non-critical virtual machines' \
        | log -p daemon.warning

    t_running=$(mktemp)
    t_critical=$(mktemp)
    for name in ${CRITICAL_VMS}; do
        [ -n "${name}" ] || continue
        echo "${name}"
    done | sort > "${t_critical}"
    for h in ${VM_HOSTS}; do
        [ -n "${h}" ] || continue
        ssh -o BatchMode=yes -o StrictHostKeyChecking=no -l root ${h} :
        uri=qemu+ssh://root@${h}/system
        virsh -c ${uri} list --name | sort > "${t_running}"
        for name in $(comm -23 "${t_running}" "${t_critical}"); do
            [ -n "${name}" ] || continue
            printf 'Stopping virtual machine: %s\n' "${name}" \
                | log -p daemon.info
            #virsh -c ${uri} shutdown ${name}
        done
    done
    
    rm -f "${t_running}" "${t_critical}"
}

printf '%s running as %d %s\n' "${0##*/}" "$(id -u)" "$(whoami)" \
    | log -p daemon.info
printf '%s %s %s\n' "${UPSNAME}" "${NOTIFYTYPE}" "${MESSAGE}" \
    | log -p daemon.info

PATH=/usr/bin:/bin:/usr/sbin:/sbin
export PATH

case "${NOTIFYTYPE}" in
    ONLINE)
        handle_online
        ;;
    ONBATT)
        handle_onbatt
        ;;
    LOWBATT)
        handle_lowbatt
        ;;
    COMMOK)
        handle_commok
        ;;
    COMMBAD)
        handle_commbad
        ;;
    REPLBATT)
        handle_replbatt
        ;;
    NOCOMM)
        handle_nocomm
        ;;
    *)
        handle_unknown
        ;;
esac