#!/usr/bin/env bash
# Helpers to run sglang.bench_serving inside the P800 Docker container.
# Usage: source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/../scripts/common/bench_client_docker.sh"

set -Eeuo pipefail

_BENCH_CLIENT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=/dev/null
source "${_BENCH_CLIENT_DIR}/lib.sh"
# shellcheck source=/dev/null
source "${_BENCH_CLIENT_DIR}/platform.sh"

# Required platform variables: CONTAINER_NAME, CONTAINER_PYTHON

# Run sglang.bench_serving inside the running container.
# All arguments are forwarded to bench_serving.
run_bench_in_container() {
  local container="${CONTAINER_NAME}"

  if ! docker inspect "$container" >/dev/null 2>&1; then
    log "ERROR: container ${container} is not running"
    return 1
  fi

  log "running bench_serving in container ${container}"
  docker exec "$container" \
    env HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_DATASETS_OFFLINE=1 \
    "${CONTAINER_PYTHON}" -m sglang.bench_serving "$@"
}

# Convenience wrapper for a single random-dataset case.
# Args:
#   $1: backend (e.g. sglang)
#   $2: port
#   $3: model path or served model name
#   $4: output jsonl path (inside the container)
#   $5: concurrency
#   $6: input length
#   $7: output length
#   $8: num prompts (optional, default 512)
#   $9: dataset path inside container (optional)
run_random_case() {
  local backend="$1"
  local port="$2"
  local model="$3"
  local output_file="$4"
  local concurrency="$5"
  local input_len="$6"
  local output_len="$7"
  local num_prompts="${8:-512}"
  local dataset_path="${9:-}"
  local warmup="${WARMUP:-100}"

  # Ensure the output directory exists inside the container.
  docker exec "${CONTAINER_NAME}" mkdir -p "$(dirname "$output_file")"

  local args=(
    --backend "$backend"
    --host 127.0.0.1
    --port "$port"
    --model "$model"
    --dataset-name random
    --random-input-len "$input_len"
    --random-output-len "$output_len"
    --num-prompts "$num_prompts"
    --max-concurrency "$concurrency"
    --warmup-requests "$warmup"
    --output-file "$output_file"
    --output-details
  )
  if [[ -n "$dataset_path" ]]; then
    args+=(--dataset-path "$dataset_path")
  fi

  run_bench_in_container "${args[@]}"
}