#!/usr/bin/env bash
# 协作存储 U7 压测入口（方案 §3 负载合同 / §5 U7；B5 seed / C2 双实例）。
#
# 起隔离栈（backend mock + 隔离 PG）→ 可选大数据量 seed → 预建引导账号
# → k6 阶梯 → SSE 可见延迟探针 → 数据不变量对账 → 结果落 results/<ts>/ → 回收。
# 绝不打生产服务：默认自建 PG 实例（svc.sh）；--dsn 复用已有隔离实例。
#
# 用法：
#   tests/load/collaboration/run_load.sh [--profile smoke|ladder|sustained60m|sustained2h]
#                                        [--users N] [--dsn <PG DSN>] [--port <app port>]
#                                        [--seed none|tiny|contract] [--instances 1|2]
#                                        [--port-b <second port>] [--keep]
#                                        [--skip-k6] [--skip-probe] [--help]
# 依赖：bash4、go（构建 backend/探针/seed）、psql（对账）、curl；
#        k6 可选（缺失时跳过 k6 段并给出安装指引，探针与对账照跑）。
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
HERE="$ROOT/tests/load/collaboration"
PROFILE=smoke USERS=30 APP_PORT=18980 APP_PORT_B=18982 PG_PORT=18981 KEEP=0 SKIP_K6=0 SKIP_PROBE=0
INSTANCES=1 SEED=none
DSN=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --profile) PROFILE="$2"; shift 2 ;;
    --users) USERS="$2"; shift 2 ;;
    --dsn) DSN="$2"; shift 2 ;;
    --port) APP_PORT="$2"; shift 2 ;;
    --port-b) APP_PORT_B="$2"; shift 2 ;;
    --instances) INSTANCES="$2"; shift 2 ;;
    --seed) SEED="$2"; shift 2 ;;
    --keep) KEEP=1; shift ;;
    --skip-k6) SKIP_K6=1; shift ;;
    --skip-probe) SKIP_PROBE=1; shift ;;
    -h|--help) sed -n '2,18p' "$0"; exit 0 ;;
    *) echo "unknown arg: $1" >&2; exit 2 ;;
  esac
done

case "$INSTANCES" in
  1|2) ;;
  *) echo "[load] --instances must be 1 or 2" >&2; exit 2 ;;
esac
case "$SEED" in
  none|tiny|contract) ;;
  *) echo "[load] --seed must be none|tiny|contract" >&2; exit 2 ;;
esac
# 避免 app B 与隔离 PG 默认端口碰撞（18981）。
if [[ "$APP_PORT_B" -eq "$PG_PORT" || "$APP_PORT_B" -eq "$APP_PORT" ]]; then
  APP_PORT_B=$((APP_PORT + 2))
  echo "[load] --port-b 与 app/PG 冲突，改用 $APP_PORT_B" >&2
fi

TS="$(date +%Y%m%d-%H%M%S)"
RESULTS="$HERE/results/$TS"
mkdir -p "$RESULTS"
OWNED_PG=0
APP_PID=""
APP_PID_B=""
ARTS_DIR=""
cleanup() {
  local ec=$?
  trap - EXIT # 防止 exit 重入本 trap
  [[ -n "${APP_PID:-}" ]] && kill "$APP_PID" 2>/dev/null || true
  [[ -n "${APP_PID_B:-}" ]] && kill "$APP_PID_B" 2>/dev/null || true
  wait "$APP_PID" 2>/dev/null || true
  wait "$APP_PID_B" 2>/dev/null || true
  if [[ $OWNED_PG -eq 1 && $KEEP -eq 0 ]]; then
    "$ROOT/scripts/svc.sh" u7load stop >/dev/null 2>&1 || true
    "$ROOT/scripts/svc.sh" rm u7load --purge >/dev/null 2>&1 || true
  fi
  if [[ $KEEP -eq 1 ]]; then
    if [[ "$INSTANCES" -eq 2 ]]; then
      echo "[load] --keep：栈保留（PG=u7load:$PG_PORT，appA=127.0.0.1:$APP_PORT，appB=127.0.0.1:$APP_PORT_B，结果 $RESULTS）"
    else
      echo "[load] --keep：栈保留（PG=u7load:$PG_PORT，app=127.0.0.1:$APP_PORT，结果 $RESULTS）"
    fi
  fi
  exit "$ec"
}
trap cleanup EXIT

# ---- 1) PG：自建隔离实例（或复用 --dsn） ----
if [[ -z "$DSN" ]]; then
  echo "[load] 创建隔离 PG 实例 u7load（端口 $PG_PORT）"
  "$ROOT/scripts/svc.sh" db create u7load --port "$PG_PORT" >/dev/null
  OWNED_PG=1
  "$ROOT/scripts/svc.sh" u7load start >/dev/null
  DSN="$("$ROOT/scripts/svc.sh" db dsn u7load --reveal)"
