#!/bin/bash
# aria-worker-stop.sh — Graceful stop + worktree cleanup for an Aria worker
#
# Usage: aria-worker-stop.sh <name>
#
# Stops the tmux session "aria-<name>", kills the auto-approve loop and watchdog,
# removes the git worktree (if applicable), and archives the state dir.
# NEVER touches the running "aria" CEO session or aria.service.
#
# SPEC: /root/aria/state/ai-company-plan/SPEC-FINAL.md (Step 7)

set -euo pipefail

# ── Constants ────────────────────────────────────────────────────────────────
WORKERS_STATE="/root/aria/state/workers"
STOP_LOG="/var/log/aria-worker-stop.log"

# ── Logging ─────────────────────────────────────────────────────────────────
log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') [aria-worker-stop] $1" | tee -a "$STOP_LOG" 2>/dev/null || true
}

die() {
    log "FATAL: $1"
    echo "ERROR: $1" >&2
    exit 1
}

# ── PATH (may be invoked from watchdog / cron without full profile) ──────────
export PATH="/root/.local/bin:/usr/local/bin:/usr/bin:/bin:$PATH"

# ── Step 1: Arg validation ────────────────────────────────────────────────────
[ $# -eq 1 ] || die "Usage: aria-worker-stop.sh <name>"
WORKER_NAME="$1"

# Name syntax check
[[ "$WORKER_NAME" =~ ^[a-zA-Z0-9][a-zA-Z0-9-]*$ ]] \
    || die "Invalid worker name: '$WORKER_NAME'"

# ── Step 2: CEO-session guard (HARD CONSTRAINT) ──────────────────────────────
# Resolve the target session name and refuse if it would be the CEO session.
WORKER_SESSION="aria-${WORKER_NAME}"

# GUARD #1 (belt): resolved session name must not equal bare "aria"
if [[ "$WORKER_SESSION" == "aria" ]]; then
    die "HARD CONSTRAINT: resolved session name is 'aria' — will NEVER stop the CEO session. Use aria-restart.sh instead."
fi

# GUARD #2 (suspenders): session name must start with "aria-"
if [[ "$WORKER_SESSION" != aria-* ]]; then
    die "HARD CONSTRAINT: session name '$WORKER_SESSION' does not start with 'aria-'. Refusing to proceed."
fi

# GUARD #3 (paranoia): name must not be empty string
if [[ -z "$WORKER_NAME" ]]; then
    die "HARD CONSTRAINT: empty worker name would target CEO session. Aborting."
fi

log "Stopping worker '$WORKER_NAME' (session: $WORKER_SESSION)"

# ── Step 3: Read state dir ────────────────────────────────────────────────────
WORKER_STATE_DIR="${WORKERS_STATE}/${WORKER_NAME}"
WORKER_LOG="${WORKER_STATE_DIR}/worker.log"

if [ ! -d "$WORKER_STATE_DIR" ]; then
    log "WARNING: State dir not found: $WORKER_STATE_DIR — will still attempt to kill session"
fi

REPO=""
WORKTREE_PATH=""
WORKTREE_BRANCH=""
WORK_DIR=""
APPROVE_PID=""

if [ -d "$WORKER_STATE_DIR" ]; then
    REPO=$(cat "${WORKER_STATE_DIR}/repo" 2>/dev/null || echo "")
    WORKTREE_PATH=$(cat "${WORKER_STATE_DIR}/worktree" 2>/dev/null || echo "")
    WORKTREE_BRANCH=$(cat "${WORKER_STATE_DIR}/branch" 2>/dev/null || echo "")
    WORK_DIR=$(cat "${WORKER_STATE_DIR}/workdir" 2>/dev/null || echo "")
    APPROVE_PID=$(cat "${WORKER_STATE_DIR}/approve_pid" 2>/dev/null || echo "")
fi

# ── Step 4: Stop the auto-approve loop ───────────────────────────────────────
if [ -n "$APPROVE_PID" ] && kill -0 "$APPROVE_PID" 2>/dev/null; then
    kill "$APPROVE_PID" 2>/dev/null || true
    log "Auto-approve loop (PID $APPROVE_PID) killed"
else
    log "Auto-approve loop not running (PID: ${APPROVE_PID:-none})"
fi

# Clear approve_pid file
[ -d "$WORKER_STATE_DIR" ] && rm -f "${WORKER_STATE_DIR}/approve_pid" 2>/dev/null || true

# ── Step 5: Kill any orphaned watchdog for this worker ───────────────────────
# Watchdog is: sleep $N; bash /root/aria/scripts/aria-worker-stop.sh <name>
# We kill the watchdog's sleep by finding the disowned subshell process group.
# This is best-effort: the watchdog may already have fired or been orphaned.
WATCHDOG_PIDS=$(pgrep -f "aria-worker-stop.sh ${WORKER_NAME}" 2>/dev/null | grep -v "^$$\$" || true)
if [ -n "$WATCHDOG_PIDS" ]; then
    for wpid in $WATCHDOG_PIDS; do
        [ "$wpid" = "$$" ] && continue
        kill "$wpid" 2>/dev/null || true
        log "Killed watchdog process (PID $wpid)"
    done
fi

# ── Step 6: Graceful stop — send /exit to Claude ─────────────────────────────
if tmux has-session -t "$WORKER_SESSION" 2>/dev/null; then
    log "Sending /exit to Claude in session $WORKER_SESSION..."
    tmux send-keys -t "$WORKER_SESSION" "/exit" Enter 2>/dev/null || true

    # Wait up to 10s for graceful exit
    for i in $(seq 1 10); do
        sleep 1
        if ! tmux has-session -t "$WORKER_SESSION" 2>/dev/null; then
            log "Session exited gracefully after ${i}s"
            break
        fi
    done

    # Force kill if still alive after 10s
    if tmux has-session -t "$WORKER_SESSION" 2>/dev/null; then
        log "Session did not exit gracefully after 10s — force killing"
        tmux kill-session -t "$WORKER_SESSION" 2>/dev/null || true
        log "tmux kill-session '$WORKER_SESSION' done"
    fi
else
    log "Session $WORKER_SESSION not found (already stopped or never started)"
fi

# Final sanity check: CEO session must still be alive (non-fatal warning only)
if ! tmux has-session -t aria 2>/dev/null; then
    log "WARNING: CEO session 'aria' is not running — this script did NOT touch it"
fi

# ── Step 7: Remove git worktree ───────────────────────────────────────────────
if [ -n "$REPO" ] && [ -n "$WORKTREE_PATH" ] && [ -d "$WORKTREE_PATH" ]; then
    log "Removing git worktree: $WORKTREE_PATH"
    git -C "$REPO" worktree remove --force "$WORKTREE_PATH" 2>/dev/null \
        || log "WARNING: git worktree remove failed — path may already be removed"
    git -C "$REPO" worktree prune 2>/dev/null || true
    log "Worktree removed and git pruned"

    # Delete worker branch ONLY if it was never pushed to remote
    if [ -n "$WORKTREE_BRANCH" ]; then
        if git -C "$REPO" ls-remote --exit-code origin "$WORKTREE_BRANCH" >/dev/null 2>&1; then
            log "Branch '$WORKTREE_BRANCH' exists on remote — leaving it (manual cleanup if needed)"
        else
            git -C "$REPO" branch -D "$WORKTREE_BRANCH" 2>/dev/null \
                && log "Branch '$WORKTREE_BRANCH' deleted (was not pushed)" \
                || log "WARNING: Could not delete local branch '$WORKTREE_BRANCH' (may already be gone)"
        fi
    fi
elif [ -n "$WORKTREE_PATH" ] && [ -d "$WORKTREE_PATH" ]; then
    log "Worktree path exists but no repo recorded — removing directory directly"
    rm -rf "$WORKTREE_PATH" 2>/dev/null \
        && log "Removed orphaned worktree directory: $WORKTREE_PATH" \
        || log "WARNING: Could not remove $WORKTREE_PATH"
elif [ -n "$WORKTREE_PATH" ]; then
    log "Worktree path $WORKTREE_PATH no longer exists (already cleaned up)"
fi

# ── Step 8: Archive state dir ─────────────────────────────────────────────────
if [ -d "$WORKER_STATE_DIR" ]; then
    ARCHIVE_TS=$(date +%s)
    ARCHIVE_PATH="${WORKER_STATE_DIR}.stopped.${ARCHIVE_TS}"
    mv "$WORKER_STATE_DIR" "$ARCHIVE_PATH" 2>/dev/null \
        && log "State dir archived: $ARCHIVE_PATH" \
        || log "WARNING: Could not archive state dir $WORKER_STATE_DIR"
else
    log "State dir already absent — nothing to archive"
fi

# ── Done ──────────────────────────────────────────────────────────────────────
echo ""
echo "============================================="
echo "  ARIA WORKER STOPPED"
echo "============================================="
echo "  Name:      $WORKER_NAME"
echo "  Session:   $WORKER_SESSION (terminated)"
[ -n "$WORKTREE_PATH" ] && echo "  Worktree:  $WORKTREE_PATH (removed)"
[ -n "$WORKTREE_BRANCH" ] && echo "  Branch:    $WORKTREE_BRANCH"
echo "  CEO session 'aria' was NOT touched."
echo "============================================="

log "Worker '$WORKER_NAME' fully stopped. CEO session 'aria' untouched."
exit 0