#!/usr/bin/env bash
# Long-context matrix for SGLang on DeepSeek-V4-Flash-INT8 (Kunlun P800).
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_BASE="${SCRIPT_DIR}/results"

log_dir_global="${RESULT_BASE}/${RUN_ID}/logs"
mkdir -p "$log_dir_global"
log_init "${log_dir_global}/orchestrator.log"

log "experiment=${EXPERIMENT_NAME}"
log "run_id=${RUN_ID}"
log "platform=${PLATFORM}"
log "hardware=${HARDWARE}"
log "model=${MODEL_PATH}"
log "server_mode=${SERVER_MODE}"
log "groups=${#CONTEXT_GROUPS[@]}"

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

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

start_server() {
  local context_length="$1"
  local max_running="$2"
  local extra_args="--context-length ${context_length} --max-running-requests ${max_running} --mem-fraction-static 0.85"
  export SGLANG_EXTRA_LAUNCH_ARGS="${SGLANG_EXTRA_LAUNCH_ARGS:-${extra_args}}"

  log "starting server (context_length=${context_length}, max_running=${max_running}, mode=${SERVER_MODE})"
  docker_server_start "${MODEL_PATH}" "${PORT}" "${log_dir_global}/server_${context_length}.outer.log" "${SERVER_MODE}" "${SGLANG_EXTRA_LAUNCH_ARGS}"
}

stop_server() {
  docker_server_stop
}

run_warmup() {
  local raw_dir="$1"
  local phase_log_dir="$2"
  local input_len="$3"
  local output_len="${4:-256}"

  log "warming up (input=${input_len}, output=${output_len}, num=1)"

  local tmp_warmup_file="/tmp/bench_outputs/${CHIP}_${ENGINE}_$(date '+%m%d')_warmup_1_${input_len}_${output_len}.jsonl"
  local host_warmup_file="${raw_dir}/${BACKEND}_warmup_$(date '+%m%d')_1_${input_len}_${output_len}.jsonl"
  local detail_log="${phase_log_dir}/warmup_${input_len}_${output_len}.log"

  if WARMUP=0 run_random_case \
      "$BACKEND" "$PORT" "$MODEL_PATH" "$tmp_warmup_file" \
      1 "$input_len" "$output_len" 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 group"
  fi
}

scenario_already_completed() {
  local output_file="$1"
  local expected="$2"
  [[ -s "$output_file" ]] || return 1
  local completed
  completed="$(python3 -c "
import json, sys
path = sys.argv[1]
try:
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:
                data = json.loads(line)
                print(data.get('completed', 0))
                break
except Exception:
    print(0)
" "$output_file")"
  [[ "${completed:-0}" -ge "$expected" ]]
}

run_group() {
  local group_label="$1"
  local context_length="$2"
  local max_running="$3"
  local -n scenarios_ref="$4"

  local result_root="${RESULT_BASE}/${RUN_ID}/${BACKEND}"
  local raw_dir="${result_root}/raw_outputs"
  local phase_log_dir="${result_root}/logs"

  mkdir -p "$raw_dir" "$phase_log_dir"

  log "===== ${BACKEND} ${group_label} START (context_length=${context_length}) ====="

  stop_server
  if ! start_server "$context_length" "$max_running"; then
    log "ERROR: server failed to start for ${group_label}; skipping group"
    stop_server
    return 1
  fi

  # Warmup with the shortest output for this input length to save time.
  local warmup_input="${scenarios_ref[0]}"
  local warmup_input_len
  warmup_input_len="$(echo "$warmup_input" | awk '{print $2}')"
  run_warmup "$raw_dir" "$phase_log_dir" "$warmup_input_len" 256

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

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

    if scenario_already_completed "$output_file" "$num_prompts"; then
      log "skipping already-completed ${BACKEND} ${group_label} scenario: c=${concurrency} i=${input_len} o=${output_len}"
      continue
    fi

    log "running ${BACKEND} ${group_label} scenario: c=${concurrency} i=${input_len} o=${output_len} n=${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}" "$output_file" || {
        log "ERROR: failed to copy ${tmp_output_file} from container"
        continue
      }
      log "finished ${BACKEND} ${group_label} scenario: output=${output_file}"
    else
      log "ERROR: ${BACKEND} ${group_label} scenario c=${concurrency} i=${input_len} o=${output_len} failed; see ${detail_log}"
      continue
    fi
  done

  stop_server
  log "===== ${BACKEND} ${group_label} DONE ====="
}

parse_backend() {
  local result_root="${RESULT_BASE}/${RUN_ID}/${BACKEND}"
  log "parsing ${BACKEND} results in ${result_root}"
  python3 "${SCRIPT_DIR}/../../scripts/common/parse_backend.py" "$result_root" --backend "$BACKEND" \
    >> "${result_root}/logs/parse.log" 2>&1 || {
      log "WARNING: parser failed for ${BACKEND}; see ${result_root}/logs/parse.log"
    }
}

write_backend_metadata() {
  local result_root="${RESULT_BASE}/${RUN_ID}/${BACKEND}"
  ensure_result_root "$result_root"
  local meta_json="${result_root}/results.json"

  write_metadata_json \
    "$meta_json" \
    "${EXPERIMENT_NAME}_${BACKEND}" \
    "$RUN_ID" \
    "$MODEL_PATH" \
    "$BACKEND" \
    "$ENGINE" \
    "$HARDWARE" \
    "$ACCELERATOR" \
    "$CHIP" \
    "experiments/${EXPERIMENT_NAME}/run_bench.sh" \
    "" \
    "Kunlun P800 long-context matrix ${BACKEND} TP=8 for DeepSeek-V4-Flash-INT8"

  # Build groups JSON from CONTEXT_GROUPS.
  local groups_json="["
  local first=1
  for group in "${CONTEXT_GROUPS[@]}"; do
    read -r label context_length max_running <<< "$group"
    if [[ "$first" -eq 1 ]]; then
      first=0
    else
      groups_json+=","
    fi
    groups_json+="{\"label\":\"${label}\",\"context_length\":${context_length},\"max_running\":${max_running}}"
  done
  groups_json+="]"

  # Embed config.
  python3 - "$meta_json" "$groups_json" <<'PY' >/dev/null
import json
import sys

path = sys.argv[1]
groups = json.loads(sys.argv[2])

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",
    "backend": data["metadata"]["backend"],
    "engine": data["metadata"]["engine"],
    "server_mode": "w8a8_int8_baseline",
    "server_start_script": f"experiments/{data['metadata']['experiment']}/start_server.sh",
    "groups": groups
}
with open(path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
PY
}

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

# Cleanup any leftovers.
stop_server

write_backend_metadata

for group in "${CONTEXT_GROUPS[@]}"; do
  read -r group_label context_length max_running <<< "$group"

  # Resolve scenario array for this group.
  scenario_array_name="SCENARIOS_${group_label^^}"
  if [[ -z "${!scenario_array_name:-}" ]]; then
    log "WARNING: no scenarios defined for group ${group_label}"
    continue
  fi

  run_group "$group_label" "$context_length" "$max_running" "$scenario_array_name"
done

parse_backend

log "all results saved to ${RESULT_BASE}/${RUN_ID}"