#!/usr/bin/env bash
# TP×DP matrix benchmark for DeepSeek-V4-Flash on vLLM.
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"
MATRIX_FILE="${MATRIX_FILE:-${SCRIPT_DIR}/matrix.json}"
MATRIX_MODE="${MATRIX_MODE:-Y}"
SCENARIO_TIMEOUT_S="${SCENARIO_TIMEOUT_S:-1800}"
GPU_MEM_SAMPLE_INTERVAL_S="${GPU_MEM_SAMPLE_INTERVAL_S:-1}"

PYTHON="${VENV_CLIENT}/bin/python"

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

log "experiment=${EXPERIMENT_NAME} run_id=${RUN_ID} platform=${PLATFORM} hardware=${HARDWARE}"
log "matrix_mode=${MATRIX_MODE} matrix_file=${MATRIX_FILE}"

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

health_port_for() {
  local dp="$1"
  if [[ "$dp" -gt 1 ]]; then
    echo "$VLLM_DP_SUPERVISOR_PORT"
  else
    echo "$VLLM_PORT"
  fi
}

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 tp="$1"
  local dp="$2"
  local pid_file="/data/user1/yy/${EXPERIMENT}_vllm_tp${tp}_dp${dp}.pid"

  if [[ -f "$pid_file" ]]; then
    local pid
    pid="$(cat "$pid_file")"
    if kill -0 "$pid" 2>/dev/null; then
      log "stopping vllm server pid=${pid} (tp=${tp}, dp=${dp})"
      kill "$pid" 2>/dev/null || true
      sleep 5
      kill -9 "$pid" 2>/dev/null || true
    fi
    rm -f "$pid_file"
  fi

  # Fallback: kill any vllm serve processes for this model.
  pkill -9 -f "vllm serve.*${MODEL_NAME}" 2>/dev/null || true
  pkill -9 -f "vllm serve.*${MODEL_PATH}" 2>/dev/null || true
  sleep 2
}

build_server_args() {
  local tp="$1"
  local dp="$2"

  local args=(
    "vllm serve" "$MODEL_PATH"
    --trust-remote-code
    --kv-cache-dtype fp8
    --block-size 256
    --tensor-parallel-size "$tp"
    --no-enable-flashinfer-autotune
    --port "$VLLM_PORT"
  )
  if [[ "$dp" -gt 1 ]]; then
    args+=(
      --data-parallel-size "$dp"
      --data-parallel-size-local "$dp"
      --data-parallel-multi-port-external-lb
    )
  fi
  printf '%s ' "${args[@]}"
}

start_server() {
  local tp="$1"
  local dp="$2"

  log "starting vllm server tp=${tp} dp=${dp}"
  bash "${SCRIPT_DIR}/start_vllm_dp.sh" "$tp" "$dp" \
    >> "${log_dir_global}/vllm_tp${tp}_dp${dp}.server.outer.log" 2>&1

  local hport
  hport="$(health_port_for "$dp")"
  if ! is_server_healthy "$hport"; then
    log "error: vllm server tp=${tp} dp=${dp} failed health check on port ${hport}"
    return 1
  fi
  log "vllm server tp=${tp} dp=${dp} is healthy on port ${hport}"
}

run_warmup() {
  local tp="$1"
  local dp="$2"
  local input_len="$3"
  local output_len="$4"
  local hport
  hport="$(health_port_for "$dp")"

  log "warming up tp=${tp} dp=${dp} (input=${input_len}, output=${output_len}, num=1)"
  "$PYTHON" "${SCRIPT_DIR}/../../scripts/common/warmup.py" \
    --backend vllm \
    --host 127.0.0.1 \
    --port "$hport" \
    --input-len "$input_len" \
    --output-len "$output_len" \
    --num 1 \
    --env-python "$PYTHON" \
    >> "${log_dir_global}/vllm_tp${tp}_dp${dp}.warmup.log" 2>&1
  log "warmup completed"
}

scenario_already_completed() {
  local output_file="$1"
  local expected="$2"
  [[ -s "$output_file" ]] || return 1
  local completed
  completed="$("$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" ]]
}

