#!/usr/bin/env bash
# memwatch -- control wrapper for the memwatch PSI watchdog service.
set -euo pipefail

UNIT=memwatch.service
CONF=/etc/memwatch.conf

# This box uses systemd's run0. pkexec and sudo are fallbacks for when the
# polkit path for run0 is unhappy.
priv() {
  if command -v run0 >/dev/null 2>&1 && run0 true >/dev/null 2>&1; then
    run0 "$@"
  elif command -v pkexec >/dev/null 2>&1; then
    pkexec "$@"
  elif command -v sudo >/dev/null 2>&1; then
    sudo "$@"
  else
    echo "memwatch: no way to escalate privileges (run0/pkexec/sudo all unusable)" >&2
    return 1
  fi
}

mode() {
  [[ -r $CONF ]] || { echo "unknown (no $CONF)"; return; }
  if grep -qE '^\s*dry_run\s*=\s*0' "$CONF"; then echo live; else echo "dry-run"; fi
}

set_dry() {
  priv sed -i -E "s/^\s*dry_run\s*=.*/dry_run = $1/" "$CONF"
}

usage() {
  cat <<'USAGE'
usage: memwatch <command>

  on        enable and start the watchdog (persists across reboots)
  off       stop and disable it
  start     start it for this boot only
  stop      stop it for this boot only
  restart   restart it
  status    service state, current mode, and live memory pressure
  logs      follow the journal (ctrl-c to quit)
  dry       switch to dry-run: log victims, kill nothing, then restart
  live      switch to live: actually freeze and kill, then restart
USAGE
}

case "${1-status}" in
  on)      priv systemctl enable --now "$UNIT" ;;
  off)     priv systemctl disable --now "$UNIT" ;;
  start)   priv systemctl start "$UNIT" ;;
  stop)    priv systemctl stop "$UNIT" ;;
  restart) priv systemctl restart "$UNIT" ;;
  logs)    journalctl -fu "$UNIT" ;;
  dry)     set_dry 1; priv systemctl try-restart "$UNIT"; echo "mode: dry-run" ;;
  live)
    echo "This lets memwatch freeze and kill real processes."
    read -rp "type yes to confirm: " ans
    [[ $ans == yes ]] || { echo "aborted"; exit 1; }
    set_dry 0; priv systemctl try-restart "$UNIT"; echo "mode: live" ;;
  status)
    printf 'enabled: %s\n' "$(systemctl is-enabled "$UNIT" 2>&1)"
    printf 'active:  %s\n' "$(systemctl is-active "$UNIT" 2>&1)"
    printf 'mode:    %s\n' "$(mode)"
    printf 'restarts: %s\n' "$(systemctl show "$UNIT" -P NRestarts 2>/dev/null)"
    echo
    echo "/proc/pressure/memory:"
    sed 's/^/  /' /proc/pressure/memory
    echo
    journalctl -u "$UNIT" --no-pager -n 5 -o cat 2>/dev/null | sed 's/^/  /'
    ;;
  -h|--help|help) usage ;;
  *) usage >&2; exit 2 ;;
esac