fi
echo "[load] PG DSN 已就绪（密码不回显）"

# ---- 2) 构建 backend、探针与 seed ----
echo "[load] 构建 backend / SSE 探针 / seed"
BIN="$HERE/.build/agent-system"
mkdir -p "$HERE/.build"
( cd "$ROOT/modules/backend" && CGO_ENABLED=0 go build -o "$BIN" ./cmd/agent-system )
( cd "$HERE/sse_probe" && go build -o "$HERE/.build/sse_probe" . )
( cd "$HERE/seed" && go build -o "$HERE/.build/seed" . )

# ---- 3) 压测库 + 引导账号（CLI 直连，终态口令避免首登改密分支） ----
ADMIN_DSN="$(python3 - "$DSN" <<'PY'
import sys
from urllib.parse import urlsplit, urlunsplit, parse_qsl, urlencode
u = urlsplit(sys.argv[1])
q = dict(parse_qsl(u.query))
q.setdefault("sslmode", "disable")
print(urlunsplit((u.scheme, u.netloc, "/" + (u.path.lstrip("/") or "postgres"), urlencode(q), "")))
PY
)"
psql "$ADMIN_DSN" -v ON_ERROR_STOP=1 -c "DROP DATABASE IF EXISTS u7load_app WITH (FORCE)" -c "CREATE DATABASE u7load_app" >/dev/null
APP_DSN="$(python3 - "$ADMIN_DSN" <<'PY'
import sys
from urllib.parse import urlsplit, urlunsplit
u = urlsplit(sys.argv[1])
print(urlunsplit((u.scheme, u.netloc, "/u7load_app", u.query, "")))
PY
)"
"$BIN" user create --dsn "$APP_DSN" --admin --owner load_admin --password load-admin-pass-Aa1 >/dev/null
"$BIN" user create --dsn "$APP_DSN" --owner probe_user --password probe-user-pass-Aa9 >/dev/null

# ---- 3b) B5 seed（可选）：先建 seed_u* 再 COPY 灌量 ----
if [[ "$SEED" != "none" ]]; then
  case "$SEED" in
    tiny) SEED_USERS=3 ;;
    contract) SEED_USERS=10 ;;
  esac
  echo "[load] seed profile=$SEED：创建 seed_u0..seed_u$((SEED_USERS-1))"
  for i in $(seq 0 $((SEED_USERS - 1))); do
    "$BIN" user create --dsn "$APP_DSN" --owner "seed_u$i" --password "seed-user-pass-${i}-Aa9" >/dev/null
  done
  echo "[load] 灌入数据集（manifest → $RESULTS/seed-manifest.json）"
  "$HERE/.build/seed" --dsn "$APP_DSN" --profile "$SEED" --manifest "$RESULTS/seed-manifest.json"
fi

# CLI 给定密码仍写 must_change=true（U4 设密门）；压测口令即终态，清标志以免探针/k6 403。
psql "$APP_DSN" -v ON_ERROR_STOP=1 -c "UPDATE users SET must_change_password=FALSE" >/dev/null

# ---- 4) 起被测实例（mock；执行原语关闭——压测不打执行面） ----
ARTS_DIR="$RESULTS/artifacts"
mkdir -p "$ARTS_DIR"
# 子进程独立 cwd，避免 ./data 落测试目录（multi-instance 同款纪律）。
APP_CWD_A="$RESULTS/cwd-a"
APP_CWD_B="$RESULTS/cwd-b"
mkdir -p "$APP_CWD_A" "$APP_CWD_B"

# 注意：不可用 APP_PID="$( … &; echo $! )"——bash 命令替换会等其子 shell
# 内后台任务结束，serve 会把编排永远卡住。$! 必须在主 shell 赋值。
echo "[load] 启动实例 A 127.0.0.1:$APP_PORT（mock / tier=team / auth=on / instance=u7load-a）"
APP_LOG="$RESULTS/app-a.log"
# exec 替换子 shell：$! 即为 serve PID，cleanup 才能真正回收（否则孤儿占端口）。
(
  cd "$APP_CWD_A"
  exec "$BIN" serve --dsn "$APP_DSN" --addr "127.0.0.1:$APP_PORT" --tier team --auth on --mock \
    --instance u7load-a --artifacts "$ARTS_DIR" --disable-exec-tools >"$APP_LOG" 2>&1
) &
APP_PID=$!

if [[ "$INSTANCES" -eq 2 ]]; then
  echo "[load] 启动实例 B 127.0.0.1:$APP_PORT_B（共享 PG/artifacts / instance=u7load-b）"
  APP_LOG_B="$RESULTS/app-b.log"
  (
    cd "$APP_CWD_B"
    exec "$BIN" serve --dsn "$APP_DSN" --addr "127.0.0.1:$APP_PORT_B" --tier team --auth on --mock \
      --instance u7load-b --artifacts "$ARTS_DIR" --disable-exec-tools >"$APP_LOG_B" 2>&1
  ) &
  APP_PID_B=$!
