#!/usr/bin/env bash
# Restore after the ways an OS stops tmux, not only `tmux kill-server`.
#
#   tests/shutdown.sh [path/to/dingger] [rounds]
#
# At a restart or logout macOS asks apps to quit (the terminal: its tmux
# client goes), then launchd sends SIGTERM to every process left in the
# session at about the same moment (tmux, agents, shells, dingger's daemon),
# and SIGKILL to what is still alive some seconds later. Each scenario below
# stops an isolated tmux server one of those ways, then logs in again: a
# fresh daemon if its own was stopped, and a terminal starting tmux with
# `tmux new-session -As 0` (a bare session 0). The layout must be offered,
# and restoring it must bring back every session, window, pane and agent.
#
# Each scenario has its own tmux servers (-L), state directory, hook socket
# and pi directory; agents are fakes named pi that record their arguments.
# Nothing here touches the user's tmux server, daemon or agents. Signals go
# only to the test's own processes: its daemon, its servers, and what runs
# in their panes.
set -u

ROOT=$(cd "$(dirname "$0")/.." && pwd)
DINGGER="${1:-$ROOT/target/release/dingger}"
case "$DINGGER" in /*) ;; *) DINGGER="$PWD/$DINGGER" ;; esac
ROUNDS="${2:-20}"
# Resolved (/tmp is a link): what a pane's $PWD says.
BASE=$(cd "$(mktemp -d /tmp/dingger-shutdown.XXXXXX)" && pwd -P)
PASSES=0
FAILS=0

pass() { PASSES=$((PASSES + 1)); echo "  ok   $1"; }
fail() { FAILS=$((FAILS + 1)); echo "  FAIL $1"; }
expect() { # description command expected [tries]  (polls every 0.1s, default 4s)
  local desc=$1 cmd=$2 want=$3 tries=${4:-40} got
  for _ in $(seq "$tries"); do
    got=$(eval "$cmd")
    [ "$got" = "$want" ] && { pass "$desc"; return 0; }
    sleep 0.1
  done
  fail "$desc: want [$want] got [$got]"
  return 1
}

# One scenario's world; set by setup().
D= SOCK= TERMSOCK= DAEMON_PID= WORLD=()
# The environment of everything in a scenario: its daemon, its servers, every
# tmux command (a session or window gets the environment of the client that
# creates it), restore. PATH has the fake pi first and no real one (the
# user's agents live under ~/.local and ~/.claude). SHELL too: tmux runs
# commands with it, and the user's zsh would read ~/.zshenv, which puts the
# real pi on PATH.
set_world() {
  WORLD=(env -i HOME="$HOME" USER="$USER" SHELL=/bin/sh LANG=en_US.UTF-8 TERM=xterm-256color
    PATH="$D/bin:/opt/homebrew/bin:/usr/bin:/bin"
    DINGGER_TMUX_SOCKET="$SOCK" DINGGER_AGENT_SOCKET="$D/agent.sock" DINGGER_STATE_DIR="$D/state"
    DINGGER_TMUX_CONFIG="$D/tmux.conf" PI_CODING_AGENT_DIR="$D/pi-agent")
}
world() { "${WORLD[@]}" "$@"; }
t() { world tmux -u -L "$SOCK" "$@"; }
term() { world tmux -u -L "$TERMSOCK" "$@"; }

cleanup_scenario() {
  [ -n "$DAEMON_PID" ] && kill "$DAEMON_PID" 2>/dev/null && wait "$DAEMON_PID" 2>/dev/null
  DAEMON_PID=
  [ -n "$SOCK" ] && t kill-server 2>/dev/null
  [ -n "$TERMSOCK" ] && term kill-server 2>/dev/null
  # A server killed with SIGKILL leaves its socket.
  [ -n "$SOCK" ] && rm -f "${TMUX_TMPDIR:-/tmp}/tmux-$(id -u)/$SOCK" "${TMUX_TMPDIR:-/tmp}/tmux-$(id -u)/$TERMSOCK"
}
cleanup() {
  cleanup_scenario
  if [ -n "${SHUTDOWN_KEEP:-}" ] || [ "$FAILS" -gt 0 ]; then echo "kept $BASE"; else rm -rf "$BASE"; fi
}
trap cleanup EXIT

start_daemon() {
  # Not through a function: $! must be the daemon itself.
  "${WORLD[@]}" RUST_LOG=debug "$DINGGER" watch --no-sounds --no-notifications >>"$D/daemon.log" 2>&1 &
  DAEMON_PID=$!
  for _ in $(seq 50); do [ -S "$D/agent.sock" ] && return; sleep 0.1; done
  echo "daemon did not start"; cat "$D/daemon.log"; exit 1
}

# A terminal opening: a client in the "terminal" server starts (or attaches
# to) the server under test, the way the user's terminal does.
open_terminal() {
  term kill-server 2>/dev/null
  world tmux -u -L "$TERMSOCK" -f /dev/null new-session -d -s term -x 160 -y 48 \
    "exec env -u TMUX tmux -u -L $SOCK -f $D/tmux.conf new-session -As 0 -c $D/w/one"
  for _ in $(seq 50); do t has-session -t =0 2>/dev/null && [ -n "$(t list-clients 2>/dev/null)" ] && return; sleep 0.1; done
  echo "the terminal's client did not attach"; exit 1
}

hook() { # pane state session-file
  python3 - "$D/agent.sock" "$1" "$2" "$3" <<'PY'
import json, socket, sys
s = socket.socket(socket.AF_UNIX)
s.connect(sys.argv[1])
s.sendall((json.dumps({"method": "pane.report_agent", "params": {"pane_id": sys.argv[2],
    "agent": "pi", "state": sys.argv[3], "agent_session_path": sys.argv[4]}}) + "\n").encode())
s.close()
PY
}

latest_snapshot() { ls "$D"/state/snapshots/*.json 2>/dev/null | tail -1; }
in_snapshot() { python3 -c 'import sys; print(sys.argv[2] in open(sys.argv[1]).read())' "$(latest_snapshot)" "$1"; }
# What restore must bring back, of the user's two sessions.
fingerprint() {
  t list-panes -a -F '#{session_name}|#{window_index}|#{window_name}|#{window_active}|#{window_zoomed_flag}|#{pane_index}|#{pane_active}|#{pane_left},#{pane_top},#{pane_width}x#{pane_height}|#{pane_current_path}' 2>/dev/null \
    | grep -E '^(0|work)\|'
}
candidate() {
  python3 -c '
import json, sys
try:
    c = json.load(open(sys.argv[1]))
except Exception:
    print("none"); sys.exit()
s = c["snapshot"]["sessions"]
panes = [p for x in s for w in x["windows"] for p in w["panes"]]
print(c["status"], sorted(x["name"] for x in s), len(panes),
      sorted((p.get("agent") or {}).get("session", "").rsplit("/", 1)[-1] for p in panes if p.get("agent")),
      sorted(p["command"] for p in panes if p.get("command")))' "$D/state/restore.json"
}
# Every pi in a test pane must be the fake.
real_agents() {
  for tty in $(t list-panes -a -F '#{pane_tty}' 2>/dev/null | sed 's|/dev/||'); do
    ps -o comm=,args= -t "$tty" | awk -v dir="$D/" '($1 ~ /(^|\/)(pi|claude)$/ || $2 ~ /^(pi|claude)$/) && index($2, dir) != 1'
  done
}
# The test's own processes: a server, and everything on its panes' ttys.
server_pid() { tmux -L "$1" display -p '#{pid}' 2>/dev/null; }
pane_pids() { for tty in $(tmux -L "$1" list-panes -a -F '#{pane_tty}' 2>/dev/null | sed 's|/dev/||'); do ps -o pid= -t "$tty"; done; }
gone() { ! kill -0 "$1" 2>/dev/null; }
wait_gone() { for _ in $(seq 100); do gone "$1" && return 0; sleep 0.1; done; return 1; }

setup() { # name
  D="$BASE/$1"
  SOCK="dgshut-$1-$$"
  TERMSOCK="dgshut-term-$1-$$"
  set_world
  mkdir -p "$D/bin/real" "$D/state" "$D/w/one" "$D/w/two" "$D/w/sessions"
  cp /bin/sleep "$D/bin/real/pi" && codesign --force --sign - "$D/bin/real/pi" 2>/dev/null
  printf '#!/bin/sh\necho "$PWD|$*" >> %s\nexec %s 100000\n' "$D/pi-args" "$D/bin/real/pi" >"$D/bin/pi"
  chmod +x "$D/bin/pi"
  # The user's shell setup, and dingger's tmux hooks, with this test's dingger.
  {
    printf 'set -g default-shell /bin/sh\nset -g default-command /bin/sh\n'
    sed -n "/^set-hook .*poke/s|\$HOME/.local/bin/dingger|$DINGGER|p" "$ROOT/packaging/tmux/dingger.tmux"
  } >"$D/tmux.conf"
  start_daemon
  open_terminal
  # Session 0 (the terminal's): an editor-like split running a command.
  t rename-window -t '=0:' edit
  t split-window -h -t '=0:edit' -c "$D/w/two"
  t send-keys -t '=0:edit.1' 'sleep 300000' Enter
  # Session work: two agents side by side, and a plain shell window.
  t new-session -d -s work -n agents -c "$D/w/one" -x 160 -y 47
  t split-window -h -t '=work:agents' -c "$D/w/two"
  t new-window -d -t '=work:' -n logs -c "$D/w/two"
  # Never the real pi: every pane that starts one must find the fake first.
  for pane in '=work:agents.0' '=work:agents.1'; do
    rm -f "$D/which-pi"
    t send-keys -t "$pane" "command -v pi > $D/which-pi" Enter
    for _ in $(seq 30); do [ -s "$D/which-pi" ] && break; sleep 0.1; done
    if [ "$(cat "$D/which-pi" 2>/dev/null)" != "$D/bin/pi" ]; then
      fail "$pane would run $(cat "$D/which-pi" 2>/dev/null), not the fake pi; stopping"
      exit 1
    fi
    t send-keys -t "$pane" 'clear; pi' Enter
  done
  A=$(t display -p -t '=work:agents.0' '#{pane_id}')
  B=$(t display -p -t '=work:agents.1' '#{pane_id}')
  touch "$D/w/sessions/a.jsonl" "$D/w/sessions/b.jsonl"
  expect "setup: both fake agents run" "t display -p -t $A '#{pane_current_command}'; t display -p -t $B '#{pane_current_command}'" "pi
pi" 50 || exit 1
  hook "$A" working "$D/w/sessions/a.jsonl"
  hook "$B" idle "$D/w/sessions/b.jsonl"
  expect "setup: the snapshot has both agents and the command" \
    "in_snapshot a.jsonl; in_snapshot b.jsonl; in_snapshot 'sleep 300000'" "True
True
True" 80 || exit 1
  BEFORE=$(fingerprint)
  # The state the user left: looking at session 0.
  WANT="offered ['0', 'work'] 5 ['a.jsonl', 'b.jsonl'] ['pi', 'pi', 'sleep 300000']"
}

# After the shutdown: the server is gone, the layout is offered, and a
# login (fresh daemon, terminal starting tmux) restores it.
after_shutdown() { # server-pid
  wait_gone "$1" && pass "the tmux server is gone" || { fail "the tmux server outlived the shutdown"; kill -KILL "$1" 2>/dev/null; }
  if [ -n "$DAEMON_PID" ] && gone "$DAEMON_PID"; then
    wait "$DAEMON_PID" 2>/dev/null
    DAEMON_PID=
  fi
  # Login: launchd starts the daemon again (if the shutdown stopped it).
  [ -z "$DAEMON_PID" ] && start_daemon
  expect "the old layout is the restore candidate" candidate "$WANT" 80
  open_terminal
  expect "the new server has a bare 0" "t list-windows -t =0 -F '#{window_panes}' | tr '\n' ' '" "1 "
  expect "the offer: 0 and work, 2 agents" \
    "python3 -c 'import json,sys; o=json.load(open(sys.argv[1])); print([s[\"name\"] for s in o[\"sessions\"]], o[\"agents\"])' $D/state/offer.json 2>/dev/null" \
    "['0', 'work'] 2" 80
  world "$DINGGER" restore >"$D/restore.out" 2>&1
  expect "restore succeeds" "grep -c FAILED $D/restore.out" 0
  expect "every window, pane, layout and directory is back" fingerprint "$BEFORE" 50
  expect "both agents resumed in their own directories and sessions" "sort $D/pi-args | grep -- --session" \
    "$D/w/one|--session $D/w/sessions/a.jsonl
$D/w/two|--session $D/w/sessions/b.jsonl" 50
  expect "the command is typed, not run" \
    "t capture-pane -p -t '=0:edit.1' | grep -c 'sleep 300000\$'; t display -p -t '=0:edit.1' '#{pane_current_command}'" "1
bash"
  expect "no real agent runs in a test pane" real_agents ""
  [ "$FAILS" -gt 0 ] && [ -n "${SHUTDOWN_VERBOSE:-}" ] && { echo "--- restore:"; cat "$D/restore.out"; }
}

scenario() { # name description stop-function
  echo "$2"
  setup "$1"
  local server
  server=$(server_pid "$SOCK")
  "$3" "$server"
  after_shutdown "$server"
  cleanup_scenario
}

# 1. launchd reaches tmux first.
term_tmux() { kill -TERM "$1"; }
# 2. The agents die first; tmux a few seconds later.
agents_first() {
  kill -TERM $(pgrep -f "^$D/bin/real/pi") 2>/dev/null
  sleep 3
  kill -TERM "$1"
}
# 3./4. Everything at once, in a random order: tmux, everything in its panes,
# the terminal's client, the daemon. What ignores SIGTERM (an interactive
# shell) is killed a few seconds later, as launchd does.
everything() { # signal server
  local pids
  pids=$( { echo "$2"; pane_pids "$SOCK"; server_pid "$TERMSOCK"; pane_pids "$TERMSOCK"; echo "$DAEMON_PID"; } | awk 'NF' | sort -u | python3 -c 'import random,sys; p=sys.stdin.read().split(); random.shuffle(p); print(" ".join(p))')
  # A random moment in the daemon's 1 s cycle.
  sleep "0.$((RANDOM % 10))"
  kill -"$1" $pids 2>/dev/null
  sleep 2
  kill -KILL $pids 2>/dev/null
}
term_everything() { everything TERM "$1"; }
kill_everything() { everything KILL "$1"; }
# 5. The daemon goes first, then tmux.
daemon_first() {
  kill -TERM "$DAEMON_PID"
  sleep 1
  kill -TERM "$1"
}
# 6. The terminal quits first (its client goes), then launchd stops tmux.
terminal_quits() {
  kill -HUP $(pane_pids "$TERMSOCK") 2>/dev/null
  term kill-server 2>/dev/null
  expect "the terminal's client is gone, the server runs on" "t list-clients | wc -l | tr -d ' '; t has-session -t =work && echo up" "0
up" 50
  sleep 2
  kill -TERM "$1"
}

scenario term-tmux "1. SIGTERM to tmux" term_tmux
scenario agents-first "2. SIGTERM to the agents, to tmux 3 s later" agents_first
scenario daemon-first "5. the daemon stopped first, then tmux" daemon_first
scenario terminal-quits "6. the terminal quits, then SIGTERM to tmux" terminal_quits
for ((i = 1; i <= ROUNDS; i++)); do
  scenario "term-all-$i" "3. SIGTERM to everything at once, round $i of $ROUNDS" term_everything
done
for ((i = 1; i <= ROUNDS; i++)); do
  scenario "kill-all-$i" "4. SIGKILL to everything at once, round $i of $ROUNDS" kill_everything
done

echo "$PASSES passed, $FAILS failed"
[ "$FAILS" -eq 0 ]