#!/usr/bin/env bash
#
# 压测的四步：干跑 → 正常负载 → 故障注入 → 记分板。
#
# 这个脚本只做「在当前 shell 里按顺序跑」，自己不建 tmux 会话、也不 attach。**长跑要在 tmux
# 里跑**（CLAUDE.md §4），所以正确的用法是人先开会话再调它：
#
#     tmux new -s polyloop-soak
#     SOAK_BUDGET_CALLS=4000 SOAK_FAULT_BUDGET_CALLS=600 SOAK_CONCURRENCY=4 \
#     SOAK_LIMIT=100 SOAK_APPWORLD_DATA_ROOT=/path/to/appworld \
#         bash tools/soak/soak.sh
#
# 会话名约定是 polyloop-soak。跑完不要急着 kill，留着给人复查。
#
# **末尾一律不接管道。** `pytest ... | tail` 的退出码来自管道最后一节，于是一次失败的跑会
# 报成 exit 0。要留日志就设 SOAK_LOG_DIR，那一路用重定向，不用 tee。
#
# 全部配置都是环境变量，没有位置参数：
#
#   必填
#     SOAK_BUDGET_CALLS         正常负载的模型调用次数上限
#     SOAK_FAULT_BUDGET_CALLS   故障注入那一步的模型调用次数上限
#     SOAK_CONCURRENCY          同时在跑的任务数上限
#     SOAK_LIMIT                每个场景最多跑几个任务（GovDoc 一个任务是三次运行）
#     SOAK_APPWORLD_DATA_ROOT   AppWorld 数据根目录（SOAK_SCENARIO 含 appworld 时）
#
#   可选
#     SOAK_OUT           产物根目录，默认 soak-out/<时间戳>
#     SOAK_SCENARIO      appworld / govdoc / both，默认 both
#     SOAK_CONTAINERS    AppWorld 容器数，默认 4
#     SOAK_SPLITS        AppWorld 划分，空格分隔，默认 "train dev"
#     SOAK_GOVDOC_DB     GovDoc 的审核点 sqlite，不给就用场景里的默认路径
#     SOAK_GOVDOC_CORPUS GovDoc 的语料目录，同上
#     SOAK_LOG_DIR       每一步的输出重定向到这里；不给就直接打屏（tmux 里能实时看）
#     SOAK_SKIP_FAULTS   非空则跳过故障注入那一步
#     SOAK_CONDA_ENV     conda 环境名，默认 PolyLoop

set -euo pipefail

STEP="启动"
trap 'echo "!!! 第「${STEP}」步失败（退出码 $?）" >&2' ERR

CONDA_ENV="${SOAK_CONDA_ENV:-PolyLoop}"
SCENARIO="${SOAK_SCENARIO:-both}"
CONTAINERS="${SOAK_CONTAINERS:-4}"
SPLITS="${SOAK_SPLITS:-train dev}"
LOG_DIR="${SOAK_LOG_DIR:-}"
OUT="${SOAK_OUT:-soak-out/$(date +%Y%m%d-%H%M%S)}"

BUDGET_CALLS="${SOAK_BUDGET_CALLS:?必须给 SOAK_BUDGET_CALLS：正常负载的模型调用次数上限}"
FAULT_BUDGET_CALLS="${SOAK_FAULT_BUDGET_CALLS:?必须给 SOAK_FAULT_BUDGET_CALLS：故障注入那一步的上限}"
CONCURRENCY="${SOAK_CONCURRENCY:?必须给 SOAK_CONCURRENCY：同时在跑的任务数上限}"
LIMIT="${SOAK_LIMIT:?必须给 SOAK_LIMIT：每个场景最多跑几个任务}"

# conda 和 Python 各缓冲一层，两层都得拆——只加其中一个，长跑命令仍然全程无输出。
RUN=(env PYTHONUNBUFFERED=1 conda run --live-stream -n "$CONDA_ENV" python)

RUNS_DIR="$OUT/runs"
FAULT_RUNS_DIR="$OUT/fault-runs"

