#!/usr/bin/env bash
# SGLang vs vLLM controlled comparison on H200.
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_BASE="${SCRIPT_DIR}/results"

log_dir_global="${RESULT_BASE}/${RUN_ID}/logs"
mkdir -p "$log_dir_global"
log_init "${log_dir_global}/orchestrator.log"

log "experiment=${EXPERIMENT_NAME}"
log "run_id=${RUN_ID}"
log "platform=${PLATFORM}"
log "hardware=${HARDWARE}"
log "model=${MODEL_PATH}"

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

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
}

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

start_server() {
  local backend="$1"
  local phase="$2"
  local start_script
  if [[ "$backend" == "sglang" ]]; then
    start_script="${SGLANG_START_SCRIPT}"
    local port="$SGLANG_PORT"
  else
    start_script="${VLLM_START_SCRIPT}"
    local port="$VLLM_PORT"
  fi

  log "starting ${backend} server (phase=${phase}) with ${start_script}"
  bash "${start_script}" "$phase" >> "${log_dir_global}/${backend}_${phase}.server.outer.log" 2>&1

  if ! is_server_healthy "$port"; then
    log "error: ${backend} server (phase=${phase}) failed to become healthy"
    return 1
  fi
  log "${backend} server is healthy on port ${port}"
}

run_warmup() {
  local backend="$1"
  local phase="$2"
  local port input_len output_len num
  if [[ "$backend" == "sglang" ]]; then
    port="$SGLANG_PORT"
  else
    port="$VLLM_PORT"
  fi

  if [[ "$phase" == "phase1" ]]; then
    input_len=4000
    output_len=512
    num=2
  else
    input_len=200000
    output_len=1024
    num=1
  fi

  log "warming up ${backend} (phase=${phase}, input=${input_len}, output=${output_len}, num=${num})"
  "${VENV_CLIENT}/bin/python" "${SCRIPT_DIR}/../../scripts/common/warmup.py" \
    --backend "$backend" \
    --port "$port" \
    --input-len "$input_len" \
    --output-len "$output_len" \
    --num "$num" \
    --env-python "${VENV_CLIENT}/bin/python" \
    >> "${log_dir_global}/${backend}_${phase}.warmup.log" 2>&1
  log "warmup for ${backend} ${phase} completed"
}

scenario_already_completed() {
  local output_file="$1"
  local expected="$2"
  [[ -s "$output_file" ]] || return 1
  local completed
  completed="$(${VENV_CLIENT}/bin/python -c "
import json, sys
path = sys.argv[1]
try:
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:
                data = json.loads(line)
                print(data.get('completed', 0))
                break
except Exception:
    print(0)
" "$output_file")"
  [[ "${completed:-0}" -ge "$expected" ]]
}

run_phase() {
  local backend="$1"
  local phase="$2"
  local result_root="${RESULT_BASE}/${RUN_ID}/${backend}"
  local raw_dir="${result_root}/raw_outputs"
  local phase_log_dir="${result_root}/logs"

  mkdir -p "$raw_dir" "$phase_log_dir"

  local -n scenarios
  local max_model_len
  if [[ "$phase" == "phase1" ]]; then
    scenarios=PHASE1_SCENARIOS
    max_model_len="$PHASE1_MAX_MODEL_LEN"
  else
    scenarios=PHASE2_SCENARIOS
    max_model_len="$PHASE2_MAX_MODEL_LEN"
  fi

  local port
  if [[ "$backend" == "sglang" ]]; then
    port="$SGLANG_PORT"
  else
    port="$VLLM_PORT"
  fi

  log "===== ${backend} ${phase} START (max_model_len=${max_model_len}) ====="

  stop_server "$backend"
  start_server "$backend" "$phase"
  run_warmup "$backend" "$phase"

  for scenario in "${scenarios[@]}"; do
    read -r concurrency input_len output_len num_prompts <<< "$scenario"
    output_file="${raw_dir}/${backend}_${phase}_$(date '+%m%d')_${concurrency}_${input_len}_${output_len}.jsonl"
    detail_log="${phase_log_dir}/${backend}_${phase}_c${concurrency}_i${input_len}_o${output_len}.log"

    if scenario_already_completed "$output_file" "$num_prompts"; then
      log "skipping already-completed ${backend} ${phase} scenario: c=${concurrency} i=${input_len} o=${output_len}"
      continue
    fi

    log "running ${backend} ${phase} scenario: c=${concurrency} i=${input_len} o=${output_len} n=${num_prompts}"

    "${VENV_CLIENT}/bin/python" -m sglang.bench_serving \
      --backend "$backend" \
      --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: ${backend} ${phase} scenario c=${concurrency} i=${input_len} o=${output_len} failed; see ${detail_log}"
        continue
      }

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

  stop_server "$backend"
  log "===== ${backend} ${phase} DONE ====="
}

