#!/usr/bin/env bash
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}/../../../scripts/common/server_docker.sh"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../../scripts/common/bench_client_docker.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}}"
RAW_DIR="${RESULT_ROOT}/raw_outputs"
LOG_DIR="${RESULT_ROOT}/logs"

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

log "experiment=${EXPERIMENT_NAME}"
log "run_id=${RUN_ID}"
log "result_root=${RESULT_ROOT}"
log "platform=${PLATFORM}"
log "chip=${CHIP}"
log "accelerator=${ACCELERATOR}"
log "engine=${ENGINE}"
log "hardware=${HARDWARE}"

# Write initial metadata for this run.
METADATA_JSON="${RESULT_ROOT}/results.json"
write_metadata_json \
  "$METADATA_JSON" \
  "$EXPERIMENT_NAME" \
  "$RUN_ID" \
  "$MODEL_PATH" \
  "$BACKEND" \
  "$ENGINE" \
  "$HARDWARE" \
  "$ACCELERATOR" \
  "$CHIP" \
  "experiments/${EXPERIMENT_NAME}/run_bench.sh" \
  "" \
  "Kunlun P800 SGLang benchmark for DeepSeek-V4-Flash-INT8"

# Start server unless SKIP_MANAGE_SERVER is set.
if [[ -n "${SKIP_MANAGE_SERVER:-}" ]]; then
  log "SKIP_MANAGE_SERVER is set, assuming server is already running on port ${PORT}"
  if ! health_check 127.0.0.1 "${PORT}" 30; then
    log "ERROR: no healthy server found at port ${PORT}"
    exit 1
  fi
else
  docker_server_start "${MODEL_PATH}" "${PORT}" "${LOG_DIR}/server.outer.log" "${SERVER_MODE}"
fi

on_exit() {
  local code=$?
  if [[ -z "${SKIP_MANAGE_SERVER:-}" ]]; then
    docker_server_stop
  fi
  log "orchestrator exiting with code=${code}"
  exit "$code"
}
trap on_exit EXIT

log "===== BENCHMARK START ====="

for scenario in "${SCENARIOS[@]}"; do
  # Allow SCENARIOS to be passed as a string like "(32 512 256 160)" from env.
  scenario="${scenario//[()]/}"
  read -r concurrency input_len output_len num_prompts <<< "$scenario"
  # Fallback for legacy 3-field scenarios.
  if [[ -z "${num_prompts:-}" ]]; then
    num_prompts="$NUM_PROMPTS"
  fi

  tmp_output_file="/tmp/bench_outputs/${CHIP}_${ENGINE}_$(date '+%m%d')_${concurrency}_${input_len}_${output_len}.jsonl"
  host_output_file="${RAW_DIR}/${CHIP}_${ENGINE}_$(date '+%m%d')_${concurrency}_${input_len}_${output_len}.jsonl"
  detail_log="${LOG_DIR}/${BACKEND}_c${concurrency}_i${input_len}_o${output_len}.log"

  log "running scenario: concurrency=${concurrency} input=${input_len} output=${output_len} num_prompts=${num_prompts}"

  run_random_case \
    "$BACKEND" \
    "$PORT" \
    "$MODEL_PATH" \
    "$tmp_output_file" \
    "$concurrency" \
    "$input_len" \
    "$output_len" \
    "$num_prompts" \
    "${DATASET_PATH:-}" \
    > "$detail_log" 2>&1 || {
      log "ERROR: scenario c=${concurrency} i=${input_len} o=${output_len} failed"
      continue
    }

  # Copy the JSONL from the container to the host result directory.
  docker cp "${CONTAINER_NAME}:${tmp_output_file}" "$host_output_file" || {
    log "ERROR: failed to copy ${tmp_output_file} from container"
    continue
  }
  log "finished scenario: output=${host_output_file}"
done

log "===== BENCHMARK DONE ====="
log "parsing results..."
"${SCRIPT_DIR}/parse_results.py" "$RESULT_ROOT"

log "all results saved to ${RESULT_ROOT}"