# 场景相关的参数拼成数组。**用数组不用字符串**：路径里有空格时字符串会在展开时被切开，
# 而表现是「找不到这个目录」，看起来像数据没准备好。
SCENARIO_ARGS=(--scenario "$SCENARIO" --limit "$LIMIT" --containers "$CONTAINERS")
if [[ "$SCENARIO" == "appworld" || "$SCENARIO" == "both" ]]; then
    APPWORLD_DATA_ROOT="${SOAK_APPWORLD_DATA_ROOT:?SOAK_SCENARIO 含 appworld 时必须给 SOAK_APPWORLD_DATA_ROOT}"
    SCENARIO_ARGS+=(--appworld-data-root "$APPWORLD_DATA_ROOT")
    for split in $SPLITS; do
        SCENARIO_ARGS+=(--split "$split")
    done
fi
# 故障注入那一步的两个 GovDoc 路径没有默认值，而它们的权威在场景模块里。**问 Python 要，
# 不在这里写第二份**：两处各写一份路径，迟早有一处被改、另一处没改，而表现是「数据不在」。
GOVDOC_DB="${SOAK_GOVDOC_DB:-$("${RUN[@]}" -c 'from tools.soak.scenarios.govdoc import DEFAULT_DATA_ROOT as R; print(R / "app.sqlite")')}"
GOVDOC_CORPUS="${SOAK_GOVDOC_CORPUS:-$("${RUN[@]}" -c 'from tools.soak.scenarios.govdoc import DEFAULT_DATA_ROOT as R; print(R / "storage" / "prepared")')}"
if [[ "$SCENARIO" == "govdoc" || "$SCENARIO" == "both" ]]; then
    SCENARIO_ARGS+=(--govdoc-db "$GOVDOC_DB" --govdoc-corpus "$GOVDOC_CORPUS")
fi

# 故障注入的参数是另一套：它自己的 `--split` 只收一个值，AppWorld 数据根目录那一项叫
# `--data-root`。这里按 `python -m tools.soak.faults --help` 的形状拼。
FAULT_ARGS=(--govdoc-db "$GOVDOC_DB" --govdoc-corpus "$GOVDOC_CORPUS")
if [[ "$SCENARIO" == "appworld" || "$SCENARIO" == "both" ]]; then
    FAULT_ARGS+=(--data-root "$APPWORLD_DATA_ROOT" --split "${SPLITS%% *}")
fi

mkdir -p "$OUT"

step() {
    local name="$1"
    shift
    STEP="$name"
    echo ""
    echo "=== [$name] $* ==="
    if [[ -n "$LOG_DIR" ]]; then
        mkdir -p "$LOG_DIR"
        "$@" >"$LOG_DIR/$name.log" 2>&1
    else
        "$@"
    fi
}

step 01-干跑 "${RUN[@]}" -m tools.soak.run_soak \
    "${SCENARIO_ARGS[@]}" \
    --budget-calls "$BUDGET_CALLS" \
    --concurrency "$CONCURRENCY" \
    --runs-dir "$RUNS_DIR" \
    --report "$OUT/dry-run.md" \
    --dry-run

step 02-正常负载 "${RUN[@]}" -m tools.soak.run_soak \
    "${SCENARIO_ARGS[@]}" \
    --budget-calls "$BUDGET_CALLS" \
    --concurrency "$CONCURRENCY" \
    --runs-dir "$RUNS_DIR" \
    --report "$OUT/normal-load.md"

# 正常负载不该有缺文件，所以这一轮不开 --allow-undetermined：判不了就是要人去看一眼。
step 03-记分板-正常负载 "${RUN[@]}" -m tools.soak.scoreboard \
    --runs-dir "$RUNS_DIR" \
    --report "$OUT/scoreboard-normal.md" \
    --completing-tool submit_finding

if [[ -n "${SOAK_SKIP_FAULTS:-}" ]]; then
    echo ""
    echo "=== [04-故障注入] 按 SOAK_SKIP_FAULTS 跳过 ==="
else
    step 04-故障注入 "${RUN[@]}" -m tools.soak.faults \
        "${FAULT_ARGS[@]}" \
        --runs-dir "$FAULT_RUNS_DIR" \
        --budget-calls "$FAULT_BUDGET_CALLS"

    # 崩溃注入那一类天然会缺文件，这一轮才该开 --allow-undetermined。
    step 05-记分板-故障注入 "${RUN[@]}" -m tools.soak.scoreboard \
        --runs-dir "$FAULT_RUNS_DIR" \
        --report "$OUT/scoreboard-faults.md" \
        --completing-tool submit_finding \
        --allow-undetermined
fi

STEP="收尾"
echo ""
echo "=== 全部步骤通过 ==="
echo "产物：$OUT"
ls -1 "$OUT"