#!/usr/bin/env bash
# Resume cases 5-8 against the four already-running TP2 base servers.
set -Eeuo pipefail

RESULT_ROOT=${RESULT_ROOT:?RESULT_ROOT is required}
REPO_ROOT=${REPO_ROOT:-/data/wxy/sskj-h3}
INPUT_ROOT=${INPUT_ROOT:-$REPO_ROOT/throughput/sglang-base/inputs/ref2va-feishu-20260831}
RECORDS=${RECORDS:-$INPUT_ROOT/records.json}
ASSETS_ROOT=${ASSETS_ROOT:-$INPUT_ROOT/assets}
MODEL=${MODEL:-/data/hf_models/MiniMax-H3}
PYTHON=${PYTHON:-/root/.miniconda3/envs/sglang/bin/python}
CLIENT_SCRIPT=${CLIENT_SCRIPT:-$REPO_ROOT/throughput/sglang-base/scripts/ref2va_feishu_bench.py}
BASE_PORT=${BASE_PORT:-34010}
PORT_STRIDE=${PORT_STRIDE:-10}
HOST=${HOST:-127.0.0.1}
SERVER_PIDS=${SERVER_PIDS:?SERVER_PIDS is required}

declare -a CLIENT_PIDS=()
read -r -a SERVERS <<< "$SERVER_PIDS"

log() { printf '[%s] %s\n' "$(date '+%F %T')" "$*"; }

stop_servers() {
  local pid alive deadline
  for pid in "${SERVERS[@]}"; do
    if [[ -r "/proc/$pid/cmdline" ]] && tr '\0' ' ' <"/proc/$pid/cmdline" | grep -Fq "$RESULT_ROOT"; then
      kill -INT "$pid" 2>/dev/null || true
    fi
  done
  deadline=$((SECONDS + 120))
  while ((SECONDS < deadline)); do
    alive=0
    for pid in "${SERVERS[@]}"; do kill -0 "$pid" 2>/dev/null && alive=1; done
    ((alive == 0)) && break
    sleep 2
  done
  for pid in "${SERVERS[@]}"; do
    if kill -0 "$pid" 2>/dev/null; then
      kill -TERM -- "-$pid" 2>/dev/null || kill -TERM "$pid" 2>/dev/null || true
      sleep 5
      kill -KILL -- "-$pid" 2>/dev/null || kill -KILL "$pid" 2>/dev/null || true
    fi
  done
}

cleanup() {
  local rc=$? pid
  trap - EXIT INT TERM
  for pid in "${CLIENT_PIDS[@]}"; do kill -TERM "$pid" 2>/dev/null || true; done
  stop_servers
  nvidia-smi --query-gpu=index,pstate,memory.used,utilization.gpu --format=csv,noheader \
    >"$RESULT_ROOT/gpu_after.csv" 2>&1 || true
  exit "$rc"
}
trap cleanup EXIT INT TERM

for replica in 0 1 2 3; do
  port=$((BASE_PORT + replica * PORT_STRIDE))
  curl -fsS --max-time 5 "http://${HOST}:${port}/health" >/dev/null \
    || { log "server not healthy on port $port"; exit 1; }
done

batch_dir="$RESULT_ROOT/batch_02"
mkdir -p "$batch_dir"
failed=0
for replica in 0 1 2 3; do
  port=$((BASE_PORT + replica * PORT_STRIDE))
  case_number=$((5 + replica))
  client_dir="$batch_dir/client_${replica}_case$(printf '%02d' "$case_number")_port${port}"
  mkdir -p "$client_dir"
  "$PYTHON" "$CLIENT_SCRIPT" run --host "$HOST" --port "$port" --replica-index "$replica" \
    --case-number "$case_number" --records "$RECORDS" --assets-root "$ASSETS_ROOT" \
    --model "$MODEL" --num-inference-steps 20 --warmup-steps 0 --short-edge 768 \
    --aspect-ratio 9:16 --duration-seconds 15 --flow-shift 12.0 --audio-flow-shift 3.0 \
    --output "$client_dir/result.json" >"$client_dir/client.log" 2>&1 &
  CLIENT_PIDS+=("$!")
  log "started case=$case_number replica=$replica port=$port pid=$!"
done
for pid in "${CLIENT_PIDS[@]}"; do wait "$pid" || failed=1; done
CLIENT_PIDS=()
((failed == 0)) || { log "batch 2 had failed requests"; exit 1; }
log "batch 2 complete"

"$PYTHON" "$CLIENT_SCRIPT" summarize --result-root "$RESULT_ROOT" \
  --output "$RESULT_ROOT/summary.json" >"$RESULT_ROOT/summary.log" 2>&1
log "summary complete"
stop_servers
nvidia-smi --query-gpu=index,pstate,memory.used,utilization.gpu --format=csv,noheader \
  >"$RESULT_ROOT/gpu_after.csv"
trap - EXIT INT TERM
log "resume complete: $RESULT_ROOT"