#!/usr/bin/env bash
set -euo pipefail

OUT_DIR="${1:-/root/work/logs/moe_tiny_matrix_$(date +%Y%m%d_%H%M%S)}"
BENCH_SCRIPT="${BENCH_SCRIPT:-/root/work/decode_microbench.py}"
MODEL_PATH="${MODEL_PATH:-/root/public-storage/models/Qwen/Qwen3.6-35B-A3B}"
BASE_URL="${BASE_URL:-http://127.0.0.1:1111}"

mkdir -p "$OUT_DIR"
echo "[matrix] out_dir=$OUT_DIR"
echo "[matrix] bench_script=$BENCH_SCRIPT"
echo "[matrix] model_path=$MODEL_PATH"

stop_server() {
  pkill -f "vllm.entrypoints.openai.api_server" 2>/dev/null || true
  pkill -f "VllmWorkerProcess" 2>/dev/null || true
  sleep 8
}

wait_health() {
  local deadline=$((SECONDS + 1200))
  while (( SECONDS < deadline )); do
    if python3 - <<'PY'
import urllib.request
try:
    urllib.request.urlopen("http://127.0.0.1:1111/health", timeout=5)
except Exception:
    raise SystemExit(1)
PY
    then
      echo "[matrix] health ok"
      return 0
    fi
    sleep 10
    echo "[matrix] waiting health: ${SECONDS}s"
  done
  echo "[matrix] ERROR: service did not become healthy" >&2
  return 1
}

start_server() {
  local impl="$1"
  local max_t="$2"
  local label="$3"
  local server_log="$OUT_DIR/${label}_server.log"
  echo "[matrix] starting $label"
  (
    cd /root
    ENGINEX_MOE_TINY_IMPL="$impl" \
    ENGINEX_MOE_TINY_MAX="$max_t" \
    VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 \
    PYTHONPATH="/usr/local/corex/lib/python3/dist-packages:/usr/local/corex/lib64/python3/dist-packages:${PYTHONPATH:-}" \
    LD_LIBRARY_PATH="/usr/local/openmpi/lib:/usr/local/corex-3.2.3/lib64:/usr/local/corex/lib64:/usr/local/corex/lib:/usr/local/iluvatar/lib64:${LD_LIBRARY_PATH:-}" \
    nohup python3 -m vllm.entrypoints.openai.api_server \
      --model "$MODEL_PATH" \
      --host 0.0.0.0 \
      --port 1111 \
      --served-model-name llm \
      --max-model-len 100000 \
      --enforce-eager \
      --trust-remote-code \
      -tp 4 \
      --gpu-memory-utilization 0.90 \
      --max-num-seqs 2 \
      --disable-log-requests \
      --disable-frontend-multiprocessing \
      --max-num-batched-tokens 8192 \
      --enable-chunked-prefill \
      --enable-prefix-caching \
      --enable-auto-tool-choice \
      --tool-call-parser qwen3_coder \
      --reasoning-parser qwen3 \
      > "$server_log" 2>&1 &
    echo $! > "$OUT_DIR/${label}_server.pid"
  )
  wait_health
}

run_case() {
  local impl="$1"
  local max_t="$2"
  local label="${impl}_max${max_t}"
  local result="$OUT_DIR/${label}_short_c2_t128_r4.json"
  local monitor="$OUT_DIR/${label}_ixsmi.json"
  local warmup="$OUT_DIR/${label}_warmup.json"

  stop_server
  start_server "$impl" "$max_t" "$label"

  echo "[matrix] warmup $label"
  python3 "$BENCH_SCRIPT" \
    --url "$BASE_URL" \
    --model llm \
    --label "${label}_warmup" \
    --prompt-mode short \
    --concurrency 1 \
    --requests 1 \
    --max-tokens 32 \
    --timeout 900 \
    --out "$warmup"

  echo "[matrix] benchmark $label"
  python3 "$BENCH_SCRIPT" \
    --url "$BASE_URL" \
    --model llm \
    --label "${label}_short_c2_t128_r4" \
    --prompt-mode short \
    --concurrency 2 \
    --requests 4 \
    --max-tokens 128 \
    --timeout 900 \
    --monitor-out "$monitor" \
    --monitor-interval 1.0 \
    --out "$result"

  stop_server
}

run_case tokenwise 4
run_case tokenwise 8
run_case tokenwise 16
run_case bmm 8
run_case bmm 16

python3 - "$OUT_DIR" <<'PY'
import json
import sys
from pathlib import Path

out_dir = Path(sys.argv[1])
rows = []
for path in sorted(out_dir.glob("*_short_c2_t128_r4.json")):
    data = json.loads(path.read_text(encoding="utf-8"))
    rows.append({
        "label": data.get("label"),
        "success_rate": data.get("success_rate"),
        "ttft_p90_sec": data.get("ttft_p90_sec"),
        "output_tps_p10_per_request": data.get("output_tps_p10_per_request"),
        "output_tps_p50_per_request": data.get("output_tps_p50_per_request"),
        "aggregate_output_tps": data.get("aggregate_output_tps"),
        "avg_gpu_util_pct": (data.get("monitor") or {}).get("avg_gpu_util_pct"),
        "max_gpu_util_pct": (data.get("monitor") or {}).get("max_gpu_util_pct"),
    })

summary = {
    "out_dir": str(out_dir),
    "rows": rows,
}
(out_dir / "summary.json").write_text(
    json.dumps(summary, ensure_ascii=False, indent=2),
    encoding="utf-8",
)

lines = [
    "# MoE tiny-batch matrix",
    "",
    "| case | success | TTFT P90 | Output TPS P10 | Output TPS P50 | Aggregate Output TPS | Avg GPU util | Max GPU util |",
    "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |",
]
for r in rows:
    lines.append(
        f"| {r['label']} | {r['success_rate']:.2%} | "
        f"{r['ttft_p90_sec']:.3f} | {r['output_tps_p10_per_request']:.3f} | "
        f"{r['output_tps_p50_per_request']:.3f} | {r['aggregate_output_tps']:.3f} | "
        f"{r['avg_gpu_util_pct']:.1f} | {r['max_gpu_util_pct']:.1f} |"
    )
(out_dir / "summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8")
print(json.dumps(summary, ensure_ascii=False, indent=2))
PY

echo "[matrix] done"
echo "[matrix] summary=$OUT_DIR/summary.md"