run_sharegpt() {
  local backend="$1"
  local phase="sharegpt"
  local result_root="${RESULT_BASE}/${RUN_ID}/${backend}"
  local raw_dir="${result_root}/raw_outputs"
  local phase_log_dir="${result_root}/logs"

  mkdir -p "$raw_dir" "$phase_log_dir"

  local port
  if [[ "$backend" == "sglang" ]]; then
    port="$SGLANG_PORT"
  else
    port="$VLLM_PORT"
  fi

  log "===== ${backend} ${phase} START (context_len=${SHAREGPT_CONTEXT_LEN}) ====="

  stop_server "$backend"
  start_server "$backend" phase1
  run_warmup "$backend" phase1

  for scenario in "${SHAREGPT_SCENARIOS[@]}"; do
    read -r concurrency context_len output_len num_prompts <<< "$scenario"
    output_file="${raw_dir}/${backend}_${phase}_$(date '+%m%d')_${concurrency}_${context_len}_${output_len}.jsonl"
    detail_log="${phase_log_dir}/${backend}_${phase}_c${concurrency}_ctx${context_len}_o${output_len}.log"

    if scenario_already_completed "$output_file" "$num_prompts"; then
      log "skipping already-completed ${backend} ${phase} scenario: c=${concurrency} ctx=${context_len} o=${output_len}"
      continue
    fi

    log "running ${backend} ${phase} scenario: c=${concurrency} ctx=${context_len} o=${output_len} n=${num_prompts}"

    "${VENV_CLIENT}/bin/python" -m sglang.bench_serving \
      --backend "$backend" \
      --host 127.0.0.1 \
      --port "$port" \
      --dataset-name sharegpt \
      --dataset-path "$SHAREGPT_DATASET" \
      --sharegpt-output-len "$output_len" \
      --sharegpt-context-len "$context_len" \
      --num-prompts "$num_prompts" \
      --max-concurrency "$concurrency" \
      --request-rate 10000 \
      --output-file "$output_file" \
      --output-details \
      > "$detail_log" 2>&1 || {
        log "ERROR: ${backend} ${phase} scenario c=${concurrency} ctx=${context_len} o=${output_len} failed; see ${detail_log}"
        continue
      }

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

  stop_server "$backend"
  log "===== ${backend} ${phase} DONE ====="
}

parse_backend() {
  local backend="$1"
  local result_root="${RESULT_BASE}/${RUN_ID}/${backend}"
  log "parsing ${backend} results in ${result_root}"
  "${VENV_CLIENT}/bin/python" "${SCRIPT_DIR}/../../scripts/common/parse_backend.py" "$result_root" --backend "$backend" \
    >> "${result_root}/logs/parse.log" 2>&1 || {
      log "WARNING: parser failed for ${backend}; see ${result_root}/logs/parse.log"
    }
}

