#!/usr/bin/env bash
# claude-sessions-shutdown — Gracefully exit all Claude Code tmux sessions.
#
# Called by the systemd user service before halt/reboot.
# Sends /exit-session to each Claude tmux session, waits for them to close,
# then force-kills any stragglers.
#
# Usage:
#   claude-sessions-shutdown [--timeout N]   # graceful timeout per session (default: 90)
#   claude-sessions-shutdown --force         # skip graceful exit, force-kill immediately

set -uo pipefail

REGISTRY="${HOME}/.local/share/claude-sessions/registry.json"
TIMEOUT=90
FORCE=0
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[0;33m'; NC='\033[0m'

log()  { echo -e "$(date +%T) [claude-shutdown] $*"; }
warn() { echo -e "$(date +%T) [claude-shutdown] ${YELLOW}$*${NC}"; }
err()  { echo -e "$(date +%T) [claude-shutdown] ${RED}$*${NC}" >&2; }

while [[ $# -gt 0 ]]; do
    case "$1" in
        --timeout) TIMEOUT="${2:-90}"; shift 2 ;;
        --force)   FORCE=1; shift ;;
        *) err "Unknown flag: $1"; exit 1 ;;
    esac
done

# ── Collect Claude sessions ──────────────────────────────────────────────────

collect_claude_sessions() {
    local seen=()
    local sessions=()

    # Source 1: registry entries with a live tmux session
    if [[ -f "$REGISTRY" ]]; then
        while IFS= read -r name; do
            [[ -n "$name" ]] || continue
            if tmux has-session -t "=$name" 2>/dev/null; then
                sessions+=("$name")
                seen+=("$name")
            fi
        done < <(python3 -c "
import json, sys
try:
    d = json.load(open('$REGISTRY'))
    for s in d.get('sessions', []):
        n = s.get('tmux_name', '')
        if n: print(n)
except Exception: pass
" 2>/dev/null)
    fi

    # Source 2: tmux sessions whose pane runs a node/claude child process
    while IFS= read -r line; do
        session="${line%% *}"
        pane_pid="${line##* }"
        [[ -n "$session" && -n "$pane_pid" ]] || continue

        # Already captured from registry?
        skip=0
        for s in "${seen[@]+"${seen[@]}"}"; do [[ "$s" == "$session" ]] && { skip=1; break; }; done
        (( skip )) && continue

        # Check if any child of the pane PID is node (Claude Code)
        if pgrep -P "$pane_pid" -x node &>/dev/null 2>&1 || \
           pgrep -a -P "$pane_pid" 2>/dev/null | grep -qE 'claude|\.claude'; then
            sessions+=("$session")
        fi
    done < <(tmux list-panes -a -F '#{session_name} #{pane_pid}' 2>/dev/null)

    printf '%s\n' "${sessions[@]+"${sessions[@]}"}"
}

# ── Main ─────────────────────────────────────────────────────────────────────

if ! command -v tmux &>/dev/null; then
    log "tmux not found — nothing to do"
    exit 0
fi

if ! tmux ls &>/dev/null 2>&1; then
    log "No tmux server running — nothing to do"
    exit 0
fi

mapfile -t claude_sessions < <(collect_claude_sessions)

if [[ ${#claude_sessions[@]} -eq 0 ]]; then
    log "No Claude tmux sessions found — nothing to do"
    exit 0
fi

log "Found ${#claude_sessions[@]} Claude session(s): ${claude_sessions[*]}"

if (( FORCE )); then
    log "Force mode: killing sessions immediately"
    for name in "${claude_sessions[@]}"; do
        tmux kill-session -t "=$name" 2>/dev/null && log "Killed: $name" || warn "Could not kill: $name"
    done
    exit 0
fi

# Send /exit-session to each Claude session
for name in "${claude_sessions[@]}"; do
    log "Sending /exit-session to: $name"
    # Clear any partial input, then send the command
    tmux send-keys -t "=$name" "" ""  2>/dev/null || true
    tmux send-keys -t "=$name" -l "/exit-session" 2>/dev/null || true
    tmux send-keys -t "=$name" Enter 2>/dev/null || true
done

# Wait for sessions to close
log "Waiting up to ${TIMEOUT}s for sessions to exit…"
deadline=$(( $(date +%s) + TIMEOUT ))
remaining=("${claude_sessions[@]}")

while [[ ${#remaining[@]} -gt 0 && $(date +%s) -lt $deadline ]]; do
    still_open=()
    for name in "${remaining[@]}"; do
        if tmux has-session -t "=$name" 2>/dev/null; then
            still_open+=("$name")
        else
            echo -e "  ${GREEN}✓${NC} $name exited cleanly"
        fi
    done
    remaining=("${still_open[@]+"${still_open[@]}"}")
    [[ ${#remaining[@]} -gt 0 ]] && sleep 2
done

# Force-kill stragglers
if [[ ${#remaining[@]} -gt 0 ]]; then
    warn "${#remaining[@]} session(s) did not exit gracefully — force-killing"
    for name in "${remaining[@]}"; do
        tmux kill-session -t "=$name" 2>/dev/null \
            && warn "Force-killed: $name" \
            || err "Failed to kill: $name"
    done
fi

log "Shutdown complete"