#!/usr/bin/env bash
# Preserve enough host state to restore the temporary RTX 3080 reservation.
set -euo pipefail
state=/host/var/lib/atlas-maintenance/lan-inference-20260928
host() { chroot /host "$@"; }
umask 077
mkdir -p "$state"

case "${1:-}" in
  reserve)
    # Wolf must be scaled down through Flux before stopping its child containers.
    for attempt in $(seq 1 120); do
      host pgrep -x wolf >/dev/null || break
      sleep 2
    done
    if host pgrep -x wolf >/dev/null; then
      echo 'Wolf still running; reservation refused' >&2
      exit 1
    fi
    if ! test -f "$state/saved"; then
      host systemctl show display-manager -p Id --value > "$state/display-unit"
      host systemctl is-active display-manager > "$state/display-active" || true
      host systemctl is-enabled display-manager > "$state/display-enabled" || true
      host docker ps --filter 'name=^/Wolf' --format '{{.ID}}' > "$state/docker-running"
      touch "$state/saved"
    fi
    unit=$(cat "$state/display-unit")
    case "$unit" in sddm.service|gdm.service|gdm3.service|lightdm.service) ;;
      *) echo 'Unexpected display manager; refusing host mutation' >&2; exit 1 ;;
    esac
    host systemctl mask --now "$unit"
    while read -r container; do
      test -z "$container" || host docker stop --time 30 "$container"
    done < "$state/docker-running"
    touch "$state/reserved"
    echo 'GPU desktop and Wolf containers stopped; restoration state saved'
    ;;
  restore)
    test -f "$state/saved"
    unit=$(cat "$state/display-unit")
    case "$unit" in sddm.service|gdm.service|gdm3.service|lightdm.service) ;;
      *) exit 1 ;;
    esac
    case "$(cat "$state/display-enabled")" in
      masked*) ;;
      *) host systemctl unmask "$unit" ;;
    esac
    if test "$(cat "$state/display-active")" = active; then
      host systemctl start "$unit"
    fi
    while read -r container; do
      test -z "$container" || host docker start "$container"
    done < "$state/docker-running"
    rm -f "$state/reserved"
    echo 'Saved display and Wolf container state restored'
    ;;
  *) echo 'Usage: gpu_pilot_reservation.sh reserve|restore' >&2; exit 2 ;;
esac