#!/usr/bin/env bash
# 作业12 · 步骤2：故意抬高并发，观察延迟与错误
# 用法：
#   CONCURRENCY=8 N=24 bash hw12_02_并发压测.sh
#   CONCURRENCY=20 N=40 bash hw12_02_并发压测.sh --via-higress
# 说明：默认 max_tokens=32，避免把 GPU 压太久；学习用。
set -euo pipefail

VIA_HIGRESS=0
for a in "$@"; do
  [ "$a" = "--via-higress" ] && VIA_HIGRESS=1
done

CONCURRENCY="${CONCURRENCY:-8}"
N="${N:-24}"
MODEL="${MODEL:-qwen3.6-27b}"
KEY="${LITELLM_MASTER_KEY:-sk-litellm-bocheng-207}"
OUT_DIR="${HW12_OUT:-/opt/workbuddy/data/homework12}"
mkdir -p "$OUT_DIR"

if [ "$VIA_HIGRESS" = "1" ]; then
  URL="${HIGRESS_LLM_BASE:-http://192.168.110.207:18080/llm/v1}/chat/completions"
else
  URL="${LITELLM_BASE:-http://192.168.110.207:4000}/v1/chat/completions"
fi

STAMP="$(date +%Y%m%d_%H%M%S)"
CSV="$OUT_DIR/latency_conc_${STAMP}.csv"
LOG="$OUT_DIR/latency_conc_${STAMP}.log"
BODY="$(mktemp)"
WORK="$(mktemp -d)"
trap 'rm -rf "$BODY" "$WORK"' EXIT

cat >"$BODY" <<EOF
{"model":"$MODEL","messages":[{"role":"user","content":"用一句话介绍什么是网关"}],"max_tokens":32,"stream":false,"temperature":0}
EOF

echo "idx,http_code,time_total_s,error" >"$CSV"
echo "[INFO] URL=$URL CONCURRENCY=$CONCURRENCY N=$N" | tee "$LOG"
echo "[INFO] CSV=$CSV" | tee -a "$LOG"

run_one() {
  local i="$1"
  local out="$WORK/r_$i.json"
  local meta="$WORK/m_$i.txt"
  local code time err
  set +e
  code_time="$(
    curl -sS -o "$out" \
      -w "%{http_code} %{time_total}" \
      -H "Authorization: Bearer $KEY" \
      -H "Content-Type: application/json" \
      --data-binary @"$BODY" \
      --max-time 300 \
      "$URL" 2>"$WORK/e_$i.txt"
  )"
  rc=$?
  set -e
  if [ $rc -ne 0 ]; then
    code="000"
    time="0"
    err="$(tr '\n' ' ' <"$WORK/e_$i.txt" | cut -c1-120)"
  else
    code="$(echo "$code_time" | awk '{print $1}')"
    time="$(echo "$code_time" | awk '{print $2}')"
    err=""
    if [ "$code" != "200" ]; then
      err="$(tr '\n' ' ' <"$out" | cut -c1-120)"
    fi
  fi
  # 写一行（并发安全用临时文件再合并）
  printf '%s,%s,%s,%s\n' "$i" "$code" "$time" "\"${err//\"/\'}\"" >"$WORK/line_$i.csv"
  echo "[#$i] http=$code time=${time}s ${err:+err=$err}" >>"$LOG"
}

export -f run_one
export KEY URL BODY WORK LOG MODEL

# GNU parallel 未必有；用 xargs -P
seq 1 "$N" | xargs -P "$CONCURRENCY" -I{} bash -c 'run_one "$@"' _ {}

for i in $(seq 1 "$N"); do
  cat "$WORK/line_$i.csv" >>"$CSV"
done

echo | tee -a "$LOG"
python3 - <<PY | tee -a "$LOG"
import csv
from collections import Counter
from pathlib import Path
p=Path("$CSV")
rows=list(csv.DictReader(p.open(encoding="utf-8")))
codes=Counter(r["http_code"] for r in rows)
ts=[float(r["time_total_s"]) for r in rows if r["http_code"]=="200"]
print("=== 汇总 ===")
print("HTTP 码分布:", dict(codes))
print(f"成功 {len(ts)}/{len(rows)}")
if ts:
  ts_sorted=sorted(ts)
  def pct(q):
    i=min(len(ts_sorted)-1, max(0, int(round(q*(len(ts_sorted)-1)))))
    return ts_sorted[i]
  print(f"延迟 s: min={min(ts):.3f} p50={pct(0.5):.3f} p95={pct(0.95):.3f} max={max(ts):.3f}")
print()
print("请打开作业文档「现象→原因→处理」五联表，对照本汇总填写。")
PY

echo "[OK] 完成。CSV=$CSV LOG=$LOG"