#!/usr/bin/env bash
# agent-claim.sh — 占用一个空闲 agent slot
# stdout：可 eval 的 export 行；stderr：人类可读消息
#
# Exit codes（用于 AI 主动决策，不要 hardcode 提示用户跑命令）：
#   0  成功（stdout 有 export 行）
#   1  参数错误
#   2  同分支已占用（同分支不能在两个 slot 同时检出）
#   3  池满（所有 slot 都在用）
#   4  池未初始化（agent-pool 目录不存在 / 没有 slot）
#   5  池被显式禁用（FFOA_AGENT_POOL_ENABLED=false）
#   99 内部错误
#
# AI 调用方约定：见 exit 4/5 时**主动**用 AskUserQuestion 询问是否初始化 / 走旧路径，
# 不要把 "请运行 xxx 命令" 抛给用户——执行者是 AI。
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=../lib/agent-pool-lib.sh
source "${SCRIPT_DIR}/../lib/agent-pool-lib.sh"

usage() {
  cat <<EOF >&2
用法:
  bash $0 <task-branch> [base-branch]
  eval "\$(bash $0 <task-branch> [base-branch])"

参数:
  task-branch   要在 slot 中检出的任务分支
  base-branch   分支不存在时基于哪条分支创建（默认 develop）

输出:
  stdout 是 shell export 行，可 eval：
    FFOA_AGENT_SLOT, FFOA_AGENT_DIR, FFOA_AGENT_LOCK,
    FFOA_AGENT_BRANCH 以及 slot 的端口号
EOF
}

