#!/usr/bin/env bash
# Find the Total-TPS saturation concurrency for each SGLang TP/DP/ISL/OSL shape.
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"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/adaptive_config.env"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../../scripts/common/adaptive_bench_lib.sh"

ENGINE="sglang"
ENGINE_PORT="$SGLANG_PORT"
RESULT_BASE="${RESULT_BASE:-${SCRIPT_DIR}/adaptive_results}"
ACTIVE_ENGINE_SERVER_LOG=""

if [[ -x "${VENV_CLIENT}/bin/python" ]]; then
  PYTHON="${VENV_CLIENT}/bin/python"
else
  PYTHON="$(command -v python3)"
fi

DOCKER_IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:latest}"

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

engine_stop_server() {
  local tp="$1"
  local dp="$2"
  local pid_file="${RUNTIME_BASE}/${EXPERIMENT}_sglang_tp${tp}_dp${dp}.pid"

  if [[ -f "$pid_file" ]]; then
    local pid
    pid="$(cat "$pid_file")"
    if [[ -n "${CONTAINER_NAME:-}" ]]; then
      if docker exec "$CONTAINER_NAME" kill -0 "$pid" 2>/dev/null; then
        log "stopping sglang in persistent container pid=${pid} tp=${tp} dp=${dp}"
        docker exec "$CONTAINER_NAME" kill "$pid" 2>/dev/null || true
        sleep 5
        docker exec "$CONTAINER_NAME" kill -9 "$pid" 2>/dev/null || true
      fi
    elif kill -0 "$pid" 2>/dev/null; then
      log "stopping sglang 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
  if [[ -z "${CONTAINER_NAME:-}" ]]; then
    docker rm -f "${EXPERIMENT}_sglang_tp${tp}_dp${dp}" >/dev/null 2>&1 || true
  fi
  ACTIVE_ENGINE_SERVER_LOG=""
  sleep 2
}

engine_build_server_args() {
  local tp="$1"
  local dp="$2"
  local -a args=(
    sglang serve --model-path "$MODEL_PATH"
    --trust-remote-code
    --tp-size "$tp"
    --moe-runner-backend "$MOE_RUNNER_BACKEND"
    --context-length "$CONTEXT_LENGTH"
    --max-running-requests "$MAX_RUNNING_REQUESTS"
    --mem-fraction-static "$MEM_FRACTION_STATIC"
    --host 0.0.0.0
    --port "$ENGINE_PORT"
  )
  if (( dp > 1 )); then
    args+=(--dp-size "$dp")
  fi
  printf '%q ' "${args[@]}"
}

engine_start_server() {
  local tp="$1"
  local dp="$2"
  local outer_log="${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log"
  log "starting sglang server tp=${tp} dp=${dp}"
  if [[ -n "${CONTAINER_NAME:-}" ]]; then
    bash "${SCRIPT_DIR}/run_sglang_in_container.sh" "$tp" "$dp" >> "$outer_log" 2>&1
  else
    bash "${SCRIPT_DIR}/start_sglang_dp.sh" "$tp" "$dp" >> "$outer_log" 2>&1
  fi
  if ! engine_is_healthy; then
    log "ERROR: sglang health check failed tp=${tp} dp=${dp}"
    return 1
  fi
  if [[ -z "${CONTAINER_NAME:-}" ]]; then
    ACTIVE_ENGINE_SERVER_LOG="$(
      find "${RUNTIME_BASE}/logs" -maxdepth 1 -type f \
        -name "${EXPERIMENT}_sglang*tp${tp}_dp${dp}_*.log" \
        -printf '%T@ %p\n' 2>/dev/null | sort -nr | head -n 1 | cut -d' ' -f2-
    )"
  fi
  log "sglang server healthy tp=${tp} dp=${dp} log=${ACTIVE_ENGINE_SERVER_LOG:-container:/tmp/sglang_server.log}"
}

engine_detect_oom() {
  local detail_log="$1"
  local tp="$2"
  local dp="$3"
  local pattern='CUDA out of memory|torch\.OutOfMemoryError|OutOfMemory|out of memory|OOM|RESOURCE_EXHAUSTED|Failed to allocate memory'
  local outer_log="${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log"
  local -a logs=("$detail_log" "$outer_log")
  if [[ -n "$ACTIVE_ENGINE_SERVER_LOG" ]]; then
    logs+=("$ACTIVE_ENGINE_SERVER_LOG")
  fi
  if grep -Eiq "$pattern" "${logs[@]}" 2>/dev/null; then
    return 0
  fi
  if [[ -n "${CONTAINER_NAME:-}" ]]; then
    docker exec "$CONTAINER_NAME" grep -Eiq "$pattern" /tmp/sglang_server.log 2>/dev/null
    return $?
  fi
  return 1
}

engine_run_bench() {
  local isl="$1"
  local osl="$2"
  local concurrency="$3"
  local num_prompts="$4"
  local output_file="$5"
  local warmup_requests
  warmup_requests="$(adaptive_warmup_request_count "$concurrency")"
  local -a bench_args=(
    --backend sglang
    --host 127.0.0.1
    --port "$ENGINE_PORT"
    --dataset-name "$BENCH_DATASET_NAME"
    --random-input-len "$isl"
    --random-output-len "$osl"
    --random-range-ratio "$RANDOM_RANGE_RATIO"
    --num-prompts "$num_prompts"
    --max-concurrency "$concurrency"
    --request-rate 10000
    --warmup-requests "$warmup_requests"
    --output-file "$output_file"
    --output-details
    --disable-tqdm
  )
  if [[ "$BENCH_DATASET_NAME" == "random" ]]; then
    bench_args+=(--dataset-path "$DATASET_PATH")
  else
    bench_args+=(--tokenize-prompt)
  fi

  if [[ "$USE_DOCKER_CLIENT" == "1" ]]; then
    local -a volume_args=(-v "${MODEL_PATH}:${MODEL_PATH}:ro" -v "${RESULT_BASE}:${RESULT_BASE}")
    if [[ "$BENCH_DATASET_NAME" == "random" ]]; then
      volume_args+=(-v "${DATASET_PATH}:${DATASET_PATH}:ro")
    fi
    docker run --rm \
      --network host \
      "${volume_args[@]}" \
      -e PYTHONUNBUFFERED=1 \
      "$DOCKER_IMAGE" \
      python -m sglang.bench_serving "${bench_args[@]}"
  else
    "$PYTHON" -m sglang.bench_serving "${bench_args[@]}"
  fi
}

export -f engine_run_bench
export ENGINE_PORT MODEL_PATH RESULT_BASE DOCKER_IMAGE USE_DOCKER_CLIENT
export BENCH_DATASET_NAME DATASET_PATH RANDOM_RANGE_RATIO BENCH_WARMUP_MAX_REQUESTS PYTHON

adaptive_main "$@"
