#!/usr/bin/env bash
# Custom benchmark orchestrator for single-service vLLM TP=8 on 8x H200.
#
# This runs 1 vLLM service (TP=8) on all 8 GPUs (0,1,2,3,4,5,6,7).
# bench_client.py sends requests to the single service.
#
# Usage:
#   bash run_bench.sh
#   SKIP_MANAGE_SERVER=1 bash run_bench.sh   # skip server start/stop

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 "services=${PORTS}"
log "lb_strategy=${LB_STRATEGY}"
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 1x vLLM TP=8 single-service benchmark using bench_client.py"

# Update metadata with config.
"${VENV_CLIENT}/bin/python" - "$METADATA_JSON" <<'PY'
import json, sys
path = sys.argv[1]
with open(path, "r", encoding="utf-8") as f:
    data = json.load(f)

data["config"] = {
    "tp": 8,
    "num_services": 1,
    "ports": "30005",
    "gpu_groups": "0,1,2,3,4,5,6,7",
    "kv_cache_dtype": "fp8",
    "block_size": 256,
    "max_num_seqs": 256,
    "client": "bench_client.py",
    "lb_strategy": "round_robin",
    "scenarios": [
        "1 512 256 20", "2 512 256 40", "4 512 256 80",
        "8 512 256 160", "16 512 256 320", "32 512 256 640",
        "64 512 256 1280", "128 512 256 2560",
        "1 2048 512 20", "8 2048 512 160", "32 2048 512 640", "128 2048 512 2560",
        "1 4096 1024 20", "8 4096 1024 160", "32 4096 1024 640", "128 4096 1024 2560"
    ]
}
with open(path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
PY

# ---------------------------------------------------------------------------
# Server helpers (single-service)
# ---------------------------------------------------------------------------

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

are_all_services_healthy() {
  local ports="$1"
  IFS=',' read -ra PORT_LIST <<< "$ports"
  for p in "${PORT_LIST[@]}"; do
    if ! is_server_healthy "$p"; then
      return 1
    fi
  done
  return 0
}

stop_all_servers() {
  log "stopping all vllm services"
  # Kill by PID files first.
  for pid_file in "${SCRIPT_DIR}"/server_port*.pid; do
    [[ -f "$pid_file" ]] || continue
    local pid
    pid="$(cat "$pid_file")"
    if kill -0 "$pid" 2>/dev/null; then
      log "killing pid=${pid}"
      kill "$pid" 2>/dev/null || true
      sleep 2
      kill -9 "$pid" 2>/dev/null || true
    fi
    rm -f "$pid_file"
  done
  # Fallback: kill any remaining vllm processes for this model.
  pkill -9 -f "vllm serve.*${MODEL_NAME}" 2>/dev/null || true
  sleep 2
}

start_all_servers() {
  log "starting all vllm services 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 ! are_all_services_healthy "$PORTS"; then
    log "error: not all vllm services became healthy"
    exit 1
  fi
  log "all vllm services are healthy"
}

# ---------------------------------------------------------------------------
# Pre-benchmark health check with detailed reporting
# ---------------------------------------------------------------------------

check_services_detailed() {
  local ports="$1"
  IFS=',' read -ra PORT_LIST <<< "$ports"
  local healthy_count=0
  local unhealthy_ports=""
  for p in "${PORT_LIST[@]}"; do
    if is_server_healthy "$p"; then
      ((healthy_count++))
    else
      unhealthy_ports="${unhealthy_ports}${p} "
    fi
  done
  echo "$healthy_count"
  if [[ -n "$unhealthy_ports" ]]; then
    echo "unhealthy: ${unhealthy_ports}" >&2
  fi
}

wait_for_all_services() {
  local ports="$1"
  local max_wait="${2:-240}"
  local interval="${3:-5}"
  local elapsed=0
  while (( elapsed < max_wait )); do
    local healthy_count
    healthy_count="$(check_services_detailed "$ports" | head -1)"
    IFS=',' read -ra PORT_LIST <<< "$ports"
    if (( healthy_count == ${#PORT_LIST[@]} )); then
      return 0
    fi
    log "waiting for services... ${healthy_count}/${#PORT_LIST[@]} healthy after ${elapsed}s"
    sleep "$interval"
    ((elapsed += interval))
  done
  return 1
}

# ---------------------------------------------------------------------------
# Mid-benchmark health check (called between scenarios)
# ---------------------------------------------------------------------------

check_and_report_services() {
  local ports="$1"
  IFS=',' read -ra PORT_LIST <<< "$ports"
  local healthy=()
  local unhealthy=()
  for p in "${PORT_LIST[@]}"; do
    if is_server_healthy "$p"; then
      healthy+=("$p")
    else
      unhealthy+=("$p")
    fi
  done
  log "service health check: healthy=[${healthy[*]}] unhealthy=[${unhealthy[*]}]"
  if [[ ${#unhealthy[@]} -gt 0 ]]; then
    log "WARNING: ${#unhealthy[@]} service(s) unhealthy: ${unhealthy[*]}"
  fi
  return ${#unhealthy[@]}
}

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

if [[ -n "${SKIP_MANAGE_SERVER:-}" ]]; then
  log "SKIP_MANAGE_SERVER is set, assuming services are already running"
  if ! wait_for_all_services "$PORTS"; then
    log "error: not all healthy services found on ports ${PORTS}"
    exit 1
  fi
else
  stop_all_servers
  start_all_servers
fi

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

for scenario in "${SCENARIOS[@]}"; do
  read -r concurrency input_len output_len num_prompts <<< "$scenario"
  # Fallback for legacy 3-field scenarios.
  if [[ -z "${num_prompts:-}" ]]; then
    num_prompts="$NUM_PROMPTS"
  fi

  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} num_prompts=${num_prompts}"

  "${VENV_CLIENT}/bin/python" "${SCRIPT_DIR}/bench_client.py" \
    --host 127.0.0.1 \
    --port "$PORT" \
    --model "$SERVED_MODEL_NAME" \
    --concurrency "$concurrency" \
    --input-len "$input_len" \
    --output-len "$output_len" \
    --num-prompts "$num_prompts" \
    --output-file "$output_file" \
    > "$detail_log" 2>&1 || {
      log "ERROR: scenario c=${concurrency} i=${input_len} o=${output_len} failed; see ${detail_log}"
      # After failure, check service health before continuing.
      check_and_report_services "$PORTS"
      continue
    }

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

  # Between scenarios, report service health so we can spot degradation early.
  check_and_report_services "$PORTS" || true
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}"