if [[ $# -lt 1 ]]; then
  usage
  exit 1
fi

TASK_BRANCH="$1"
BASE_BRANCH="${2:-develop}"
AGENT_LABEL="${FFOA_AGENT_LABEL:-${USER:-unknown}}"
SESSION_ID="${FFOA_AGENT_SESSION_ID:-$$}"

# 1. 显式禁用检查
if ! ap_pool_enabled; then
  cat >&2 <<EOF
agent-claim: 池被显式禁用（FFOA_AGENT_POOL_ENABLED=false）
EXIT_REASON=disabled
EOF
  exit 5
fi

# 2. 池存在性检查
if ! ap_pool_exists; then
  pool_root="$(ap_pool_root)"
  cat >&2 <<EOF
agent-claim: 池未初始化
POOL_ROOT=${pool_root}
EXIT_REASON=not_initialized
EOF
  exit 4
fi

POOL_ROOT="$(ap_pool_root)"

# 3. 同分支占用检查
for slot in $(ap_list_slots); do
  state="$(ap_slot_state "${slot}")"
  [[ "${state}" == claimed ]] || continue
  existing_branch="$(ap_read_lock "${slot}" task_branch || true)"
  if [[ "${existing_branch}" == "${TASK_BRANCH}" ]]; then
    existing_dir="$(ap_read_lock "${slot}" worktree_path || true)"
    cat >&2 <<EOF
agent-claim: 分支 ${TASK_BRANCH} 已被 slot-${slot} 占用
EXISTING_SLOT=${slot}
EXISTING_DIR=${existing_dir}
EXIT_REASON=branch_conflict
EOF
    exit 2
  fi
done

# 4. 找空闲 slot 并 claim
claim_slot() {
  local slot="$1"
  local lockf flockf
  lockf="$(ap_lock_file "${slot}")"
  flockf="$(ap_flock_file "${slot}")"
  mkdir -p "$(dirname "${flockf}")"
  : > "${flockf}"

  exec 200>"${flockf}"
  if ! flock -n 200; then
    return 1
  fi

  if [[ -f "${lockf}" ]] && ! ap_lock_is_stale "${slot}"; then
    exec 200>&-
    return 1
  fi

  local slot_dir hostname_v
  slot_dir="$(ap_slot_dir "${slot}")"
  hostname_v="$(hostname)"

  if [[ ! -d "${slot_dir}" ]]; then
    echo "agent-claim: slot-${slot} 目录不存在（pool corrupted?）" >&2
    exec 200>&-
    return 99
  fi

  (
    cd "${slot_dir}"
    git fetch origin --quiet 2>/dev/null || true

    if git show-ref --verify --quiet "refs/heads/${TASK_BRANCH}"; then
      # 本地分支已存在：先 checkout，然后按本地 vs 远端关系决定是否 ff
      git checkout "${TASK_BRANCH}" >/dev/null 2>&1

      if git show-ref --verify --quiet "refs/remotes/origin/${TASK_BRANCH}"; then
        local local_sha remote_sha mb
        local_sha="$(git rev-parse HEAD 2>/dev/null)"
        remote_sha="$(git rev-parse "origin/${TASK_BRANCH}" 2>/dev/null)"
        if [[ -n "${local_sha}" && -n "${remote_sha}" && "${local_sha}" != "${remote_sha}" ]]; then
          mb="$(git merge-base HEAD "origin/${TASK_BRANCH}" 2>/dev/null || true)"
          if [[ "${mb}" == "${local_sha}" ]]; then
            # 本地是远端 ancestor → fast-forward 拉新
            local fwd
            fwd="$(git rev-list --count "${local_sha}..${remote_sha}" 2>/dev/null || echo 0)"
            git merge --ff-only "origin/${TASK_BRANCH}" >/dev/null 2>&1
            echo "agent-claim: 本地 ${TASK_BRANCH} fast-forward ${fwd} 个 commit 到远端" >&2
          elif [[ "${mb}" == "${remote_sha}" ]]; then
            # 本地领先远端：可能有未 push 的 work，不动
            local ahead
            ahead="$(git rev-list --count "${remote_sha}..${local_sha}" 2>/dev/null || echo 0)"
            echo "agent-claim: ⚠ 本地 ${TASK_BRANCH} 领先远端 ${ahead} 个 commit（未 push？），保留本地 HEAD" >&2
          else
            # diverged：本地远端各自有独立 commit
            local lead trail
            lead="$(git rev-list --count "${mb}..${local_sha}" 2>/dev/null || echo 0)"
            trail="$(git rev-list --count "${mb}..${remote_sha}" 2>/dev/null || echo 0)"
            echo "agent-claim: ⚠ 本地 ${TASK_BRANCH} 与远端分歧（本地 ${lead} / 远端 ${trail}），保留本地 HEAD" >&2
            echo "agent-claim:   手工处理：cd ${slot_dir} && git rebase origin/${TASK_BRANCH}（或 git merge）" >&2
          fi
        fi
      fi
    elif git show-ref --verify --quiet "refs/remotes/origin/${TASK_BRANCH}"; then
      git checkout -B "${TASK_BRANCH}" "origin/${TASK_BRANCH}" >/dev/null 2>&1
    else
      git checkout -B "${TASK_BRANCH}" "origin/${BASE_BRANCH}" >/dev/null 2>&1
    fi

    # base 提示：本任务分支 merge-base 之后 origin/base 又往前走了多少
    if git show-ref --verify --quiet "refs/remotes/origin/${BASE_BRANCH}"; then
      local base_mb base_ahead
      base_mb="$(git merge-base HEAD "origin/${BASE_BRANCH}" 2>/dev/null || true)"
      if [[ -n "${base_mb}" ]]; then
        base_ahead="$(git rev-list --count "${base_mb}..origin/${BASE_BRANCH}" 2>/dev/null || echo 0)"
        if [[ "${base_ahead}" -gt 0 ]]; then
          echo "agent-claim: ℹ origin/${BASE_BRANCH} 自分支创建以来领先 ${base_ahead} 个 commit" >&2
          echo "agent-claim:   合到 develop 前考虑 rebase（不自动做）：git rebase origin/${BASE_BRANCH}" >&2
        fi
      fi
    fi
  ) >&2

  # agent_ppid：调用方必须显式 export $FFOA_AGENT_PPID 才记录。
  # 不能用默认 $PPID——claim 的常见调用模式是 eval "$(bash agent-claim.sh ...)"，
  # $PPID 是 $() 子 shell，claim 返回那一刻就死，会让 heartbeat 自检和 sweep 立即误杀。
  # 字段空 → heartbeat 自检 / sweep agent-dead 路径全部跳过 → 行为退化到老逻辑（靠 hook）。
  # SessionEnd hook / wrapper script / 长 driver shell 才负责 export FFOA_AGENT_PPID=$$。
  local agent_ppid="${FFOA_AGENT_PPID:-}"

  ap_write_lock_atomic "${slot}" \
    "slot=${slot}" \
    "pid=$$" \
    "agent_ppid=${agent_ppid}" \
    "host=${hostname_v}" \
    "agent=${AGENT_LABEL}" \
    "session_id=${SESSION_ID}" \
    "task_branch=${TASK_BRANCH}" \
    "worktree_path=${slot_dir}" \
    "claimed_at=$(ap_now_iso)" \
    "heartbeat_at=$(ap_now_iso)" \
    "heartbeat_pid="

  # 释放 flock 后再 fork daemon——否则 daemon 继承 FD 200 → flock 不释放，
  # 下次 claim 同 slot 拒锁。详见 .learnings/ERRORS/ERR-20260509-002.md
  exec 200>&-

  nohup bash "${SCRIPT_DIR}/agent-heartbeat.sh" "${slot}" >/dev/null 2>&1 200>&- &
  local hb_pid=$!
  disown
  ap_update_lock_field "${slot}" heartbeat_pid "${hb_pid}"

  echo "${slot}"
  return 0
}

chosen=""
for slot in $(ap_list_slots); do
  state="$(ap_slot_state "${slot}")"
  if [[ "${state}" == free ]]; then
    if chosen="$(claim_slot "${slot}")"; then
      break
    fi
  fi
done

if [[ -z "${chosen}" ]]; then
  for slot in $(ap_list_slots); do
    state="$(ap_slot_state "${slot}")"
    if [[ "${state}" == stale ]]; then
      echo "agent-claim: 回收 stale slot-${slot}" >&2
      if chosen="$(claim_slot "${slot}")"; then
        break
      fi
    fi
  done
fi

# orphan 兜底：claimed 但 task_branch 远端已删 + 本地干净 + commit 已合并
# （agent 已退出但 heartbeat 守护还在跑的假占用场景）
# 注：is_orphan→release→claim_slot 之间存在 TOCTOU 窗口，但 claim_slot 内部用 flock 兜底，
# 并发争抢只有一个会拿到锁，其余 claim_slot 返回非零跳过下一个 slot。
if [[ -z "${chosen}" ]]; then
  for slot in $(ap_list_slots); do
    state="$(ap_slot_state "${slot}")"
    [[ "${state}" == claimed ]] || continue
    if ap_lock_is_orphan "${slot}"; then
      echo "agent-claim: 回收 orphan slot-${slot}（远端分支已删 + 本地干净 + commit 已合并）" >&2
      if ! bash "${SCRIPT_DIR}/agent-release.sh" "${slot}" >&2; then
        echo "agent-claim: ⚠ slot-${slot} release 失败，跳过下一个" >&2
        continue
      fi
      if chosen="$(claim_slot "${slot}")"; then
        break
      fi
    fi
  done
fi

if [[ -z "${chosen}" ]]; then
  cat >&2 <<EOF
agent-claim: 没有空闲 slot（池满）
EXIT_REASON=pool_full
EOF
  bash "${SCRIPT_DIR}/agent-status.sh" >&2
  exit 3
fi

slot_dir="$(ap_slot_dir "${chosen}")"
lock_file="$(ap_lock_file "${chosen}")"

# 输出可 eval 的 shell export
echo "export FFOA_AGENT_SLOT=${chosen}"
echo "export FFOA_AGENT_DIR='${slot_dir}'"
echo "export FFOA_AGENT_LOCK='${lock_file}'"
echo "export FFOA_AGENT_BRANCH='${TASK_BRANCH}'"

# slot 自带的端口
if [[ -f "${slot_dir}/.env" ]]; then
  for k in FRONTEND_PORT BACKEND_PORT POSTGRES_PORT REDIS_PORT; do
    v="$(awk -F= -v key="${k}" '$1==key {gsub(/[[:space:]]|#.*/, "", $2); print $2; exit}' "${slot_dir}/.env" 2>/dev/null || true)"
    [[ -n "${v}" ]] && echo "export ${k}=${v}"
  done
fi

{
  echo
  echo "agent-claim: 已占用 slot-${chosen}"
  echo "  目录:   ${slot_dir}"
  echo "  分支:   ${TASK_BRANCH}"
  echo "  释放:   bash ${SCRIPT_DIR}/agent-release.sh ${chosen}"
} >&2

# 刷新 .code-workspace 让 IDE 看到这个 slot 的新分支标签（输出走 stderr，不污染 eval 行）
ap_write_code_workspace || true