#!/usr/bin/env bash
# TP×DP matrix benchmark for DeepSeek-V4-Flash-INT8 on Kunlun P800 SGLang (Docker).
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')}"
RESULT_BASE="${SCRIPT_DIR}/results"
MATRIX_FILE="${MATRIX_FILE:-${SCRIPT_DIR}/matrix.json}"
MATRIX_MODE="${MATRIX_MODE:-Y}"
SCENARIO_TIMEOUT_S="${SCENARIO_TIMEOUT_S:-1800}"
GPU_MEM_SAMPLE_INTERVAL_S="${GPU_MEM_SAMPLE_INTERVAL_S:-1}"
DRY_RUN="${DRY_RUN:-0}"
GRID_LIMIT="${GRID_LIMIT:-0}"

PYTHON="${PYTHON:-$(command -v python3)}"

# Export variables used inside functions that are called via bash -c subshells.
export CONTAINER_NAME CONTAINER_PYTHON MODEL_PATH SGLANG_PORT CONTAINER_DATASET_PATH

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

log "experiment=${EXPERIMENT_NAME} run_id=${RUN_ID} platform=${PLATFORM} hardware=${HARDWARE}"
log "matrix_mode=${MATRIX_MODE} matrix_file=${MATRIX_FILE} dry_run=${DRY_RUN} grid_limit=${GRID_LIMIT}"

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

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

container_running() {
  docker inspect "$CONTAINER_NAME" >/dev/null 2>&1
}

stop_server() {
  local tp="$1"
  local dp="$2"

  # The server only ever runs inside this experiment's dedicated container,
  # so removing the container is a complete and safe stop.
  if container_running; then
    log "stopping container ${CONTAINER_NAME} (tp=${tp}, dp=${dp})"
    docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true
  else
    docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true
  fi
  sleep 2
}

build_server_args() {
  # Keep in sync with the launch args in start_sglang_docker.sh.
  local tp="$1"
  local dp="$2"

  local args="--host 0.0.0.0 --port ${SGLANG_PORT} --model-path /models --attention-backend nsa --nsa-prefill klxdsa --nsa-decode klxdsa --trust-remote-code --disable-custom-all-reduce --page-size 64 --mem-fraction-static ${MEM_FRACTION_STATIC} --tensor-parallel-size ${tp} --disable-shared-experts-fusion --quantization w8a8_int8 --kv-cache-dtype float16 --disable-piecewise-cuda-graph --cuda-graph-max-bs 32 --watchdog-timeout 3000000 --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 --constrained-json-disable-any-whitespace --enable-metrics --enable-request-time-stats-logging --context-length ${CONTEXT_LENGTH}"
  if [[ "$dp" -gt 1 ]]; then
    args="${args} --dp-size ${dp}"
  fi
  printf '%s' "$args"
}

start_server() {
  local tp="$1"
  local dp="$2"

  log "starting sglang server tp=${tp} dp=${dp}"
  bash "${SCRIPT_DIR}/start_sglang_docker.sh" "$tp" "$dp" \
    >> "${log_dir_global}/sglang_tp${tp}_dp${dp}.server.outer.log" 2>&1

  if ! is_server_healthy; then
    log "error: sglang server tp=${tp} dp=${dp} failed health check on port ${SGLANG_PORT}"
    return 1
  fi
  log "sglang server tp=${tp} dp=${dp} is healthy on port ${SGLANG_PORT}"
}

restart_server() {
  local tp="$1"
  local dp="$2"
  log "restarting sglang server tp=${tp} dp=${dp} after non-OOM failure"
  stop_server "$tp" "$dp"
  sleep 10
  start_server "$tp" "$dp"
}

run_bench_serving() {
  # Run sglang.bench_serving inside the server container (the P800 image
  # carries the only working client environment for this fork).
  docker exec "$CONTAINER_NAME" \
    env HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_DATASETS_OFFLINE=1 \
    "$CONTAINER_PYTHON" -m sglang.bench_serving "$@"
}
export -f run_bench_serving

run_warmup() {
  local input_len="$1"
  local output_len="$2"

  log "warming up (input=${input_len}, output=${output_len}, num=1)"
  bash -c '
    run_bench_serving \
      --backend sglang \
      --host 127.0.0.1 \
      --port "'"$SGLANG_PORT"'" \
      --model "'"$MODEL_PATH"'" \
      --dataset-name random \
      --dataset-path "'"$CONTAINER_DATASET_PATH"'" \
      --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 /dev/null \
      --output-details \
      >> "'"${log_dir_global}/warmup.log"'" 2>&1
  ' || log "WARNING: warmup failed; continuing with config"
  log "warmup completed"
}

