#!/bin/bash
# aria-worker.sh — Spawn a scoped Aria worker session
#
# Usage: aria-worker.sh <name> <project-dir-or-repo> <role>
#
# Spawns tmux session "aria-<name>" running Claude Code scoped to the project.
# If project-dir is a git repo: creates an isolated git worktree + branch.
# Does NOT touch or affect the running "aria" CEO session (triple-guarded).
#
# Phase 1: no Telegram channel. Workers driven via panel/tmux only.
# SPEC: /root/aria/state/ai-company-plan/SPEC-FINAL.md (Step 6)
# All CF-1 through CF-12 fixes are baked in.

set -euo pipefail

# ── Constants ───────────────────────────────────────────────────────────────
ARIA_STATE="/root/aria/state"
WORKERS_STATE="${ARIA_STATE}/workers"
WORKTREES_BASE="${ARIA_STATE}/worktrees"
ROLES_DIR="/root/aria/brain/workers/roles"
ENV_FILE="/root/aria/.env"
ENV_ARIA="/root/aria/scripts/.env.aria"
SPAWN_LOG="/var/log/aria-worker-spawn.log"

# Default MAX_RUNTIME: 3600s (1 hour). Override via env var (CF-8).
MAX_RUNTIME="${ARIA_WORKER_MAX_RUNTIME:-3600}"

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

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

# ── PATH (tmux doesn't inherit user profile) ─────────────────────────────────
export PATH="/root/.local/bin:/usr/local/bin:/usr/bin:/bin:$PATH"