fi

# auth=on 时 /api/health 同样要求鉴权：200/401 都视为已就绪。
app_ready() {
  local port="$1"
  local code
  code=$(curl -s -o /dev/null -w "%{http_code}" "http://127.0.0.1:$port/api/health")
  [[ "$code" == 200 || "$code" == 401 ]]
}
for _ in $(seq 1 60); do app_ready "$APP_PORT" && break; sleep 0.5; done
app_ready "$APP_PORT" || { echo "[load] app A 未就绪，日志 $APP_LOG" >&2; exit 1; }
if [[ "$INSTANCES" -eq 2 ]]; then
  for _ in $(seq 1 60); do app_ready "$APP_PORT_B" && break; sleep 0.5; done
  app_ready "$APP_PORT_B" || { echo "[load] app B 未就绪，日志 $APP_LOG_B" >&2; exit 1; }
fi

BASE_A="http://127.0.0.1:$APP_PORT"
BASE_B="http://127.0.0.1:$APP_PORT_B"
BASE="$BASE_A"

# ---- 5) k6 阶梯 ----
if [[ $SKIP_K6 -eq 1 ]]; then
  echo "[load] --skip-k6：跳过 k6 段"
elif ! command -v k6 >/dev/null 2>&1; then
  cat >"$RESULTS/k6-skipped.txt" <<'EOF'
k6 未安装：本段跳过（探针与对账照跑）。
安装（本机 linuxbrew）：brew install k6
安装后重跑：tests/load/collaboration/run_load.sh --profile <同参数>
EOF
  echo "[load] k6 未安装，跳过（见 $RESULTS/k6-skipped.txt）"
else
  echo "[load] k6 profile=$PROFILE users=$USERS instances=$INSTANCES"
  K6_ENV=(BASE_URL="$BASE_A" ADMIN_OWNER=load_admin ADMIN_PASS=load-admin-pass-Aa1 USERS="$USERS" PROFILE="$PROFILE")
  if [[ "$INSTANCES" -eq 2 ]]; then
    K6_ENV+=(BASE_URLS="$BASE_A,$BASE_B")
  fi
  ( cd "$HERE" && env "${K6_ENV[@]}" \
      k6 run --summary-export "$RESULTS/k6-summary.json" k6/collab_api.js ) | tee "$RESULTS/k6-output.txt"
fi

# ---- 6) SSE 可见延迟探针 ----
if [[ $SKIP_PROBE -eq 0 ]]; then
  echo "[load] SSE 可见延迟探针 A（60s，10 会话）"
  "$HERE/.build/sse_probe" -base "$BASE_A" -dur 60s -convs 10 -out "$RESULTS/sse-probe-a.json"
  if [[ "$INSTANCES" -eq 2 ]]; then
    echo "[load] SSE 可见延迟探针 B（60s，10 会话）"
    "$HERE/.build/sse_probe" -base "$BASE_B" -dur 60s -convs 10 -out "$RESULTS/sse-probe-b.json"
  fi
fi

# ---- 7) 数据不变量对账（方案 §3：零违反） ----
echo "[load] 数据不变量对账"
psql "$APP_DSN" -v ON_ERROR_STOP=1 \
  -c "SELECT 'users' k, COUNT(*) v FROM users UNION ALL SELECT 'projects', COUNT(*) FROM projects UNION ALL SELECT 'conversations', COUNT(*) FROM conversations UNION ALL SELECT 'messages', COUNT(*) FROM messages UNION ALL SELECT 'events', COUNT(*) FROM events UNION ALL SELECT 'knowledge', COUNT(*) FROM knowledge UNION ALL SELECT 'runs', COUNT(*) FROM runs" \
  -c "SELECT 'dup_seq_per_conv' k, COUNT(*) v FROM (SELECT conversation_id, seq FROM events GROUP BY 1,2 HAVING COUNT(*)>1) d" \
  -c "SELECT 'nonterminal_orphan_runs' k, COUNT(*) v FROM runs WHERE status IN ('pending','running','waiting_hitl')" \
  | tee "$RESULTS/invariants.txt"

DUP=$(psql "$APP_DSN" -tAc "SELECT COUNT(*) FROM (SELECT conversation_id, seq FROM events GROUP BY 1,2 HAVING COUNT(*)>1) d")
if [[ "$DUP" != "0" ]]; then
  echo "[load] 对账失败：events 每会话 seq 出现重复（$DUP 组）" >&2
  exit 1
fi
echo "[load] 对账通过：无重复 seq；孤儿 run 报告见 invariants.txt（探针正常收口应为 0）"

echo "[load] 完成。结果目录：$RESULTS"