#!/usr/bin/env bash
# H200 native vLLM baseline benchmark for DeepSeek-V4-Flash.
set -Eeuo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")"

# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../scripts/common/lib.sh"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../scripts/common/platform.sh"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/config.env"

RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}"
RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}"
RAW_DIR="${RESULT_ROOT}/raw_outputs"
LOG_DIR="${RESULT_ROOT}/logs"

ensure_result_root "$RESULT_ROOT"
log_init "${LOG_DIR}/orchestrator.log"

log "experiment=${EXPERIMENT_NAME}"
log "run_id=${RUN_ID}"
log "result_root=${RESULT_ROOT}"
log "platform=${PLATFORM}"
log "chip=${CHIP}"
log "accelerator=${ACCELERATOR}"
log "engine=${ENGINE}"
log "hardware=${HARDWARE}"
log "model=${MODEL_PATH}"
log "server_start_script=${SERVER_START_SCRIPT}"

# Write initial metadata for this run.
METADATA_JSON="${RESULT_ROOT}/results.json"
write_metadata_json \
  "$METADATA_JSON" \
  "$EXPERIMENT_NAME" \
  "$RUN_ID" \
  "$MODEL_PATH" \
  "$BACKEND" \
  "$ENGINE" \
  "$HARDWARE" \
  "$ACCELERATOR" \
  "$CHIP" \
  "experiments/${EXPERIMENT_NAME}/run_bench.sh" \
  "$VENV_SERVER" \
  "H200 native vLLM baseline benchmark for DeepSeek-V4-Flash"

# Update metadata with config.
"${VENV_CLIENT}/bin/python" - "$METADATA_JSON" <<'PY'
import json
import sys

path = sys.argv[1]
with open(path, "r", encoding="utf-8") as f:
    data = json.load(f)

data["config"] = {
    "tp": 4,
    "kv_cache_dtype": "fp8",
    "block_size": 256,
    "max_num_seqs": 256,
    "port": 30005,
    "num_prompts": 128,
    "scenarios": ["32 512 256", "128 512 256", "32 4000 512"]
}
with open(path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
PY

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

stop_server() {
  local pid_file="/data/user1/yy/dsv4_h200_vllm.pid"
  if [[ -f "$pid_file" ]]; then
    local pid
    pid="$(cat "$pid_file")"
    if kill -0 "$pid" 2>/dev/null; then
      log "stopping existing vllm server pid=${pid}"
      kill "$pid" 2>/dev/null || true
      sleep 5
      kill -9 "$pid" 2>/dev/null || true
    fi
    rm -f "$pid_file"
  fi
  pkill -9 -f "vllm serve.*DeepSeek-V4-Flash" 2>/dev/null || true
  sleep 2
}

start_server() {
  log "starting vllm server with ${SERVER_START_SCRIPT}"
  if [[ ! -x "${SERVER_START_SCRIPT}" ]]; then
    log "error: start script not found or not executable: ${SERVER_START_SCRIPT}"
    exit 1
  fi
  bash "${SERVER_START_SCRIPT}" >> "${LOG_DIR}/server.outer.log" 2>&1
  if ! is_server_healthy; then
    log "error: vllm server failed to become healthy"
    exit 1
  fi
  log "vllm server is healthy"
}

on_exit() {
  local code=$?
  log "orchestrator exiting with code=${code}"
  if [[ -z "${SKIP_MANAGE_SERVER:-}" ]]; then
    stop_server
  fi
  exit "$code"
}
trap on_exit EXIT

if [[ -n "${SKIP_MANAGE_SERVER:-}" ]]; then
  log "SKIP_MANAGE_SERVER is set, assuming server is already running on port ${PORT}"
  if ! is_server_healthy; then
    log "error: no healthy server found at port ${PORT}"
    exit 1
  fi
else
  stop_server
  start_server
fi

log "===== BENCHMARK START ====="

for scenario in "${SCENARIOS[@]}"; do
  read -r concurrency input_len output_len <<< "$scenario"
  output_file="${RAW_DIR}/vllm_$(date '+%m%d')_${concurrency}_${input_len}_${output_len}.jsonl"
  detail_log="${LOG_DIR}/vllm_c${concurrency}_i${input_len}_o${output_len}.log"

  log "running scenario: concurrency=${concurrency} input=${input_len} output=${output_len}"

  "${VENV_CLIENT}/bin/python" -m sglang.bench_serving \
    --backend vllm \
    --host 127.0.0.1 \
    --port "$PORT" \
    --dataset-name random \
    --random-input-len "$input_len" \
    --random-output-len "$output_len" \
    --num-prompts "$NUM_PROMPTS" \
    --max-concurrency "$concurrency" \
    --request-rate 10000 \
    --output-file "$output_file" \
    --output-details \
    > "$detail_log" 2>&1 || {
      log "ERROR: scenario c=${concurrency} i=${input_len} o=${output_len} failed; see ${detail_log}"
      continue
    }

  log "finished scenario: output=${output_file}"
done

log "===== BENCHMARK DONE ====="

log "parsing results"
"${VENV_CLIENT}/bin/python" "${SCRIPT_DIR}/parse_results.py" "$RESULT_ROOT" >> "${LOG_DIR}/parse.log" 2>&1 || {
  log "WARNING: parser failed; see ${LOG_DIR}/parse.log"
}

log "all results saved to ${RESULT_ROOT}"