#!/usr/bin/env bash
# mem-guard-root.sh — declared root-level memory posture for the workstation.
#
#   mem-guard-root.sh          apply the posture (root)
#   mem-guard-root.sh --audit  report drift, no mutation, no root needed (exit 1 on drift)
#
# Posture:
#   * systemd-oomd kills on sustained memory PSI only. Swap-fullness killing is OFF:
#     zram is priority-100 swap, so "swap 90% full" is normal steady state here and
#     the victim is whichever cgroup holds the most cold pages — always a long-lived
#     agent or terminal session.
#   * No session-wide MemoryHigh/MemoryMax on user-.slice.
#   * agent.slice is governed by sustained memory pressure only — never by a byte or
#     percentage ceiling, and never by swap fullness. Agent work is legitimately heavy.
#   * The owner's desktop and sessions (app.slice, human.slice) carry user.oomd_omit, so
#     the memory-pressure backstop on user@.service can only ever select agent work.
#   * One swappiness sysctl file.
set -uo pipefail

# SYSROOT + OOMCTL exist so the audit can be exercised against a fixture tree.
SYSROOT="${SYSROOT:-}"
OOMCTL="${OOMCTL:-oomctl}"

OOMD_CONF="$SYSROOT/etc/systemd/oomd.conf.d/99-early.conf"
ROOT_SLICE_CONF="$SYSROOT/etc/systemd/system/-.slice.d/10-oomd.conf"
USER_SLICE_CONF="$SYSROOT/etc/systemd/system/user.slice.d/10-oomd.conf"
MEMCAP_CONF="$SYSROOT/etc/systemd/system/user-.slice.d/10-memcap.conf"
USER_MANAGER_CONF_DIR="$SYSROOT/etc/systemd/system/user@.service.d"
SWAPPINESS_CONF="$SYSROOT/etc/sysctl.d/99-swappiness.conf"
SWAPPINESS_PROC="$SYSROOT/proc/sys/vm/swappiness"
# agent.slice is written by three independent layers — the repo unit, the root drop-in,
# and `systemctl --user set-property` runtime drop-ins — so the audit reads all of them.
AGENT_CONF_HOME="${SYSROOT}${HOME:-/root}"
CGROUP_ROOT="${SYSROOT}/sys/fs/cgroup"
OWNER_UID="${OWNER_UID:-${SUDO_UID:-$(id -u)}}"
SWAPPINESS=10
# oomd's own defaults, applied to every monitored cgroup that declares no limit of its own
# (user@0, user@997 — agent/build/user@1000 all declare theirs). Declared here so that
# applying the posture cannot silently move numbers a previous apply left behind.
OOMD_DEFAULTS='[OOM]
DefaultMemoryPressureLimit=40%
DefaultMemoryPressureDurationSec=10s
SwapUsedLimit=45%'

drift=0
ok()  { printf 'OK    %-16s %s\n' "$1" "$2"; }
bad() { printf 'DRIFT %-16s %s\n' "$1" "$2"; drift=1; }

