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

# Compare vllm-dspark with --spec-tokens 3 vs 5.
# Uses the parameterized start script and focused benchmark.

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
START_SCRIPT="${SCRIPT_DIR}/../start_dsv4_dspark_8card_param.sh"
BENCH_SCRIPT="${SCRIPT_DIR}/bench_dspark_focused.sh"
RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}"
RESULT_ROOT="${RESULT_ROOT:-/data/user1/yy/bench_results/dspark_st_comparison_${RUN_ID}}"
LOG_DIR="${RESULT_ROOT}/logs"
PORT="${PORT:-30004}"
PID_FILE="/data/user1/yy/dsv4_flash_dspark.pid"
WARMUP_REQUESTS="${WARMUP_REQUESTS:-100}"

mkdir -p "${LOG_DIR}"

log() {
  echo "[$(date --iso-8601=seconds)] $*" | tee -a "${LOG_DIR}/orchestrator.log"
}

is_server_healthy() {
  curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1
}

stop_server() {
  if [[ -f "${PID_FILE}" ]]; then
    local pid
    pid="$(cat "${PID_FILE}" 2>/dev/null || true)"
    if [[ -n "${pid}" ]] && kill -0 "${pid}" 2>/dev/null; then
      log "stopping existing dspark server pid=${pid}"
      kill "${pid}" 2>/dev/null || true
      sleep 5
      kill -9 "${pid}" 2>/dev/null || true
    fi
    rm -f "${PID_FILE}"
  fi
  local pids
  pids="$(pgrep -f "VLLM::|vllm serve.*DeepSeek-V4-Flash-DSpark" 2>/dev/null || true)"
  if [[ -n "${pids}" ]]; then
    for pid in ${pids}; do
      if [[ "$pid" != "$$" ]]; then kill -9 "$pid" 2>/dev/null || true; fi
    done
  fi
  sleep 2
}

start_server() {
  local spec_tokens="$1"
  log "starting dspark server with spec-tokens=${spec_tokens}"
  if [[ ! -x "${START_SCRIPT}" ]]; then
    log "error: start script not found or not executable: ${START_SCRIPT}"
    exit 1
  fi
  SPEC_TOKENS="${spec_tokens}" bash "${START_SCRIPT}" >> "${LOG_DIR}/server_st${spec_tokens}.outer.log" 2>&1
  if ! is_server_healthy; then
    log "error: dspark server (st=${spec_tokens}) failed to become healthy"
    exit 1
  fi
  log "dspark server (st=${spec_tokens}) is healthy"
}

run_benchmark() {
  local spec_tokens="$1"
  local log_file="${LOG_DIR}/focused_st${spec_tokens}.log"
  log "running focused benchmark with spec-tokens=${spec_tokens}, warmup=${WARMUP_REQUESTS}"
  RUN_ID="${RUN_ID}_st${spec_tokens}" RESULT_ROOT="${RESULT_ROOT}" WARMUP_REQUESTS="${WARMUP_REQUESTS}" bash "${BENCH_SCRIPT}" 2>&1 | tee "${log_file}"
  log "finished focused benchmark with spec-tokens=${spec_tokens}"
}

on_exit() {
  local code="$?"
  log "orchestrator exiting with code=${code}"
  stop_server
  exit "${code}"
}
trap on_exit EXIT

main() {
  log "run_id=${RUN_ID}"
  log "result_root=${RESULT_ROOT}"
  log "warmup_requests=${WARMUP_REQUESTS}"

  log "===== SPEC-TOKENS 3 START ====="
  stop_server
  start_server 3
  run_benchmark 3
  stop_server

  log "cooldown 30s"
  sleep 30

  log "===== SPEC-TOKENS 5 START ====="
  start_server 5
  run_benchmark 5
  stop_server

  log "===== COMPARISON DONE ====="
  log "all results saved to ${RESULT_ROOT}"
}

main "$@"