#!/usr/bin/env bash
# Max context length exploration for SGLang and vLLM 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_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}"
LOG_DIR="${RESULT_ROOT}/logs"
RAW_DIR="${RESULT_ROOT}/raw_outputs"

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

log "experiment=${EXPERIMENT_NAME}"
log "run_id=${RUN_ID}"
log "platform=${PLATFORM}"
log "hardware=${HARDWARE}"
log "model=${MODEL_PATH}"
log "candidate_lengths=${CANDIDATE_LENS[*]}"

# ---------------------------------------------------------------------------
# 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/${EXPERIMENT_NAME}_${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
  sleep 2
}

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

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

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

build_server_args() {
  local backend="$1"
  local target_len="$2"
  local max_len=$(( target_len + CONTEXT_PAD ))
  if [[ "$backend" == "sglang" ]]; then
    echo "sglang serve --trust-remote-code --model-path $MODEL_PATH --tp $TP --moe-runner-backend marlin --context-length $max_len --max-running-requests 1 --mem-fraction-static 0.88 --host 0.0.0.0 --port $SGLANG_PORT"
  else
    echo "vllm serve $MODEL_PATH --trust-remote-code --tensor-parallel-size $TP --kv-cache-dtype fp8 --max-model-len $max_len --max-num-seqs 1 --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
}

append_scenario() {
  local scenario_json="$1"
  "${VENV_CLIENT}/bin/python" - "$RESULT_JSON" "$scenario_json" <<'PY'
import json
import sys

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

data["scenarios"].append(json.loads(scenario_json))
with open(results_path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
PY
}

run_backend() {
  local backend="$1"
  local port
  if [[ "$backend" == "sglang" ]]; then
    port="$SGLANG_PORT"
  else
    port="$VLLM_PORT"
  fi

  log "===== ${backend} max context exploration START ====="

  for target_len in "${CANDIDATE_LENS[@]}"; do
    local max_len=$(( target_len + CONTEXT_PAD ))
    local output_file="${RAW_DIR}/${backend}_ctx_${target_len}.jsonl"
    local detail_log="${LOG_DIR}/${backend}_ctx_${target_len}.log"

    log "trying ${backend} target_len=${target_len} (max_model/context_len=${max_len})"

    stop_server "$backend"
    if ! start_server "$backend" "$target_len"; then
      local err="server failed to start"
      log "ERROR: ${err}"
      append_scenario "$(jq -n \
        --arg backend "$backend" \
        --argjson target_len "$target_len" \
        --argjson max_model_len "$max_len" \
        --arg server_args "$(build_server_args "$backend" "$target_len")" \
        --arg error "$err" \
        '{backend:$backend, target_len:$target_len, max_model_len:$max_model_len, success:false, server_args:$server_args, metrics:null, error:$error}')"
      break
    fi

    local bench_rc=0
    "${VENV_CLIENT}/bin/python" -m sglang.bench_serving \
      --backend vllm \
      --host 127.0.0.1 \
      --port "$port" \
      --dataset-name random \
      --random-input-len "$target_len" \
      --random-output-len "$OUTPUT_LEN" \
      --random-range-ratio 1.0 \
      --num-prompts "$NUM_PROMPTS" \
      --max-concurrency "$MAX_CONCURRENCY" \
      --request-rate "$REQUEST_RATE" \
      --output-file "$output_file" \
      --output-details \
      > "$detail_log" 2>&1 || bench_rc=$?

    if [[ "$bench_rc" -ne 0 ]]; then
      local err="bench_serving exited with code ${bench_rc}; see ${detail_log}"
      log "ERROR: ${backend} target_len=${target_len} failed; ${err}"
      append_scenario "$(jq -n \
        --arg backend "$backend" \
        --argjson target_len "$target_len" \
        --argjson max_model_len "$max_len" \
        --arg server_args "$(build_server_args "$backend" "$target_len")" \
        --arg error "$err" \
        '{backend:$backend, target_len:$target_len, max_model_len:$max_model_len, success:false, server_args:$server_args, metrics:null, error:$error}')"
      stop_server "$backend"
      break
    fi

    # Verify the request actually completed.
    local completed
    completed="$(${VENV_CLIENT}/bin/python -c "
import json
with open('${output_file}') as f:
    for line in f:
        data = json.loads(line)
        print(data.get('completed', 0))
        break
")"
    if [[ "${completed:-0}" -lt "$NUM_PROMPTS" ]]; then
      local err="only ${completed}/${NUM_PROMPTS} requests completed"
      log "ERROR: ${backend} target_len=${target_len} ${err}"
      append_scenario "$(jq -n \
        --arg backend "$backend" \
        --argjson target_len "$target_len" \
        --argjson max_model_len "$max_len" \
        --arg server_args "$(build_server_args "$backend" "$target_len")" \
        --arg error "$err" \
        '{backend:$backend, target_len:$target_len, max_model_len:$max_model_len, success:false, server_args:$server_args, metrics:null, error:$error}')"
      stop_server "$backend"
      break
    fi

    local metrics_json
    metrics_json="$(${VENV_CLIENT}/bin/python "${SCRIPT_DIR}/extract_metrics.py" "$output_file")"

    append_scenario "$(jq -n \
      --arg backend "$backend" \
      --argjson target_len "$target_len" \
      --argjson max_model_len "$max_len" \
      --arg server_args "$(build_server_args "$backend" "$target_len")" \
      --argjson metrics "$metrics_json" \
      '{backend:$backend, target_len:$target_len, max_model_len:$max_model_len, success:true, server_args:$server_args, metrics:$metrics, error:null}')"

    log "${backend} target_len=${target_len} succeeded"
    stop_server "$backend"
  done

  log "===== ${backend} max context exploration DONE ====="
}

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

mkdir -p "$RAW_DIR" "$LOG_DIR"

RESULT_JSON="${RESULT_ROOT}/results.json"
write_metadata_json \
  "$RESULT_JSON" \
  "$EXPERIMENT_NAME" \
  "$RUN_ID" \
  "$MODEL_PATH" \
  "sglang+vllm" \
  "sglang+vllm" \
  "$HARDWARE" \
  "$ACCELERATOR" \
  "$CHIP" \
  "experiments/${EXPERIMENT_NAME}/run_bench.sh" \
  "$VENV_CLIENT" \
  "H200 max context length exploration for DeepSeek-V4-Flash"

# Embed config.
"${VENV_CLIENT}/bin/python" - "$RESULT_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": 8,
    "cuda_visible_devices": "0,1,2,3,4,5,6,7",
    "output_len": 1,
    "num_prompts": 1,
    "max_concurrency": 1,
    "request_rate": 10000,
    "context_pad": 1024,
}
with open(path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
PY

# Ensure jq is available.
if ! command -v jq >/dev/null 2>&1; then
  log "ERROR: jq is required for JSON manipulation in this script"
  exit 1
fi

stop_server sglang
stop_server vllm

run_backend sglang
run_backend vllm

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

log "all results saved to ${RESULT_ROOT}"