#!/usr/bin/env bash
# athena-claude — take the coding task's Claude session over in this terminal.
#
# Installed by the environment bake at /opt/athena/bin/athena-claude. It asks the
# in-guest coding worker (athena-claude-coder, loopback port 46100) to hand its
# Claude session to you: the running turn is interrupted and drained, the parent
# agent sees a `blocked` result with blocker `human_active`, and Claude Code
# resumes the SAME session here — same CLAUDE_CONFIG_DIR, same checkout. When you
# leave Claude, however you leave it, the session is handed back with the id you
# ended on, and the next SDK run resumes exactly there.
#
# The lock the worker hands out is a lease (`owner_since`) bound to this
# launcher's PID: the return must present the lease, so a late or duplicate
# return never unlocks a newer holder, and the worker releases the lock by itself
# once this process is gone (closed terminal, suspended computer, killed shell).
# If another terminal already holds the session this launcher refuses to start a
# second writer; ATHENA_CLAUDE_TAKEOVER=1 re-mints the lease and takes it over.
#
# The guest token is read from the running server's pm2 environment and is never
# printed. Extra arguments are passed through to `claude`.
#
# Knobs (all optional): ATHENA_CODING_WORKER_PORT, ATHENA_CODING_WORKER_STATE_DIR,
# ATHENA_CODING_WORKER_PM2_HOME, ATHENA_CODING_WORKER_THREAD_ID (pick a thread;
# the default is the live one, else the most recent session), ATHENA_CLAUDE_BIN,
# ATHENA_CLAUDE_TAKEOVER=1.

set -euo pipefail

PORT="${ATHENA_CODING_WORKER_PORT:-46100}"
BASE_URL="http://127.0.0.1:${PORT}"
STATE_DIR="${ATHENA_CODING_WORKER_STATE_DIR:-/workspace/.claude-state}"
CODER_PM2_HOME="${ATHENA_CODING_WORKER_PM2_HOME:-/workspace/.pm2-athena-coder}"
CODER_PM2_NAME="${ATHENA_CODING_WORKER_PM2_NAME:-athena-claude-coder}"
CLAUDE_BIN="${ATHENA_CLAUDE_BIN:-/opt/athena/bin/claude}"
THREAD_ID="${ATHENA_CODING_WORKER_THREAD_ID:-}"
TAKEOVER="${ATHENA_CLAUDE_TAKEOVER:-0}"

fail() { printf 'athena-claude: %s\n' "$1" >&2; exit "${2:-1}"; }
json_field() { python3 -c 'import json, sys
value = json.load(sys.stdin).get(sys.argv[1])
print("" if value is None else value)' "$1"; }

for tool in curl python3; do
  command -v "$tool" >/dev/null 2>&1 || fail "$tool is required"
done
command -v pm2 >/dev/null 2>&1 || fail "pm2 is not installed; the coding worker is not running here" 2

# 1. The shared secret exists only in the running server's environment.
TOKEN="$(PM2_HOME="$CODER_PM2_HOME" pm2 jlist 2>/dev/null | python3 -c 'import json, sys
name = sys.argv[1]
try:
    procs = json.load(sys.stdin)
except Exception:
    procs = []
for proc in procs:
    env = proc.get("pm2_env") or {}
    if proc.get("name") == name and env.get("status") == "online":
        print(env.get("ATHENA_CODING_WORKER_TOKEN", ""))
        break' "$CODER_PM2_NAME" || true)"
[ -n "$TOKEN" ] || fail "the coding worker is not running (no online '$CODER_PM2_NAME' under $CODER_PM2_HOME); start a coding task from Athena first" 2

# api METHOD PATH JSON OUT_FILE -> writes the response body to OUT_FILE and prints
# the HTTP status code (000 when the server could not be reached).
api() {
  local method="$1" path="$2" body="$3" out="$4" code
  code="$(curl -s -o "$out" -w '%{http_code}' -X "$method" "$BASE_URL$path" \
    -H 'Content-Type: application/json' \
    -H "X-Athena-Coding-Worker-Token: $TOKEN" \
    --data "$body" 2>/dev/null)" || true
  printf '%s' "${code:-000}"
}

# 2. Request the handoff; the server picks the live thread when none is named.
#    Our PID rides along so the worker can tell a dead launcher from a live one.
REQUEST_BODY="$(python3 -c 'import json, sys
thread, takeover, pid = sys.argv[1], sys.argv[2] == "1", int(sys.argv[3])
payload = {"launcher_pid": pid}
if thread:
    payload["thread_id"] = thread
if takeover:
    payload["takeover"] = True
print(json.dumps(payload))' "$THREAD_ID" "$TAKEOVER" "$$")"
RESPONSE_FILE="$(mktemp)"
API_STATUS="$(api POST /internal/handoff/request "$REQUEST_BODY" "$RESPONSE_FILE")"
RESPONSE="$(cat "$RESPONSE_FILE")"
rm -f "$RESPONSE_FILE"
if [ "$API_STATUS" != "200" ]; then
  DETAIL="$(printf '%s' "$RESPONSE" | json_field detail 2>/dev/null || true)"
  fail "handoff refused (HTTP $API_STATUS): ${DETAIL:-$RESPONSE}" 2
fi

# 3. From here on the lock may be ours: arm the give-back BEFORE parsing anything
#    that could fail. It only fires once we know we acquired the lease; a lock we
#    cannot account for is released by the worker itself when this PID is gone.
ACQUIRED="False"
LEASE=""
SESSION_ID=""
give_back() {
  [ "$ACQUIRED" = "True" ] || return 0
  local body attempt code
  body="$(python3 -c 'import json, sys
thread, session, lease = sys.argv[1], sys.argv[2], sys.argv[3]
payload = {"thread_id": thread}
if session:
    payload["claude_session_id"] = session
if lease:
    payload["owner_since"] = lease
print(json.dumps(payload))' "$THREAD_ID" "$SESSION_ID" "$LEASE")"
  for attempt in 1 2 3 4 5; do
    code="$(api POST /internal/handoff/return "$body" /dev/null)"
    case "$code" in
      2*) return 0 ;;
      4*) printf 'athena-claude: handoff return refused (HTTP %s)\n' "$code" >&2; return 0 ;;
    esac
    [ "$attempt" -lt 5 ] && sleep 1
  done
  printf 'athena-claude: could not hand the session back (last HTTP %s); the worker releases the lock once this launcher (pid %s) is gone\n' "$code" "$$" >&2
}
trap give_back EXIT
trap 'exit 130' INT
trap 'exit 143' TERM HUP

