#!/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}"
log "server_mode=${SERVER_MODE}"
log "extra_launch_args=${SGLANG_EXTRA_LAUNCH_ARGS:-<none>}"

# 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 256k-input 4k-output OOM probe for DeepSeek-V4-Flash-INT8"

# Embed config into metadata.
"${PYTHON:-python3}" - "$METADATA_JSON" <<'PY' >/dev/null
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,
    "xpu_visible_devices": "0,1,2,3,4,5,6,7",
    "context_length": 270000,
    "server_mode": "w8a8_int8_baseline",
    "server_start_script": f"experiments/{data['metadata']['experiment']}/start_server.sh",
    "extra_launch_args": "--context-length 270000 --max-running-requests 2 --mem-fraction-static 0.85",
    "scenarios": ["2 262144 4096"],
    "num_prompts": 4
}
with open(path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
PY

# Allow the benchmark client to run in a different container than the managed
# server (useful when reusing an already-running deployment).
BENCH_CONTAINER_NAME="${BENCH_CONTAINER_NAME:-${CONTAINER_NAME}}"
if [[ -n "${SKIP_MANAGE_SERVER:-}" && "$BENCH_CONTAINER_NAME" != "$CONTAINER_NAME" ]]; then
  log "using benchmark client container: ${BENCH_CONTAINER_NAME}"
  CONTAINER_NAME="$BENCH_CONTAINER_NAME"
fi

# 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}" "${SGLANG_EXTRA_LAUNCH_ARGS:-}"
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 ====="

# Optional small warmup to exercise the long-context path. Disabled by default
# because the probe scenario itself is already small (4 prompts).
if [[ -z "${SKIP_WARMUP:-}" ]]; then
  log "running small warmup (1 prompt, 64k input, 256 output)"
  tmp_warmup_file="/tmp/bench_outputs/${CHIP}_${ENGINE}_$(date '+%m%d')_warmup_1_65536_256.jsonl"
  host_warmup_file="${RAW_DIR}/${CHIP}_${ENGINE}_$(date '+%m%d')_warmup_1_65536_256.jsonl"
  detail_log="${LOG_DIR}/warmup.log"

  if WARMUP=0 NUM_PROMPTS=1 run_random_case \
      "$BACKEND" "$PORT" "$MODEL_PATH" "$tmp_warmup_file" \
      1 65536 256 1 "${DATASET_PATH:-}" \
      > "$detail_log" 2>&1; then
    docker cp "${CONTAINER_NAME}:${tmp_warmup_file}" "$host_warmup_file" 2>/dev/null || true
    log "warmup completed"
  else
    log "WARNING: warmup failed; continuing with main scenario"
  fi
fi

for scenario in "${SCENARIOS[@]}"; do
  # Allow SCENARIOS to be passed as a string like "(2 262144 4096)" from env.
  scenario="${scenario//[()]/}"
  read -r concurrency input_len output_len <<< "$scenario"

  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}"

  if run_random_case \
      "$BACKEND" "$PORT" "$MODEL_PATH" "$tmp_output_file" \
      "$concurrency" "$input_len" "$output_len" "$NUM_PROMPTS" \
      "${DATASET_PATH:-}" \
      > "$detail_log" 2>&1; then
    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}"
  else
    log "ERROR: scenario c=${concurrency} i=${input_len} o=${output_len} failed; see ${detail_log}"
  fi
done

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

log "all results saved to ${RESULT_ROOT}"