#!/bin/bash
# ============================================================================
# agent-pane-track.sh — Agent Team tmux Pane 追踪 Hook
#
# 用于 PreToolUse(Agent) 和 PostToolUse(Agent) 两个事件
# - PreToolUse: 快照当前所有 tmux pane（仅首次），记录 PM pane
# - PostToolUse: 对比快照，将新增 pane 记录到 .claude/team-panes.json
#
# 仅当 tool_input.team_name 非空时（Agent Team 场景）才执行
# ============================================================================

set -euo pipefail

# 项目根目录
PROJECT_ROOT="${CLAUDE_PROJECT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}"
STATE_FILE="${PROJECT_ROOT}/.claude/team-panes.json"
LOG_FILE="${PROJECT_ROOT}/.claude/team-panes.log"

# 日志函数
log() {
    local msg="[$(date '+%Y-%m-%d %H:%M:%S')] agent-pane-track: $*"
    echo "${msg}" >&2
    echo "${msg}" >> "${LOG_FILE}" 2>/dev/null || true
}

# 从 stdin 读取 JSON（3 秒超时防止阻塞）
INPUT_JSON=$(timeout 3 cat 2>/dev/null) || exit 0

if [ -z "${INPUT_JSON}" ]; then
    exit 0
fi

# 提取关键字段
HOOK_EVENT=$(echo "${INPUT_JSON}" | jq -r '.hook_event_name // empty' 2>/dev/null)
SESSION_ID=$(echo "${INPUT_JSON}" | jq -r '.session_id // empty' 2>/dev/null)
TEAM_NAME=$(echo "${INPUT_JSON}" | jq -r '.tool_input.team_name // empty' 2>/dev/null)

# 非 team agent → 跳过
if [ -z "${TEAM_NAME}" ]; then
    exit 0
fi

log "event=${HOOK_EVENT} team=${TEAM_NAME} session=${SESSION_ID}"

# 快照文件路径（使用 team_name 确保同一 team 共享一个快照）
SNAPSHOT_FILE="/tmp/claude-pane-snapshot-${SESSION_ID}-${TEAM_NAME}"

# 临时目录确保存在
mkdir -p "${PROJECT_ROOT}/.claude"

# ── PreToolUse: 快照当前 pane 列表 ──────────────────────────
if [ "${HOOK_EVENT}" == "PreToolUse" ]; then
    # 记录 PM pane（当前 pane）
    PM_PANE_ID=""
    if [ -n "${TMUX_PANE:-}" ]; then
        PM_PANE_ID="${TMUX_PANE}"
    else
        PM_PANE_ID=$(tmux display-message -p '#{pane_id}' 2>/dev/null || echo "")
    fi

    # 仅在快照不存在时创建（防止并发 Agent 调用覆盖快照）
    if [ ! -f "${SNAPSHOT_FILE}" ]; then
        tmux list-panes -a -F '#{pane_id}' 2>/dev/null | sort > "${SNAPSHOT_FILE}" || true
        log "快照已创建: $(wc -l < "${SNAPSHOT_FILE}") panes, PM=${PM_PANE_ID}"
    else
        log "快照已存在，跳过覆盖"
    fi

    # 初始化状态文件（如果不存在）
    if [ ! -f "${STATE_FILE}" ]; then
        TIMESTAMP=$(date -u '+%Y-%m-%dT%H:%M:%S')
        jq -n \
            --arg team "${TEAM_NAME}" \
            --arg pm "${PM_PANE_ID}" \
            --arg ts "${TIMESTAMP}" \
            '{
                team_name: $team,
                pm_pane_id: $pm,
                agent_pane_ids: [],
                created_at: $ts
            }' > "${STATE_FILE}"
        log "状态文件已初始化: team=${TEAM_NAME}, pm=${PM_PANE_ID}"
    fi

    exit 0
fi

# ── PostToolUse: 对比快照，记录新增 pane ────────────────────
if [ "${HOOK_EVENT}" == "PostToolUse" ]; then
    # 读取快照
    if [ ! -f "${SNAPSHOT_FILE}" ]; then
        log "快照文件不存在，跳过 diff"
        exit 0
    fi

    # 获取当前 pane 列表
    CURRENT_PANES=$(tmux list-panes -a -F '#{pane_id}' 2>/dev/null | sort || true)
    if [ -z "${CURRENT_PANES}" ]; then
        exit 0
    fi

    # Diff: 找出新增 pane
    NEW_PANES=$(comm -13 "${SNAPSHOT_FILE}" <(echo "${CURRENT_PANES}") 2>/dev/null || true)

    if [ -z "${NEW_PANES}" ]; then
        log "无新增 pane"
        # 注意: 不删除快照，因为后续 Agent 调用可能还需要它
        exit 0
    fi

    log "发现新增 pane: $(echo "${NEW_PANES}" | tr '\n' ' ')"

    # 读取 PM pane ID
    PM_PANE_ID=$(jq -r '.pm_pane_id // empty' "${STATE_FILE}" 2>/dev/null || echo "")

    # ── 将 PM pane 的模型 pane option 复制到新增 pane ──
    # 确保子 agent pane 继承正确的模型配置（同 session 不同 pane 隔离）
    if [ -n "${PM_PANE_ID}" ]; then
        opt_val="$(tmux show-option -p -t "${PM_PANE_ID}" @anthropic_model 2>/dev/null | sed 's/^[^ ]* //' || true)"
        if [ -n "${opt_val}" ]; then
            while IFS= read -r pane_id; do
                [ -z "${pane_id}" ] && continue
                tmux set-option -p -t "${pane_id}" @anthropic_model "${opt_val}" 2>/dev/null || true
            done <<< "${NEW_PANES}"
            log "已复制 PM pane 模型配置 (@anthropic_model=${opt_val}) 到新增 pane"
        fi
    fi

    # 读取已有 agent pane IDs
    EXISTING_IDS=$(jq -r '.agent_pane_ids[]?' "${STATE_FILE}" 2>/dev/null || true)

    # 构建更新后的 agent_pane_ids（去重，排除 PM pane）
    UPDATED_IDS="${EXISTING_IDS}"
    while IFS= read -r pane_id; do
        [ -z "${pane_id}" ] && continue
        # 排除 PM pane
        [ "${pane_id}" == "${PM_PANE_ID}" ] && continue
        # 去重检查
        if ! echo "${UPDATED_IDS}" | grep -qx "${pane_id}"; then
            UPDATED_IDS="${UPDATED_IDS}"$'\n'"${pane_id}"
        fi
    done <<< "${NEW_PANES}"

    # 转换为 JSON 数组并更新状态文件
    JSON_ARRAY=$(echo "${UPDATED_IDS}" | grep -v '^$' | sort -u | jq -R . | jq -s .)
    jq --argjson ids "${JSON_ARRAY}" '.agent_pane_ids = $ids' "${STATE_FILE}" > "${STATE_FILE}.tmp" \
        && mv "${STATE_FILE}.tmp" "${STATE_FILE}"

    log "状态文件已更新: agent_panes=$(echo "${UPDATED_IDS}" | grep -v '^$' | sort -u | tr '\n' ' ')"

    # 注意: 不删除快照文件，让所有 Agent 调用共享同一初始快照
    # 快照将在 TeamDelete 清理时统一删除

    exit 0
fi

exit 0