PARSED_THREAD="$(printf '%s' "$RESPONSE" | json_field thread_id 2>/dev/null || true)"
ACQUIRED="$(printf '%s' "$RESPONSE" | json_field acquired 2>/dev/null || printf 'False')"
LEASE="$(printf '%s' "$RESPONSE" | json_field owner_since 2>/dev/null || true)"
SESSION_ID="$(printf '%s' "$RESPONSE" | json_field claude_session_id 2>/dev/null || true)"
TASK_CWD="$(printf '%s' "$RESPONSE" | json_field cwd 2>/dev/null || true)"
if [ -z "$PARSED_THREAD" ] || [ -z "$TASK_CWD" ] || [ -z "$LEASE" ]; then
  ACQUIRED="False"
  fail "malformed handoff response ($RESPONSE); the worker releases any lock held by this pid ($$) once it exits" 2
fi
THREAD_ID="$PARSED_THREAD"
if [ "$ACQUIRED" != "True" ]; then
  fail "another terminal already holds this session (since $LEASE); finish there, or run ATHENA_CLAUDE_TAKEOVER=1 athena-claude to take it over" 2
fi

# 4. Exactly the environment the SDK runs Claude with, and no ambient credential:
#    the key comes from the apiKeyHelper named in settings.json.
export CLAUDE_CONFIG_DIR="$STATE_DIR"
CLAUDE_CODE_PROJECT_DIR_NAME="$(basename "$TASK_CWD")"
export CLAUDE_CODE_PROJECT_DIR_NAME
export ENABLE_TOOL_SEARCH=true
if [ -r "$STATE_DIR/credential.json" ]; then
  GATEWAY_URL="$(json_field anthropic_base_url < "$STATE_DIR/credential.json" 2>/dev/null || true)"
  if [ -n "$GATEWAY_URL" ]; then
    export ANTHROPIC_BASE_URL="$GATEWAY_URL"
  fi
fi
unset ANTHROPIC_API_KEY ANTHROPIC_AUTH_TOKEN CLAUDE_CODE_OAUTH_TOKEN CLAUDE_CODE_API_KEY_FILE_DESCRIPTOR
cd "$TASK_CWD"

# 5. Resume the task's session — or start one whose id we choose, so it can be
#    reported back and resumed by the SDK later. Not `exec`: the EXIT trap above
#    must still run after Claude exits, whatever the exit code.
if [ -n "$SESSION_ID" ]; then
  printf 'athena-claude: resuming Claude session %s on thread %s (cwd %s)\n' "$SESSION_ID" "$THREAD_ID" "$TASK_CWD" >&2
  "$CLAUDE_BIN" --settings "$STATE_DIR/settings.json" --resume "$SESSION_ID" "$@"
else
  SESSION_ID="$(python3 -c 'import uuid; print(uuid.uuid4())')"
  printf 'athena-claude: starting Claude session %s on thread %s (cwd %s)\n' "$SESSION_ID" "$THREAD_ID" "$TASK_CWD" >&2
  "$CLAUDE_BIN" --settings "$STATE_DIR/settings.json" --session-id "$SESSION_ID" "$@"
fi