scenario_already_completed() {
  local output_file="$1"
  local expected="$2"
  [[ -s "$output_file" ]] || return 1
  local completed
  completed="$("$PYTHON" -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" ]]
}

scenario_already_processed() {
  local result_root="$1"
  local scenario_name="$2"
  local json_path="${result_root}/results.json"
  [[ -f "$json_path" ]] || return 1
  "$PYTHON" -c "
import json, sys
path, name = sys.argv[1], sys.argv[2]
try:
    with open(path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    for s in data.get('scenarios', []):
        if s.get('name') == name:
            if s.get('status') or s.get('metrics', {}).get('success', 0) > 0:
                sys.exit(0)
except Exception:
    pass
sys.exit(1)
" "$json_path" "$scenario_name"
}

detect_oom() {
  local detail_log="$1"
  local pattern='CUDA out of memory|torch\.OutOfMemoryError|OutOfMemory|out of memory|OOM|RESOURCE_EXHAUSTED|Failed to allocate memory'
  if grep -Eiq "$pattern" "$detail_log" 2>/dev/null; then
    return 0
  fi
  docker logs --tail 2000 "$CONTAINER_NAME" 2>/dev/null | grep -Eiq "$pattern"
}

start_gpu_monitor() {
  # Sample XPU memory/utilization with xpu-smi and emit nvidia-smi-style CSV
  # rows so scripts/common/parse_backend.py can read them unchanged.
  local csv_path="$1"
  mkdir -p "$(dirname "$csv_path")"
  (
    echo "timestamp, index, memory.used [MiB], memory.total [MiB], utilization.gpu [%]"
    while true; do
      xpu-smi 2>/dev/null | awk -F'|' -v ts="$(date '+%Y/%m/%d %H:%M:%S')" '
        /^\|[[:space:]]*[0-9]+[[:space:]]/ {
          l=$0; sub(/^\|[[:space:]]*/, "", l); sub(/[[:space:]].*$/, "", l); idx=l
        }
        /MiB/ && $4 ~ /%/ {
          split($3, a, "/"); u=a[1]; t=a[2]
          gsub(/[^0-9.]/, "", u); gsub(/[^0-9.]/, "", t)
          r=$4; gsub(/[^0-9.]/, "", r)
          printf "%s, %s, %s MiB, %s MiB, %s %%\n", ts, idx, u, t, r
        }'
      sleep "$GPU_MEM_SAMPLE_INTERVAL_S"
    done
  ) >> "$csv_path" 2>/dev/null &
  echo $!
}

stop_gpu_monitor() {
  local pid="$1"
  if kill -0 "$pid" 2>/dev/null; then
    kill "$pid" 2>/dev/null || true
    sleep 1
    kill -9 "$pid" 2>/dev/null || true
  fi
}

scenario_config_json() {
  "$PYTHON" -c "
import json, sys
print(json.dumps({
    'phase': 'main',
    'concurrency': int(sys.argv[1]),
    'input_len': int(sys.argv[2]),
    'output_len': int(sys.argv[3]),
    'dataset': 'random',
    'num_prompts': int(sys.argv[4]),
}))
" "$1" "$2" "$3" "$4"
}

append_scenario_record() {
  local result_root="$1"
  local json_path="$result_root/results.json"
  shift
  local scenario_json
  scenario_json="$("$PYTHON" -c "
import json, sys
pairs = [a.split('=', 1) for a in sys.argv[1:]]
d = {}
for k, v in pairs:
    try:
        d[k] = json.loads(v)
    except json.JSONDecodeError:
        d[k] = v
print(json.dumps(d, ensure_ascii=False))
" "$@")"
  PYTHON="$PYTHON" append_scenario_to_json "$json_path" "$scenario_json"
}

record_skipped_csv() {
  local csv_path="$1"
  shift
  # Args: key=value
  local row
  row="$("$PYTHON" -c "
import csv, json, sys, io
pairs = [a.split('=', 1) for a in sys.argv[1:]]
d = {}
for k, v in pairs:
    try:
        d[k] = json.loads(v)
    except json.JSONDecodeError:
        d[k] = v
buf = io.StringIO()
writer = csv.DictWriter(buf, fieldnames=['engine','tp','dp','mark','isl','osl','concurrency','status','reason','detail_log'], extrasaction='ignore')
writer.writerow(d)
print(buf.getvalue().strip())
" "$@")"
  echo "$row" >> "$csv_path"
}

skip_remaining_scenarios() {
  local result_root="$1"
  local scenario_tsv="$2"
  local start_index="$3"
  local status="$4"
  local reason="$5"
  local tp="$6"
  local dp="$7"
  local skipped_csv="${RESULT_BASE}/${RUN_ID}/skipped_after_oom.csv"

  local i=0
  tail -n +2 "$scenario_tsv" | while IFS=$'\t' read -r mark isl osl conc num; do
    if (( i < start_index )); then
      i=$((i + 1))
      continue
    fi
    i=$((i + 1))
    local sname="c${conc}_i${isl}_o${osl}"
    if scenario_already_processed "$result_root" "$sname"; then
      continue
    fi
    append_scenario_record "$result_root" \
      "name=${sname}" \
      "config=$(scenario_config_json "$conc" "$isl" "$osl" "$num")" \
      "status=\"${status}\"" \
      "note=\"${reason}\""
    record_skipped_csv "$skipped_csv" \
      "engine=sglang" "tp=${tp}" "dp=${dp}" "mark=${mark}" "isl=${isl}" "osl=${osl}" "concurrency=${conc}" "status=${status}" "reason=${reason}"
  done
}

# ---------------------------------------------------------------------------
# Per-configuration runner
# ---------------------------------------------------------------------------

run_parallel_config() {
  local tp="$1"
  local dp="$2"
  local config_label="tp${tp}_dp${dp}"
  local result_root="${RESULT_BASE}/${RUN_ID}/${config_label}"
  local raw_dir="${result_root}/raw_outputs"
  local gpu_log_dir="${result_root}/gpu_logs"
  local phase_log_dir="${result_root}/logs"
  mkdir -p "$raw_dir" "$gpu_log_dir" "$phase_log_dir"

  log "===== ${config_label} START ====="

  # Generate scenario list for this config.
  local scenario_tsv="${result_root}/scenarios.tsv"
  "$PYTHON" "${SCRIPT_DIR}/generate_scenarios.py" \
    --matrix "$MATRIX_FILE" \
    --mode "$MATRIX_MODE" \
    > "$scenario_tsv"
  local total_scenarios
  total_scenarios="$(tail -n +2 "$scenario_tsv" | wc -l)"
  log "generated ${total_scenarios} scenarios for ${config_label}"

  # Write metadata.
  ensure_result_root "$result_root"
  write_metadata_json \
    "${result_root}/results.json" \
    "${EXPERIMENT_NAME}_${config_label}" \
    "$RUN_ID" \
    "$MODEL_PATH" \
    "sglang" \
    "sglang-xpu" \
    "$HARDWARE" \
    "$ACCELERATOR" \
    "$CHIP" \
    "experiments/p800/${EXPERIMENT_NAME}/run_bench.sh" \
    "$DOCKER_IMAGE" \
    "P800 SGLang TP×DP matrix for DeepSeek-V4-Flash-INT8"

  local server_args_str
  server_args_str="$(build_server_args "$tp" "$dp")"
  "$PYTHON" - "$tp" "$dp" "$server_args_str" "${result_root}/results.json" <<'PY'
import json, os, sys

tp, dp, server_args, path = int(sys.argv[1]), int(sys.argv[2]), sys.argv[3], sys.argv[4]
with open(path, "r", encoding="utf-8") as f:
    data = json.load(f)
data["config"] = {
    "tp": tp,
    "dp": dp,
    "xpu_visible_devices": os.environ.get("XPU_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7"),
    "backend": "sglang",
    "engine": "sglang-xpu",
    "server_start_script": "experiments/p800/dsv4_p800_sglang_tp_dp_matrix/start_sglang_docker.sh",
    "server_args": server_args,
}
with open(path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
PY

  if [[ "$DRY_RUN" == "1" ]]; then
    log "DRY_RUN: would start server with args: ${server_args_str}"
    local line
    tail -n +2 "$scenario_tsv" | while IFS=$'\t' read -r mark isl osl conc num; do
      log "DRY_RUN: ${config_label} scenario mark=${mark} c=${conc} i=${isl} o=${osl} n=${num}"
    done
    log "===== ${config_label} DONE (dry run) ====="
    return 0
  fi

  # Initialize skipped_after_oom.csv for this run.
  local skipped_csv="${RESULT_BASE}/${RUN_ID}/skipped_after_oom.csv"
  if [[ ! -f "$skipped_csv" ]]; then
    echo "engine,tp,dp,mark,isl,osl,concurrency,status,reason,detail_log" > "$skipped_csv"
  fi

  # Start server once for this TP×DP config.
  if ! start_server "$tp" "$dp"; then
    log "ERROR: ${config_label} failed to start; skipping all scenarios"
    skip_remaining_scenarios "$result_root" "$scenario_tsv" 0 "SKIPPED_SERVICE_START_FAILED" "service failed to start" "$tp" "$dp"
    log "===== ${config_label} DONE ====="
    return 0
  fi

  # Warmup with a small prompt before the first scenario.
  run_warmup 1024 128 || true

  # Read scenarios into an array so we can skip remaining entries on failure.
  local -a scenarios=()
  while IFS= read -r line; do
    scenarios+=("$line")
  done < <(tail -n +2 "$scenario_tsv")

  local i mark isl osl conc num
  local output_file container_output detail_log gpu_csv sname bench_rc
  for (( i = 0; i < ${#scenarios[@]}; i++ )); do
    IFS=$'\t' read -r mark isl osl conc num <<< "${scenarios[$i]}"

    if [[ "$GRID_LIMIT" -gt 0 && "$i" -ge "$GRID_LIMIT" ]]; then
      log "GRID_LIMIT=${GRID_LIMIT} reached; skipping remaining scenarios"
      skip_remaining_scenarios "$result_root" "$scenario_tsv" "$i" "SKIPPED_GRID_LIMIT" "GRID_LIMIT reached" "$tp" "$dp"
      break
    fi

    sname="c${conc}_i${isl}_o${osl}"
    output_file="${raw_dir}/sglang_main_${conc}_${isl}_${osl}.jsonl"
    container_output="/tmp/bench_outputs/sglang_main_${conc}_${isl}_${osl}.jsonl"
    detail_log="${phase_log_dir}/sglang_${config_label}_${sname}.log"
    gpu_csv="${gpu_log_dir}/gpu_mem_${conc}_${isl}_${osl}.csv"

    if scenario_already_completed "$output_file" "$num" || scenario_already_processed "$result_root" "$sname"; then
      log "skipping already-processed ${config_label} scenario: ${sname}"
      continue
    fi

    if ! container_running; then
      log "WARNING: container ${CONTAINER_NAME} died, skipping remaining scenarios in ${config_label}"
      skip_remaining_scenarios "$result_root" "$scenario_tsv" "$i" "SKIPPED_CONTAINER_DIED" "container died" "$tp" "$dp"
      break
    fi

    log "running ${config_label} scenario: mark=${mark} c=${conc} i=${isl} o=${osl} n=${num}"

    local gpu_pid
    gpu_pid="$(start_gpu_monitor "$gpu_csv")"

    bench_rc=0
    docker exec "$CONTAINER_NAME" mkdir -p /tmp/bench_outputs
    timeout "$SCENARIO_TIMEOUT_S" bash -c '
      run_bench_serving \
        --backend sglang \
        --host 127.0.0.1 \
        --port "'"$SGLANG_PORT"'" \
        --model "'"$MODEL_PATH"'" \
        --dataset-name random \
        --dataset-path "'"$CONTAINER_DATASET_PATH"'" \
        --random-input-len "'"$isl"'" \
        --random-output-len "'"$osl"'" \
        --random-range-ratio 1.0 \
        --num-prompts "'"$num"'" \
        --max-concurrency "'"$conc"'" \
        --request-rate 10000 \
        --output-file "'"$container_output"'" \
        --output-details \
        > "'"$detail_log"'" 2>&1
    ' || bench_rc=$?

    stop_gpu_monitor "$gpu_pid"

    if [[ "$bench_rc" -eq 0 ]]; then
      if docker cp "${CONTAINER_NAME}:${container_output}" "$output_file"; then
        log "finished ${config_label} scenario: output=${output_file}"
        append_scenario_record "$result_root" \
          "name=${sname}" \
          "config=$(scenario_config_json "$conc" "$isl" "$osl" "$num")" \
          "status=\"completed\"" \
          "note=\"benchmark finished successfully\""
        continue
      fi
      log "ERROR: failed to copy ${container_output} from container"
      bench_rc=1
    fi

    # Failure handling.
    if detect_oom "$detail_log"; then
      log "ERROR: ${config_label} scenario ${sname} triggered OOM; stopping config"
      append_scenario_record "$result_root" \
        "name=${sname}" \
        "config=$(scenario_config_json "$conc" "$isl" "$osl" "$num")" \
        "status=\"OOM\"" \
        "note=\"detected out-of-memory\""
      record_skipped_csv "$skipped_csv" \
        "engine=sglang" "tp=${tp}" "dp=${dp}" "mark=${mark}" "isl=${isl}" "osl=${osl}" "concurrency=${conc}" "status=OOM" "reason=detected out-of-memory" "detail_log=${detail_log}"
      stop_server "$tp" "$dp"
      skip_remaining_scenarios "$result_root" "$scenario_tsv" "$((i + 1))" "SKIPPED_AFTER_OOM" "previous case OOM" "$tp" "$dp"
      break
    fi

    log "ERROR: ${config_label} scenario ${sname} failed (rc=${bench_rc}); see ${detail_log}"
    if [[ "$mark" == "P" ]]; then
      log "optional (P) scenario failed; recording as skipped and continuing"
      append_scenario_record "$result_root" \
        "name=${sname}" \
        "config=$(scenario_config_json "$conc" "$isl" "$osl" "$num")" \
        "status=\"skipped_optional\"" \
        "note=\"optional scenario failed (rc=${bench_rc})\""
      record_skipped_csv "$skipped_csv" \
        "engine=sglang" "tp=${tp}" "dp=${dp}" "mark=${mark}" "isl=${isl}" "osl=${osl}" "concurrency=${conc}" "status=skipped_optional" "reason=optional scenario failed (rc=${bench_rc})" "detail_log=${detail_log}"
      continue
    fi

    # Mandatory scenario failed but not OOM: try to restart the server.
    if restart_server "$tp" "$dp"; then
      run_warmup 1024 128 || true
      log "resuming ${config_label} after server restart"
      continue
    fi

    log "ERROR: ${config_label} server restart failed; skipping remaining scenarios"
    append_scenario_record "$result_root" \
      "name=${sname}" \
      "config=$(scenario_config_json "$conc" "$isl" "$osl" "$num")" \
      "status=\"FAILED\"" \
      "note=\"scenario failed and server restart failed (rc=${bench_rc})\""
    record_skipped_csv "$skipped_csv" \
      "engine=sglang" "tp=${tp}" "dp=${dp}" "mark=${mark}" "isl=${isl}" "osl=${osl}" "concurrency=${conc}" "status=FAILED" "reason=scenario failed and server restart failed" "detail_log=${detail_log}"
    skip_remaining_scenarios "$result_root" "$scenario_tsv" "$((i + 1))" "SKIPPED_RESTART_FAILED" "server restart failed" "$tp" "$dp"
    break
  done

  stop_server "$tp" "$dp"

  # Parse results.
  log "parsing ${config_label} results"
  "$PYTHON" "${SCRIPT_DIR}/../../../scripts/common/parse_backend.py" "$result_root" --backend sglang \
    >> "${phase_log_dir}/parse.log" 2>&1 || {
      log "WARNING: parser failed for ${config_label}; see ${phase_log_dir}/parse.log"
    }

  log "===== ${config_label} DONE ====="
}

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

# Cleanup any leftovers.
for cfg in "${PARALLEL_CONFIGS[@]}"; do
  read -r tp dp <<< "$cfg"
  stop_server "$tp" "$dp"
done

# Run each parallel configuration.
for cfg in "${PARALLEL_CONFIGS[@]}"; do
  read -r tp dp <<< "$cfg"
  run_parallel_config "$tp" "$dp"
done

# Generate cross-configuration comparison.
log "generating comparison report"
"$PYTHON" "${SCRIPT_DIR}/compare.py" \
  --run-root "${RESULT_BASE}/${RUN_ID}" \
  --output "${RESULT_BASE}/${RUN_ID}/comparison.md" \
  >> "${log_dir_global}/compare.log" 2>&1 || {
    log "WARNING: comparison script failed; see ${log_dir_global}/compare.log"
  }

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