#!/bin/sh
# Stay alive while any zellij session exists; hold a systemd-inhibit
# lock only while at least one of those zellij sessions was spawned from
# an SSH context.
#
# Rationale: a zellij session started locally (e.g. from a sway terminal)
# is the user actively sitting in front of the laptop — that should NOT
# inhibit suspend. Only zellij sessions started while SSH'd in deserve
# the lock, so the host stays awake across detach + disconnect but
# normal local-attended suspend still works.
#
# Detection: zellij's daemonised server is exec'd by the client and
# inherits the client's environment. Linux preserves that exec-time
# environment in /proc/<pid>/environ for the life of the process, even
# after the original SSH session is gone. So an "ssh-spawned" zellij is
# one whose environ contains SSH_CONNECTION=.
#
# Lifecycle: the .path unit starts this script when the zellij socket
# directory becomes non-empty. The script then polls and stays alive as
# long as any zellij socket exists, so the .path unit never re-triggers
# the service while zellij is up (which previously caused a start-rate
# limit storm when only local zellij was around). When the last zellij
# exits, this script exits, the service stops, and the .path resumes
# watching for the next session.
set -eu

poll=${ZELLIJ_INHIBIT_POLL:-15}
sock_dir="${XDG_RUNTIME_DIR:-/run/user/$(id -u)}/zellij"

has_ssh_zellij() {
  pids=$(pgrep -x zellij 2>/dev/null) || return 1
  for pid in $pids; do
    [ -r "/proc/$pid/environ" ] || continue
    if tr '\0' '\n' <"/proc/$pid/environ" 2>/dev/null |
      grep -q '^SSH_CONNECTION='; then
      return 0
    fi
  done
  return 1
}

any_zellij_socket() {
  [ -d "$sock_dir" ] || return 1
  found=$(ls -A "$sock_dir" 2>/dev/null || true)
  [ -n "$found" ]
}

inhibit_pid=
release_inhibit() {
  pid=$inhibit_pid
  inhibit_pid=
  [ -n "$pid" ] || return 0
  kill "$pid" 2>/dev/null || true
  wait "$pid" 2>/dev/null || true
}
trap release_inhibit EXIT INT TERM

inhibit_alive() {
  [ -n "$inhibit_pid" ] || return 1
  kill -0 "$inhibit_pid" 2>/dev/null
}

acquire_inhibit() {
  if inhibit_alive; then return 0; fi
  systemd-inhibit \
    --what=sleep:idle:handle-lid-switch \
    --who=zellij \
    --why='Active SSH-spawned zellij sessions' \
    --mode=block \
    sleep infinity &
  inhibit_pid=$!
}

while any_zellij_socket; do
  if has_ssh_zellij; then
    acquire_inhibit
  else
    release_inhibit
  fi
  sleep "$poll"
done