audit() {
  if command -v "$OOMCTL" >/dev/null 2>&1; then
    local monitored
    monitored=$("$OOMCTL" 2>/dev/null | sed -n '/^Swap Monitored CGroups:/,/^Memory Pressure Monitored CGroups:/p' | grep -c '^[[:space:]]*Path:')
    if [ "$monitored" -eq 0 ]; then
      ok oomd-swap-kill "disarmed (0 swap-monitored cgroups)"
    else
      bad oomd-swap-kill "$monitored cgroup(s) armed for swap-fullness SIGKILL — sessions will be killed"
    fi
  else
    bad oomd-swap-kill "oomctl missing"
  fi

  if [ "$(cat "$OOMD_CONF" 2>/dev/null)" = "$OOMD_DEFAULTS" ]; then
    ok oomd-defaults "as declared"
  else
    bad oomd-defaults "$OOMD_CONF differs from the declared defaults"
  fi

  local swap_kill_files
  swap_kill_files=$(grep -ls 'ManagedOOMSwap=kill' \
    "$ROOT_SLICE_CONF" "$USER_SLICE_CONF" "$USER_MANAGER_CONF_DIR"/*.conf 2>/dev/null)
  if [ -n "$swap_kill_files" ]; then
    bad oomd-dropins "ManagedOOMSwap=kill present in $(printf '%s' "$swap_kill_files" | tr '\n' ' ')"
  else
    ok oomd-dropins "no ManagedOOMSwap=kill"
  fi

  if grep -qsE '^\s*Memory(High|Max)=' "$MEMCAP_CONF"; then
    bad user-slice-cap "$MEMCAP_CONF caps the whole session (reclaim storm)"
  else
    ok user-slice-cap "uncapped"
  fi

  local agent_files agent_hits
  agent_files=$(ls -d "$SYSROOT/etc/systemd/user/agent.slice" \
    "$SYSROOT/etc/systemd/user/agent.slice.d"/*.conf \
    "$AGENT_CONF_HOME/.config/systemd/user/agent.slice" \
    "$AGENT_CONF_HOME/.config/systemd/user/agent.slice.d"/*.conf \
    "$AGENT_CONF_HOME/.config/systemd/user.control/agent.slice.d"/*.conf \
    "$SYSROOT/run/user/$OWNER_UID/systemd/user.control/agent.slice.d"/*.conf \
    2>/dev/null)
  agent_hits=$(printf '%s\n' "$agent_files" | grep -v '^$' \
    | xargs -r grep -lP '^\s*(Memory(High|Max|SwapMax)=(?!infinity)|ManagedOOMSwap=kill)' 2>/dev/null)
  if [ -n "$agent_hits" ]; then
    bad agent-slice-cap "agent.slice is capped by a number or armed for swap-kill in: $(printf '%s' "$agent_hits" | tr '\n' ' ')"
  else
    ok agent-slice-cap "pressure-only (no memory ceiling, no swap kill)"
  fi

  # The pressure backstop on user@.service can only ever pick a cgroup that is not marked
  # omit, so the owner's desktop and sessions (app.slice) and human.slice must carry the
  # mark. The kernel xattr is the truth: the drop-in alone is inert until the user manager
  # reloads, and `systemctl --user show` reports the parsed unit either way.
  local user_cg="$CGROUP_ROOT/user.slice/user-$OWNER_UID.slice/user@$OWNER_UID.service"
  local unomitted=""
  for s in app.slice human.slice; do
    [ -d "$user_cg/$s" ] || continue
    [ "$(getfattr --only-values -n user.oomd_omit "$user_cg/$s" 2>/dev/null)" = 1 ] \
      || unomitted="$unomitted $s"
  done
  if [ -n "$unomitted" ]; then
    bad owner-oom-omit "oomd may select the owner's session:$unomitted lack user.oomd_omit"
  else
    ok owner-oom-omit "app.slice/human.slice omitted from oomd candidates"
  fi

  local files count live
  files=$(grep -lsE '^[[:space:]]*vm\.swappiness' "$SYSROOT"/etc/sysctl.d/*.conf "$SYSROOT/etc/sysctl.conf" 2>/dev/null)
  count=$(printf '%s' "$files" | grep -c .)
  live=$(cat "$SWAPPINESS_PROC" 2>/dev/null)
  if [ "$count" -le 1 ] && [ "$live" = "$SWAPPINESS" ]; then
    ok swappiness "$SWAPPINESS (${files:-none})"
  else
    bad swappiness "live=${live:-unknown} want=$SWAPPINESS, ${count} declaring file(s): $(printf '%s' "$files" | tr '\n' ' ')"
  fi

  return "$drift"
}

apply() {
  [ "$EUID" -eq 0 ] || { echo "run as root: deck-sudo bash $0" >&2; exit 1; }

  install -d -m0755 "$(dirname "$OOMD_CONF")" "$(dirname "$ROOT_SLICE_CONF")" "$(dirname "$MEMCAP_CONF")"

  printf '%s\n' "$OOMD_DEFAULTS" > "$OOMD_CONF"

  cat > "$ROOT_SLICE_CONF" <<'EOF'
[Slice]
# kill => oomd SIGKILLs the largest-swap cgroup (an agent or terminal session) at 90% swap.
ManagedOOMSwap=auto
EOF

  rm -f "$USER_SLICE_CONF"

  cat > "$MEMCAP_CONF" <<'EOF'
[Slice]
# No session-wide cap: MemoryHigh=48G here caused 8.7M reclaim events and froze the box.
EOF

  printf 'vm.swappiness=%s\n' "$SWAPPINESS" > "$SWAPPINESS_CONF"
  while IFS= read -r f; do
    [ -n "$f" ] && [ "$f" != "$SWAPPINESS_CONF" ] && sed -i -E '/^\s*vm\.swappiness/d' "$f"
  done < <(grep -lsE '^[[:space:]]*vm\.swappiness' "$SYSROOT"/etc/sysctl.d/*.conf "$SYSROOT/etc/sysctl.conf" 2>/dev/null)
  sysctl -q -w "vm.swappiness=$SWAPPINESS"

  systemctl daemon-reload
  systemctl set-property --runtime -- -.slice ManagedOOMSwap=auto 2>/dev/null
  systemctl restart systemd-oomd

  audit
}

case "${1:-apply}" in
  --audit) audit;;
  apply)   apply;;
  *)       echo "usage: $0 [--audit]" >&2; exit 2;;
esac