# ── Step 1: Argument validation ──────────────────────────────────────────────
[ $# -eq 3 ] || die "Usage: aria-worker.sh <name> <project-dir-or-repo> <role>"

WORKER_NAME="$1"
PROJECT_DIR="$2"
ROLE="$3"

# Name must match regex ^[a-zA-Z0-9][a-zA-Z0-9-]*$
[[ "$WORKER_NAME" =~ ^[a-zA-Z0-9][a-zA-Z0-9-]*$ ]] \
    || die "name must be alphanumeric+dash (no leading dash), got: '$WORKER_NAME'"

# Name must not be bare "aria" (would form session "aria-aria" — that's fine, but
# the double-guard below also checks the resolved session name explicitly)
[ "$WORKER_NAME" != "aria" ] \
    || die "Worker name 'aria' disallowed — would be ambiguous."

# Role: alphanumeric + dash + underscore
[[ "$ROLE" =~ ^[a-zA-Z0-9][a-zA-Z0-9_-]*$ ]] \
    || die "role must be alphanumeric+dash+underscore, got: '$ROLE'"

# Role file must exist (unknown role guard)
ROLE_FILE="${ROLES_DIR}/${ROLE}.md"
[ -f "$ROLE_FILE" ] \
    || die "Unknown role '$ROLE': role file not found at $ROLE_FILE"

# Project dir must exist and be absolute
[[ "$PROJECT_DIR" == /* ]] || die "project-dir must be an absolute path, got: '$PROJECT_DIR'"
[ -d "$PROJECT_DIR" ] || die "project-dir does not exist: '$PROJECT_DIR'"

# Canonicalize and restrict to /root for safety
REAL_PROJECT_DIR=$(realpath "$PROJECT_DIR" 2>/dev/null) \
    || die "Cannot canonicalize: '$PROJECT_DIR'"
[[ "$REAL_PROJECT_DIR" == /root/* || "$REAL_PROJECT_DIR" == /root ]] \
    || die "project-dir must be inside /root for safety, got: '$REAL_PROJECT_DIR'"
PROJECT_DIR="$REAL_PROJECT_DIR"

# ── Step 2: Resolve worker session name + TRIPLE DASH-GUARD ─────────────────
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 target CEO session. Aborting."
fi

# GUARD #2 (suspenders): session name must contain a dash
if [[ "$WORKER_SESSION" != *-* ]]; then
    die "HARD CONSTRAINT: session name '$WORKER_SESSION' contains no dash — refusing to proceed."
fi

# GUARD #3 (triple): session name must start with exactly "aria-" (aria + dash)
if [[ "$WORKER_SESSION" != aria-* ]]; then
    die "HARD CONSTRAINT: session name '$WORKER_SESSION' does not start with 'aria-'. Aborting."
fi

log "Spawning worker '$WORKER_NAME' | project: '$PROJECT_DIR' | role: '$ROLE'"
log "Target tmux session: $WORKER_SESSION"

# ── Step 3: State directory ───────────────────────────────────────────────────
WORKER_STATE_DIR="${WORKERS_STATE}/${WORKER_NAME}"
mkdir -p "$WORKER_STATE_DIR"
WORKER_LOG="${WORKER_STATE_DIR}/worker.log"
WORKER_STDERR_LOG="${WORKER_STATE_DIR}/stderr.log"

# Redefine log to also write to the per-worker log
log() {
    local msg
    msg="$(date '+%Y-%m-%d %H:%M:%S') [aria-worker/$WORKER_NAME] $1"
    echo "$msg" | tee -a "$SPAWN_LOG" >> "$WORKER_LOG" 2>/dev/null || true
}

# ── Step 4: Concurrent worker cap (CF-3) ────────────────────────────────────
# Count ONLY active tmux sessions starting with "aria-" (not archived state dirs)
MAX=$(cat "${WORKERS_STATE}/max_concurrent" 2>/dev/null || echo 5)
ACTIVE=$(tmux list-sessions -F '#{session_name}' 2>/dev/null | grep -c '^aria-' || echo 0)
if [[ "$ACTIVE" -ge "$MAX" ]]; then
    die "Max concurrent workers ($MAX) reached — $ACTIVE aria-* sessions running. (CF-3)"
fi

# ── Step 5: Idempotency — abort if session already exists ────────────────────
if tmux has-session -t "$WORKER_SESSION" 2>/dev/null; then
    die "Session '$WORKER_SESSION' is already running. Use aria-worker-stop.sh $WORKER_NAME first."
fi

# ── Step 6: Lock file — prevent double-spawn (CF-12) ────────────────────────
LOCK_FILE="/tmp/aria-worker-spawn-${WORKER_NAME}.lock"
exec 9>"$LOCK_FILE"
flock -n 9 || die "Another spawn of '$WORKER_NAME' is already in progress (CF-12 lock held: $LOCK_FILE)"

# ── Step 7: Git worktree setup (if project is a git repo) ────────────────────
IS_GIT=0
WORK_DIR="$PROJECT_DIR"
WORKTREE_PATH=""
WORKTREE_BRANCH=""
REPO_ROOT=""

if git -C "$PROJECT_DIR" rev-parse --git-dir >/dev/null 2>&1; then
    IS_GIT=1
    REPO_ROOT=$(git -C "$PROJECT_DIR" rev-parse --show-toplevel 2>/dev/null)
    log "Git repo detected: $REPO_ROOT"

    # Defensive prune before adding (handles stale worktrees from prior crashes)
    git -C "$REPO_ROOT" worktree prune 2>/dev/null || true

    TIMESTAMP=$(date +%Y%m%d%H%M%S)
    WORKTREE_BRANCH="worker/${WORKER_NAME}/${TIMESTAMP}"

    mkdir -p "$WORKTREES_BASE"
    REPO_BASENAME=$(basename "$REPO_ROOT")
    WORKTREE_PATH="${WORKTREES_BASE}/${REPO_BASENAME}-${WORKER_NAME}-${TIMESTAMP}"

    # Determine base branch (prefer origin/main, fallback origin/master, fallback HEAD)
    BASE_REF="HEAD"
    if git -C "$REPO_ROOT" rev-parse origin/main >/dev/null 2>&1; then
        BASE_REF="origin/main"
    elif git -C "$REPO_ROOT" rev-parse origin/master >/dev/null 2>&1; then
        BASE_REF="origin/master"
    fi
    log "Creating git worktree: $WORKTREE_PATH | branch: $WORKTREE_BRANCH | base: $BASE_REF"

    git -C "$REPO_ROOT" worktree add -b "$WORKTREE_BRANCH" "$WORKTREE_PATH" "$BASE_REF" \
        || die "git worktree add failed. Stale path? Try: git -C $REPO_ROOT worktree prune"

    WORK_DIR="$WORKTREE_PATH"
    log "Worktree ready: $WORK_DIR"
else
    log "Not a git repo — using project-dir as-is: $WORK_DIR"
fi

# ── Step 8: Source env — then IMMEDIATELY unset ANTHROPIC_API_KEY (CF-2) ────
# Workers MUST bill OAuth (Max plan), not the API key.
# Order: source first, then unset — keeps CLAUDE_CODE_OAUTH_TOKEN, drops the key.
# shellcheck source=/dev/null
source "$ENV_FILE" 2>/dev/null || true
# CF-2: MUST happen immediately after source — prevents API-key billing in workers
unset ANTHROPIC_API_KEY
# Source additional env (no API keys expected here, but unset again defensively)
# shellcheck source=/dev/null
source "$ENV_ARIA" 2>/dev/null || true
unset ANTHROPIC_API_KEY  # belt-and-suspenders: unset again after second source

# Ensure OAuth token is exported for the claude process
export CLAUDE_CODE_OAUTH_TOKEN

# ── Step 9: CF-1 — Export ARIA_WORKER_NAME BEFORE tmux new-session ──────────
# session-start.sh hook reads this env var and skips CEO brain build + Telegram ping.
# MUST be exported BEFORE tmux new-session so the hook fires with this var set.
export ARIA_WORKER_NAME="${WORKER_NAME}"

# ── Step 10: Read role file for default_mode and startup prompt ───────────────
# Extract default_mode from frontmatter
ROLE_DEFAULT_MODE=$(grep '^default_mode:' "$ROLE_FILE" 2>/dev/null \
    | head -1 | sed 's/default_mode:[[:space:]]*//' | tr -d '[:space:]"' || echo "assist")

# Validate extracted mode — fail-safe to assist
case "$ROLE_DEFAULT_MODE" in
    assist|auto|yolo) ;;
    *) ROLE_DEFAULT_MODE="assist" ;;
esac

# Extract startup prompt between ---START PROMPT--- and ---END PROMPT---
STARTUP_PROMPT_RAW=$(awk '/^---START PROMPT---$/{found=1; next} /^---END PROMPT---$/{found=0} found{print}' \
    "$ROLE_FILE" 2>/dev/null || true)

if [ -z "$STARTUP_PROMPT_RAW" ]; then
    log "WARNING: No ---START PROMPT---/---END PROMPT--- section found in $ROLE_FILE — using generic prompt"
    STARTUP_PROMPT_RAW="You are an Aria worker agent with role: ${ROLE}. Working directory: ${WORK_DIR}. Work autonomously. Report only when blocked or done."
fi

# Substitute <name> placeholder in prompt
STARTUP_PROMPT_RAW=$(echo "$STARTUP_PROMPT_RAW" | sed "s|<name>|${WORKER_NAME}|g")

# Sanitize: strip carriage returns and null bytes (Risk 4 mitigation)
STARTUP_PROMPT=$(printf '%s' "$STARTUP_PROMPT_RAW" | tr -d '\r\000')

# Append auto-resume instruction (Step 5.3 of SPEC)
STARTUP_PROMPT="${STARTUP_PROMPT}

WICHTIG AUTO-RESUME:
1. Lies /root/aria/state/workers/${WORKER_NAME}/active-task.md.
2. Falls Blockers-Feld leer und Next Action gesetzt: fuehre SOFORT aus, kein Warten.
3. Pruefe 'Done So Far' bevor du irreversible Aktionen wiederholst: TOKEN schon eingetragen = ueberspringen.
4. Melde Kais nur wenn: Blockers gesetzt, Task done, oder kein active-task.md.

Working directory: ${WORK_DIR}
$([ "$IS_GIT" = "1" ] && echo "Git branch: ${WORKTREE_BRANCH}" || true)
WORKER_NAME: ${WORKER_NAME} | ROLE: ${ROLE}"

# ── Step 11: Write mode file and active-task.md BEFORE approve loop (CF-9) ──
# CF-9: mode file MUST be written before the approve loop starts.
echo "${ROLE_DEFAULT_MODE}" > "${WORKER_STATE_DIR}/mode"
log "Mode file written: ${ROLE_DEFAULT_MODE}"

# Write initial active-task.md scaffold (auto-resume foundation)
cat > "${WORKER_STATE_DIR}/active-task.md" << ACTIVETASK
---
schema: active-task-v1
updated: $(date -u '+%Y-%m-%dT%H:%M:%SZ')
session_id: (pending)
kar_issue:
worker_name: ${WORKER_NAME}
---

## Goal
(Set by Kais via tmux send-keys after spawn)

## Current Step
Step 0/N: Initializing

## Next Action
Read task brief and begin Step 1.

## Done So Far (idempotency log)
<!-- APPEND-ONLY. Format: [ISO-timestamp] ACTION_TOKEN: description -->
<!-- grep for ACTION_TOKEN before re-running any irreversible action -->

## Relevant Files / Branch
- branch: ${WORKTREE_BRANCH:-}
- repo: ${REPO_ROOT:-}
- worktree: ${WORKTREE_PATH:-}

## Blockers
ACTIVETASK
log "active-task.md written"

# Write initial supporting state files
echo "$WORKER_NAME"   > "${WORKER_STATE_DIR}/name"
echo "$ROLE"          > "${WORKER_STATE_DIR}/role"
echo "$PROJECT_DIR"   > "${WORKER_STATE_DIR}/scope"
echo "$WORK_DIR"      > "${WORKER_STATE_DIR}/workdir"
echo ""               > "${WORKER_STATE_DIR}/pid"
echo ""               > "${WORKER_STATE_DIR}/approve_pid"

if [ "$IS_GIT" = "1" ]; then
    echo "$REPO_ROOT"       > "${WORKER_STATE_DIR}/repo"
    echo "$WORKTREE_PATH"   > "${WORKER_STATE_DIR}/worktree"
    echo "$WORKTREE_BRANCH" > "${WORKER_STATE_DIR}/branch"
fi

# ── Step 12: Build worker brain snapshot ─────────────────────────────────────
# Derives vault slug from role file frontmatter (defaults to "core")
VAULT_SLUG=$(grep '^vault:' "$ROLE_FILE" 2>/dev/null \
    | head -1 | sed 's/vault:[[:space:]]*//' | tr -d '[:space:]"' || echo "core")
log "Building brain snapshot (vault: $VAULT_SLUG)..."
bash /root/aria/scripts/aria-worker-start.sh "$WORKER_NAME" "$VAULT_SLUG" \
    >> "$WORKER_LOG" 2>&1 || log "WARNING: aria-worker-start.sh failed — worker will start without snapshot"

# ── Step 13: Spawn tmux session ───────────────────────────────────────────────
# Note: ARIA_WORKER_NAME is already exported (Step 9) — it will be inherited by
# the shell that starts claude, so the session-start hook fires with the var set.
log "Spawning tmux session: $WORKER_SESSION in $WORK_DIR"

# Build the claude command. No --channels flag (Phase 1 = panel/tmux only).
CLAUDE_CMD="claude --add-dir $(printf '%q' "$WORK_DIR") 2>>'$(printf '%s' "$WORKER_STDERR_LOG")'"

tmux new-session -d -s "$WORKER_SESSION" -c "$WORK_DIR" \
    "bash -c 'export PATH=/root/.local/bin:/usr/local/bin:/usr/bin:/bin:\$PATH; \
              source /root/aria/.env 2>/dev/null || true; \
              unset ANTHROPIC_API_KEY; \
              export CLAUDE_CODE_OAUTH_TOKEN; \
              export ARIA_WORKER_NAME=$(printf '%q' "$WORKER_NAME"); \
              ${CLAUDE_CMD}; \
              echo \"[worker exited \$?]\" >> $(printf '%q' "$WORKER_LOG"); \
              sleep 60'"

log "tmux session $WORKER_SESSION started"

# ── Step 14: Start auto-approve loop AFTER mode file is written (CF-9) ───────
bash /root/aria/scripts/aria-worker-auto-approve.sh \
    "$WORKER_NAME" "$WORKER_SESSION" >> "$WORKER_LOG" 2>&1 &
APPROVE_PID=$!
echo "$APPROVE_PID" > "${WORKER_STATE_DIR}/approve_pid"
log "Auto-approve loop started (PID $APPROVE_PID)"

# ── Step 15: Wall-clock watchdog (CF-8) ──────────────────────────────────────
# Kills the worker after MAX_RUNTIME seconds regardless of state.
# disown prevents watchdog from dying if this parent shell exits.
(sleep "$MAX_RUNTIME"; bash /root/aria/scripts/aria-worker-stop.sh "${WORKER_NAME}" \
    >> "${WORKER_STATE_DIR}/worker.log" 2>&1) &
WATCHDOG_PID=$!
disown "$WATCHDOG_PID"
log "Watchdog started (PID $WATCHDOG_PID): worker auto-stops after ${MAX_RUNTIME}s (CF-8)"

# ── Step 16: Wait for Claude to be ready (❯ prompt) ──────────────────────────
log "Waiting for Claude prompt (max 120s)..."
READY=0
for i in $(seq 1 60); do
    sleep 2
    if ! tmux has-session -t "$WORKER_SESSION" 2>/dev/null; then
        log "ERROR: Session $WORKER_SESSION died before becoming ready"
        kill "$APPROVE_PID" 2>/dev/null || true
        # Clean up worktree on failure
        if [ "$IS_GIT" = "1" ] && [ -n "$WORKTREE_PATH" ]; then
            git -C "$REPO_ROOT" worktree remove --force "$WORKTREE_PATH" 2>/dev/null || true
            git -C "$REPO_ROOT" worktree prune 2>/dev/null || true
        fi
        die "Worker session '$WORKER_SESSION' died during startup"
    fi

    SCREEN=$(tmux capture-pane -t "$WORKER_SESSION" -p 2>/dev/null | sed 's/\x1b\[[0-9;]*m//g')

    # Handle trust dialog if it appears
    if echo "$SCREEN" | grep -qi "trust this folder\|trust.*directory"; then
        tmux send-keys -t "$WORKER_SESSION" "" Enter 2>/dev/null || true
        log "Trust dialog confirmed at iteration $i"
        sleep 2
        continue
    fi

    if echo "$SCREEN" | grep -q "^❯"; then
        READY=1
        log "Claude prompt ready after $((i * 2))s"
        break
    fi
done

if [ "$READY" = "0" ]; then
    log "WARNING: ❯ prompt not detected after 120s — session may still be loading. Injecting prompt anyway."
fi

# ── Step 17: Inject startup prompt via tmux send-keys -l ────────────────────
# -l flag passes the string literally (no special key interpretation).
# Small stabilisation sleep before inject.
sleep 2
log "Injecting startup prompt..."
tmux send-keys -l -t "$WORKER_SESSION" "$STARTUP_PROMPT" 2>/dev/null || true
tmux send-keys -t "$WORKER_SESSION" "" Enter 2>/dev/null || true
log "Startup prompt sent"

# ── Step 18: Record Claude PID (best-effort) ─────────────────────────────────
CLAUDE_PID=$(pgrep -f "claude.*--add-dir" 2>/dev/null | head -1 || echo "")
echo "${CLAUDE_PID:-unknown}" > "${WORKER_STATE_DIR}/pid"
log "Claude PID: ${CLAUDE_PID:-unknown}"

# ── Done ──────────────────────────────────────────────────────────────────────
echo ""
echo "============================================="
echo "  ARIA WORKER SPAWNED"
echo "============================================="
echo "  Name:        $WORKER_NAME"
echo "  Session:     $WORKER_SESSION"
echo "  Role:        $ROLE"
echo "  Mode:        $ROLE_DEFAULT_MODE"
echo "  WorkDir:     $WORK_DIR"
[ "$IS_GIT" = "1" ] && echo "  Branch:      $WORKTREE_BRANCH"
echo "  State:       $WORKER_STATE_DIR"
echo "  Watchdog:    ${MAX_RUNTIME}s (PID $WATCHDOG_PID)"
echo "  ApprovePID:  $APPROVE_PID"
echo ""
echo "  Attach:      tmux attach -t $WORKER_SESSION"
echo "  Monitor:     tmux capture-pane -t $WORKER_SESSION -p"
echo "  Stop:        bash /root/aria/scripts/aria-worker-stop.sh $WORKER_NAME"
echo "  Panel:       http://127.0.0.1:9999/team"
echo "============================================="

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