#!/usr/bin/env bash
# Long-context matrix for P800 SGLang INT8.
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"

RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}"
RUN_DATE="$(date '+%m%d')"
RESULT_BASE="${SCRIPT_DIR}/results"
RESULT_ROOT="${RESULT_BASE}/${RUN_ID}"
LOG_DIR="${RESULT_ROOT}/logs"
RAW_DIR="${RESULT_ROOT}/raw_outputs"
CONTAINER_PYTHON="${CONTAINER_PYTHON:-/root/miniconda/envs/python310_torch25_cuda/bin/python}"

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

log "experiment=${EXPERIMENT_NAME}"
log "run_id=${RUN_ID}"
log "platform=${PLATFORM}"
log "hardware=${HARDWARE}"
log "model=${MODEL_PATH}"
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
}

stop_server() {
  log "stopping container ${CONTAINER_NAME}"
  docker rm -f "$CONTAINER_NAME" 2>/dev/null || true
  pkill -9 -f 'sglang.launch_server' 2>/dev/null || true
  sleep 2
}

start_server() {
  local max_context_len="$1"
  local max_running="$2"
  log "starting server (context-length=${max_context_len}, max-running=${max_running})"
  bash "${SCRIPT_DIR}/start_sglang.sh" "$max_context_len" "$max_running" >> "${LOG_DIR}/start_server_${max_context_len}.log" 2>&1
}

run_warmup() {
  local input_len="$1"
  local output_len="${2:-256}"
  local container_output="/tmp/bench_outputs/sglang_warmup_${input_len}_${output_len}.jsonl"
  local detail_log="${LOG_DIR}/warmup_${input_len}_${output_len}.log"

  log "warming up (input=${input_len}, output=${output_len}, num=1)"
  docker exec "$CONTAINER_NAME" mkdir -p "$(dirname "$container_output")"

  docker exec "$CONTAINER_NAME" \
    env HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_DATASETS_OFFLINE=1 \
    "${CONTAINER_PYTHON}" -m sglang.bench_serving \
      --backend sglang \
      --host 127.0.0.1 \
      --port "$PORT" \
      --model "$MODEL_PATH" \
      --dataset-name random \
      --dataset-path /workspace/dummy_sharegpt.json \
      --random-input-len "$input_len" \
      --random-output-len "$output_len" \
      --random-range-ratio 1.0 \
      --num-prompts 1 \
      --max-concurrency 1 \
      --request-rate 10000 \
      --output-file "$container_output" \
      --output-details \
      > "$detail_log" 2>&1 || {
        log "WARNING: warmup failed; see ${detail_log}"
        return 1
      }
  log "warmup completed"
}

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 max_context_len="$2"
  local max_running="$3"
  local -n scenarios_ref="$4"

  mkdir -p "$RAW_DIR" "$LOG_DIR"

  log "===== ${group_label} START (context-length=${max_context_len}, max-running=${max_running}) ====="

  stop_server
  if ! start_server "$max_context_len" "$max_running"; then
    log "ERROR: failed to start server 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 "$warmup_input_len" 256

  for scenario in "${scenarios_ref[@]}"; do
    read -r concurrency input_len output_len num_prompts <<< "$scenario"

    local host_output_file="${RAW_DIR}/sglang_${group_label}_${RUN_DATE}_${concurrency}_${input_len}_${output_len}.jsonl"
    local container_output="/tmp/bench_outputs/sglang_${group_label}_${RUN_DATE}_${concurrency}_${input_len}_${output_len}.jsonl"
    local detail_log="${LOG_DIR}/sglang_${group_label}_c${concurrency}_i${input_len}_o${output_len}.log"

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

    log "running scenario: group=${group_label} c=${concurrency} i=${input_len} o=${output_len} n=${num_prompts}"

    docker exec "$CONTAINER_NAME" mkdir -p "$(dirname "$container_output")"

    local bench_rc=0
    docker exec "$CONTAINER_NAME" \
      env HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_DATASETS_OFFLINE=1 \
      "${CONTAINER_PYTHON}" -m sglang.bench_serving \
        --backend sglang \
        --host 127.0.0.1 \
        --port "$PORT" \
        --model "$MODEL_PATH" \
        --dataset-name random \
        --dataset-path /workspace/dummy_sharegpt.json \
        --random-input-len "$input_len" \
        --random-output-len "$output_len" \
        --random-range-ratio 1.0 \
        --num-prompts "$num_prompts" \
        --max-concurrency "$concurrency" \
        --request-rate 10000 \
        --output-file "$container_output" \
        --output-details \
        > "$detail_log" 2>&1 || bench_rc=$?

    if [[ "$bench_rc" -ne 0 ]]; then
      log "ERROR: scenario group=${group_label} c=${concurrency} i=${input_len} o=${output_len} failed; see ${detail_log}"
      continue
    fi

    docker cp "${CONTAINER_NAME}:${container_output}" "$host_output_file" || {
      log "ERROR: failed to copy output from container for scenario group=${group_label} c=${concurrency} i=${input_len} o=${output_len}"
      continue
    }

    # Verify the request actually completed.
    local completed
    completed="$(python3 -c "
import json
with open('${host_output_file}') as f:
    for line in f:
        data = json.loads(line)
        print(data.get('completed', 0))
        break
")"
    if [[ "${completed:-0}" -lt "$num_prompts" ]]; then
      log "ERROR: scenario group=${group_label} c=${concurrency} i=${input_len} o=${output_len} only completed ${completed}/${num_prompts}"
      continue
    fi

    log "finished scenario: output=${host_output_file}"
  done

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

parse_backend() {
  log "parsing results in ${RESULT_ROOT}"
  python3 "${SCRIPT_DIR}/../../scripts/common/parse_backend.py" "$RESULT_ROOT" --backend sglang \
    >> "${LOG_DIR}/parse.log" 2>&1 || {
      log "WARNING: parse_backend.py failed; see ${LOG_DIR}/parse.log"
    }
}

write_metadata() {
  ensure_result_root "$RESULT_ROOT"
  local meta_json="${RESULT_ROOT}/results.json"

  write_metadata_json \
    "$meta_json" \
    "$EXPERIMENT_NAME" \
    "$RUN_ID" \
    "$MODEL_PATH" \
    "sglang" \
    "sglang-xpu" \
    "$HARDWARE" \
    "$ACCELERATOR" \
    "$CHIP" \
    "experiments/${EXPERIMENT_NAME}/run_bench.sh" \
    "" \
    "P800 long-context matrix for DeepSeek-V4-Flash-INT8"

  # Embed config.
  python3 - "$meta_json" <<'PY'
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,
    "ep": 8,
    "xpu_visible_devices": "0,1,2,3,4,5,6,7",
    "platform": "kunlun_p800",
    "container_name": "sglang-dsv4-flash",
    "context_pad": 1024,
    "groups": [
        {"label": "64k", "max_context_len": 80000, "max_running": 10},
        {"label": "128k", "max_context_len": 140000, "max_running": 5}
    ]
}
with open(path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
PY
}

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

stop_server
write_metadata

for group in "${CONTEXT_GROUPS[@]}"; do
  read -r group_label max_context_len 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" "$max_context_len" "$max_running" "$scenario_array_name"
done

parse_backend

log "all results saved to ${RESULT_ROOT}"