start_gpu_monitor() {
  local csv_path="$1"
  mkdir -p "$(dirname "$csv_path")"
  nvidia-smi \
    --query-gpu=timestamp,index,memory.used,memory.total,utilization.gpu \
    --format=csv \
    -l "$GPU_MEM_SAMPLE_INTERVAL_S" \
    > "$csv_path" 2>/dev/null &
  echo $!
}

stop_gpu_monitor() {
  local pid="$1"
  if kill -0 "$pid" 2>/dev/null; then
    kill "$pid" 2>/dev/null || true
    sleep 1
    kill -9 "$pid" 2>/dev/null || true
  fi
}

append_scenario_record() {
  local result_root="$1"
  local json_path="$result_root/results.json"
  shift
  # Remaining args are key=value pairs.
  local scenario_json
  scenario_json="$("$PYTHON" -c "
import json, sys
pairs = [a.split('=', 1) for a in sys.argv[1:]]
d = {k: json.loads(v) for k, v in pairs}
print(json.dumps(d, ensure_ascii=False))
" "$@")"
  PYTHON="$PYTHON" append_scenario_to_json "$json_path" "$scenario_json"
}

# ---------------------------------------------------------------------------
# Per-configuration runner
# ---------------------------------------------------------------------------

run_parallel_config() {
  local tp="$1"
  local dp="$2"
  local config_label="tp${tp}_dp${dp}"
  local result_root="${RESULT_BASE}/${RUN_ID}/${config_label}"
  local raw_dir="${result_root}/raw_outputs"
  local gpu_log_dir="${result_root}/gpu_logs"
  local phase_log_dir="${result_root}/logs"
  mkdir -p "$raw_dir" "$gpu_log_dir" "$phase_log_dir"

  log "===== ${config_label} START ====="

  # Generate scenario list for this config.
  local scenario_tsv="${result_root}/scenarios.tsv"
  "$PYTHON" "${SCRIPT_DIR}/generate_scenarios.py" \
    --matrix "$MATRIX_FILE" \
    --mode "$MATRIX_MODE" \
    > "$scenario_tsv"
  local total_scenarios
  total_scenarios="$(tail -n +2 "$scenario_tsv" | wc -l)"
  log "generated ${total_scenarios} scenarios for ${config_label}"

  if [[ "$total_scenarios" -eq 0 ]]; then
    log "no scenarios for ${config_label}; skipping"
    return 0
  fi

  # Write metadata.
  ensure_result_root "$result_root"
  write_metadata_json \
    "${result_root}/results.json" \
    "${EXPERIMENT_NAME}_${config_label}" \
    "$RUN_ID" \
    "$MODEL_PATH" \
    "vllm" \
    "vllm" \
    "$HARDWARE" \
    "$ACCELERATOR" \
    "$CHIP" \
    "experiments/${EXPERIMENT_NAME}/run_bench.sh" \
    "$VENV_VLLM" \
    "H200 vLLM TP×DP matrix for DeepSeek-V4-Flash"

  # Embed static config and the exact server command line used.
  local server_args_str
  server_args_str="$(build_server_args "$tp" "$dp")"
  jq --arg tp "$tp" --arg dp "$dp" --arg cuda "$CUDA_VISIBLE_DEVICES" --arg args "$server_args_str" \
    '.config = {
       "tp": ($tp | tonumber),
       "dp": ($dp | tonumber),
       "cuda_visible_devices": $cuda,
       "backend": "vllm",
       "server_start_script": "experiments/'${EXPERIMENT_NAME}'/start_vllm_dp.sh",
       "server_args": $args
     }' "${result_root}/results.json" > "${result_root}/results.json.tmp" && \
    mv "${result_root}/results.json.tmp" "${result_root}/results.json"

  # Group scenarios by input_len.  We restart the server for each ISL to keep
  # memory state isolated and to warm up per context length.
  local current_isl=""
  local group_started=false

  tail -n +2 "$scenario_tsv" | while IFS=$'\t' read -r mark isl dsl conc num; do
    # When the input length changes, restart the server for the new group.
    if [[ "$isl" != "$current_isl" ]]; then
      if [[ "$group_started" == true ]]; then
        stop_server "$tp" "$dp"
      fi
      current_isl="$isl"

      stop_server "$tp" "$dp"
      if ! start_server "$tp" "$dp"; then
        log "ERROR: ${config_label} failed to start for ISL=${isl}; skipping this group"
        group_started=false
        continue
      fi
      group_started=true

      # Warmup with the shortest output for this ISL.
      run_warmup "$tp" "$dp" "$isl" 128
    fi

    # If the server for this group did not start, skip all scenarios in it.
    if [[ "$group_started" != true ]]; then
      log "skipping ${config_label} ISL=${isl} scenario (server not started)"
      continue
    fi

    local output_file="${raw_dir}/vllm_main_c${conc}_i${isl}_o${dsl}.jsonl"
    local detail_log="${phase_log_dir}/vllm_${config_label}_c${conc}_i${isl}_o${dsl}.log"
    local gpu_csv="${gpu_log_dir}/gpu_mem_c${conc}_i${isl}_o${dsl}.csv"

    if scenario_already_completed "$output_file" "$num"; then
      log "skipping already-completed ${config_label} scenario: c=${conc} i=${isl} o=${dsl}"
      continue
    fi

    log "running ${config_label} scenario: mark=${mark} c=${conc} i=${isl} o=${dsl} n=${num}"

    local gpu_pid
    gpu_pid="$(start_gpu_monitor "$gpu_csv")"

    local bench_rc=0
    timeout "$SCENARIO_TIMEOUT_S" \
      "$PYTHON" -m sglang.bench_serving \
        --backend vllm \
        --host 127.0.0.1 \
        --port "$(health_port_for "$dp")" \
        --dataset-name random \
        --random-input-len "$isl" \
        --random-output-len "$dsl" \
        --num-prompts "$num" \
        --max-concurrency "$conc" \
        --request-rate 10000 \
        --output-file "$output_file" \
        --output-details \
        > "$detail_log" 2>&1 || bench_rc=$?

    stop_gpu_monitor "$gpu_pid"

    if [[ "$bench_rc" -ne 0 ]]; then
      log "ERROR: ${config_label} scenario c=${conc} i=${isl} o=${dsl} failed (rc=${bench_rc}); see ${detail_log}"
      if [[ "$mark" == "P" ]]; then
        log "optional (P) scenario failed; recording as skipped_oom and continuing"
        append_scenario_record "$result_root" \
          "name=c${conc}_i${isl}_o${dsl}" \
          "config=$(jq -n --arg phase main --argjson c "$conc" --argjson i "$isl" --argjson o "$dsl" --arg dataset random --argjson n "$num" '{phase: $phase, concurrency: $c, input_len: $i, output_len: $o, dataset: $dataset, num_prompts: $n}')" \
          "status=\"skipped_oom\"" \
          "note=\"bench command failed or timed out (rc=${bench_rc})\""
        continue
      else
        log "mandatory (Y) scenario failed; aborting current ISL group"
        stop_server "$tp" "$dp"
        group_started=false
        continue
      fi
    fi

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

  stop_server "$tp" "$dp"

  # Parse results.
  log "parsing ${config_label} results"
  "$PYTHON" "${SCRIPT_DIR}/../../scripts/common/parse_backend.py" "$result_root" --backend vllm \
    >> "${phase_log_dir}/parse.log" 2>&1 || {
      log "WARNING: parser failed for ${config_label}; see ${phase_log_dir}/parse.log"
    }

  log "===== ${config_label} DONE ====="
}

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

# Cleanup any leftovers.
for cfg in "${PARALLEL_CONFIGS[@]}"; do
  read -r tp dp <<< "$cfg"
  stop_server "$tp" "$dp"
done

# Run each parallel configuration.
for cfg in "${PARALLEL_CONFIGS[@]}"; do
  read -r tp dp <<< "$cfg"
  run_parallel_config "$tp" "$dp"
done

# Generate cross-configuration comparison.
log "generating comparison report"
"$PYTHON" "${SCRIPT_DIR}/compare.py" \
  --run-root "${RESULT_BASE}/${RUN_ID}" \
  --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}"