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

# Reproducible SGLang vs vLLM backend comparison for DeepSeek-V4-Flash.
# Produces JSONL raw outputs compatible with the legacy 2026-07-07 sweep.
#
# Usage:
#   bash scripts/benchmark_dsv4_backend_comparison.sh [vllm|sglang|all]
#
# Output:
#   bench_results/dsv4_backend_comparison_${RUN_ID}/raw_outputs/

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"

BACKEND_FILTER="${1:-all}"
RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}"
RESULT_ROOT="${RESULT_ROOT:-${ROOT_DIR}/bench_results/dsv4_backend_comparison_${RUN_ID}}"
RAW_DIR="${RESULT_ROOT}/raw_outputs"
LOG_DIR="${RESULT_ROOT}/logs"

mkdir -p "${RAW_DIR}" "${LOG_DIR}"

BENCH_PY="${ROOT_DIR}/envs/sglang/bin/python"
DATASET="${ROOT_DIR}/datasets/ShareGPT_filtered_chat.json"
MODEL="/data/models/DeepSeek-V4-Flash"
SERVED_MODEL_NAME="deepseek-v4-flash"
NUM_PROMPTS=512
WARMUP=100

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

run_case() {
  local backend="$1"
  local port="$2"
  local concurrency="$3"
  local input_len="$4"
  local output_len="$5"
  local output_file="${RAW_DIR}/${backend}_$(date '+%m%d')_${concurrency}_${input_len}_${output_len}.jsonl"
  local detail_log="${LOG_DIR}/${backend}_c${concurrency}_i${input_len}_o${output_len}.log"

  log "running ${backend}: concurrency=${concurrency} input=${input_len} output=${output_len} -> ${output_file}"

  "${BENCH_PY}" -m sglang.bench_serving \
    --backend "${backend}" \
    --host 127.0.0.1 \
    --port "${port}" \
    --model "${MODEL}" \
    --served-model-name "${SERVED_MODEL_NAME}" \
    --dataset-name random \
    --dataset-path "${DATASET}" \
    --random-input-len "${input_len}" \
    --random-output-len "${output_len}" \
    --num-prompts "${NUM_PROMPTS}" \
    --max-concurrency "${concurrency}" \
    --warmup-requests "${WARMUP}" \
    --output-file "${output_file}" \
    --output-details \
    > "${detail_log}" 2>&1

  log "finished ${backend} c=${concurrency} i=${input_len} o=${output_len}"
}

main() {
  log "run_id=${RUN_ID}"
  log "result_root=${RESULT_ROOT}"
  log "backend_filter=${BACKEND_FILTER}"

  # This script assumes a healthy server is already running on the target port.
  if [[ "${BACKEND_FILTER}" == "all" || "${BACKEND_FILTER}" == "sglang" ]]; then
    log "===== SGLang sweep (port 30000) ====="
    run_case sglang 30000 32 512 256
    run_case sglang 30000 128 512 256
    run_case sglang 30000 256 512 256
    run_case sglang 30000 512 512 256
    run_case sglang 30000 512 512 2000
    run_case sglang 30000 32 4000 512
    run_case sglang 30000 128 4000 512
    run_case sglang 30000 512 4000 512
    run_case sglang 30000 32 16000 512
    run_case sglang 30000 128 16000 512
    run_case sglang 30000 512 1000 256
    run_case sglang 30000 768 1000 256
    run_case sglang 30000 1024 1000 256
    run_case sglang 30000 512 1000 1000
  fi

  if [[ "${BACKEND_FILTER}" == "all" || "${BACKEND_FILTER}" == "vllm" ]]; then
    log "===== vLLM sweep (port 8000) ====="
    run_case vllm 8000 32 512 256
    run_case vllm 8000 128 512 256
    run_case vllm 8000 256 512 256
    run_case vllm 8000 128 512 2000
    run_case vllm 8000 256 512 2000
    run_case vllm 8000 512 512 2000
    run_case vllm 8000 32 1000 256
    run_case vllm 8000 128 1000 256
    run_case vllm 8000 256 1000 256
    run_case vllm 8000 512 1000 256
    run_case vllm 8000 768 1000 256
    run_case vllm 8000 1024 1000 256
    run_case vllm 8000 32 1000 1000
    run_case vllm 8000 128 1000 1000
    run_case vllm 8000 256 1000 1000
    run_case vllm 8000 512 1000 1000
    run_case vllm 8000 1024 1000 1000
    run_case vllm 8000 32 4000 512
    run_case vllm 8000 128 4000 512
    run_case vllm 8000 256 4000 512
    run_case vllm 8000 512 4000 512
    run_case vllm 8000 32 8000 1000
    run_case vllm 8000 128 8000 1000
    run_case vllm 8000 256 8000 1000
    run_case vllm 8000 512 8000 1000
    run_case vllm 8000 32 16000 512
    run_case vllm 8000 128 16000 512
    run_case vllm 8000 256 16000 512
    run_case vllm 8000 32 32000 512
    run_case vllm 8000 128 32000 512
  fi

  log "===== Done ====="
  log "all outputs saved to ${RAW_DIR}"
}

main "$@"