write_backend_metadata() {
  local backend="$1"
  local result_root="${RESULT_BASE}/${RUN_ID}/${backend}"
  ensure_result_root "$result_root"
  local meta_json="${result_root}/results.json"

  local env_path
  if [[ "$backend" == "sglang" ]]; then
    env_path="$VENV_SGLANG"
  else
    env_path="$VENV_VLLM"
  fi

  # Capture the exact server launch args for reproducibility.
  local phase1_args phase2_args
  if [[ "$backend" == "sglang" ]]; then
    phase1_args="sglang serve --trust-remote-code --model-path $MODEL_PATH --tp 8 --moe-runner-backend marlin --context-length $PHASE1_MAX_MODEL_LEN --max-running-requests $PHASE1_MAX_RUNNING --mem-fraction-static 0.88 --host 0.0.0.0 --port $SGLANG_PORT"
    phase2_args="sglang serve --trust-remote-code --model-path $MODEL_PATH --tp 8 --moe-runner-backend marlin --context-length $PHASE2_MAX_MODEL_LEN --max-running-requests $PHASE2_MAX_RUNNING --mem-fraction-static 0.88 --host 0.0.0.0 --port $SGLANG_PORT"
  else
    phase1_args="vllm serve $MODEL_PATH --trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8 --max-model-len $PHASE1_MAX_MODEL_LEN --max-num-seqs $PHASE1_MAX_NUM_SEQS --block-size 256 --gpu-memory-utilization 0.90 --tokenizer-mode deepseek_v4 --reasoning-parser deepseek_v4 --no-disable-hybrid-kv-cache-manager --disable-uvicorn-access-log --port $VLLM_PORT"
    phase2_args="vllm serve $MODEL_PATH --trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8 --max-model-len $PHASE2_MAX_MODEL_LEN --max-num-seqs $PHASE2_MAX_NUM_SEQS --block-size 256 --gpu-memory-utilization 0.90 --tokenizer-mode deepseek_v4 --reasoning-parser deepseek_v4 --no-disable-hybrid-kv-cache-manager --disable-uvicorn-access-log --port $VLLM_PORT"
  fi

  write_metadata_json \
    "$meta_json" \
    "${EXPERIMENT_NAME}_${backend}" \
    "$RUN_ID" \
    "$MODEL_PATH" \
    "$backend" \
    "$backend" \
    "$HARDWARE" \
    "$ACCELERATOR" \
    "$CHIP" \
    "experiments/${EXPERIMENT_NAME}/run_bench.sh" \
    "$env_path" \
    "H200 ${backend} TP=8 comparison benchmark for DeepSeek-V4-Flash"

  # Embed config and server args.
  jq --arg backend "$backend" \
     --arg phase1_args "$phase1_args" \
     --arg phase2_args "$phase2_args" \
     '.config = {
        "tp": 8,
        "cuda_visible_devices": "0,1,2,3,4,5,6,7",
        "phase1_max_model_len": 32768,
        "phase2_max_model_len": 210000,
        "backend": $backend,
        "server_start_script": "experiments/dsv4_h200_sglang_vs_vllm/start_\($backend).sh",
        "phase1_server_args": $phase1_args,
        "phase2_server_args": $phase2_args
      }' "$meta_json" > "${meta_json}.tmp" && mv "${meta_json}.tmp" "$meta_json"
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

# Cleanup any leftovers.
stop_server sglang
stop_server vllm

write_backend_metadata sglang
write_backend_metadata vllm

# Run SGLang.
run_phase sglang phase1
run_sharegpt sglang
run_phase sglang phase2
parse_backend sglang

# Run vLLM.
run_phase vllm phase1
run_sharegpt vllm
run_phase vllm phase2
parse_backend vllm

# Generate comparison.
log "generating comparison report"
"${VENV_CLIENT}/bin/python" "${SCRIPT_DIR}/../../scripts/common/compare.py" \
  --sglang "${RESULT_BASE}/${RUN_ID}/sglang" \
  --vllm "${RESULT_BASE}/${RUN_ID}/vllm" \
  --output "${RESULT_BASE}/${RUN_ID}/comparison.md" \
  >> "${log_dir_global}/compare.log" 2>&1 || {
    log "WARNING: comparison script failed; see ${log_dir_global}/compare.log"
  }

log "all results saved to ${RESULT_BASE}/${RUN_ID}"