#!/usr/bin/env bash
# Max context length exploration 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')}"
RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}"
LOG_DIR="${RESULT_ROOT}/logs"
RAW_DIR="${RESULT_ROOT}/raw_outputs"
CONTAINER_NAME="${CONTAINER_NAME:-sglang-dsv4-flash}"
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 "result_root=${RESULT_ROOT}"
log "platform=${PLATFORM}"
log "hardware=${HARDWARE}"
log "model=${MODEL_PATH}"
log "candidate_lengths=${CANDIDATE_LENS[*]}"

# ---------------------------------------------------------------------------
# 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 target_len="$1"
  log "starting server (target_len=${target_len})"
  bash "${SCRIPT_DIR}/start_server.sh" "$target_len" >> "${LOG_DIR}/start_server_${target_len}.log" 2>&1
}

build_server_args() {
  local target_len="$1"
  local max_context_len=$(( target_len + CONTEXT_PAD ))
  echo "python -m sglang.launch_server --host 0.0.0.0 --port ${PORT} --model-path /models --attention-backend nsa --nsa-prefill klxdsa --nsa-decode klxdsa --trust-remote-code --disable-custom-all-reduce --chunked-prefill-size 8192 --page-size 64 --mem-fraction-static 0.8 --max-prefill-tokens ${max_context_len} --max-running-requests 64 --tensor-parallel-size 8 --ep-size 8 --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 --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --constrained-json-disable-any-whitespace --enable-metrics --enable-request-time-stats-logging --context-length ${max_context_len}"
}

append_scenario() {
  local scenario_json="$1"
  python3 - "$RESULT_JSON" "$scenario_json" <<'PY'
import json
import sys

results_path, scenario_json = sys.argv[1], sys.argv[2]
with open(results_path, "r", encoding="utf-8") as f:
    data = json.load(f)

data["scenarios"].append(json.loads(scenario_json))
with open(results_path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
PY
}

# Build a scenario JSON blob safely with Python.
build_scenario_json() {
  local target_len="$1"
  local max_context_len="$2"
  local success="$3"
  local server_args="$4"
  local error_msg="$5"
  local metrics_json="${6:-null}"

  python3 - "$target_len" "$max_context_len" "$success" "$server_args" "$error_msg" "$metrics_json" <<'PY'
import json
import sys

target_len, max_context_len, success, server_args, error_msg, metrics_json = sys.argv[1:7]
metrics = json.loads(metrics_json) if metrics_json not in ("null", "") else None
print(json.dumps({
    "backend": "sglang",
    "target_len": int(target_len),
    "max_context_len": int(max_context_len),
    "success": success == "true",
    "server_args": server_args,
    "metrics": metrics,
    "error": error_msg or None,
}, ensure_ascii=False))
PY
}

run_length() {
  local target_len="$1"
  local max_context_len=$(( target_len + CONTEXT_PAD ))
  local output_file="${RAW_DIR}/sglang_ctx_${target_len}.jsonl"
  local container_output="/tmp/bench_outputs/sglang_ctx_${target_len}.jsonl"
  local detail_log="${LOG_DIR}/sglang_ctx_${target_len}.log"

  log "trying sglang target_len=${target_len} (context-length=${max_context_len})"

  stop_server
  local server_args
  server_args="$(build_server_args "$target_len")"
  if ! start_server "$target_len"; then
    local err="server failed to start"
    log "ERROR: ${err}"
    append_scenario "$(build_scenario_json "$target_len" "$max_context_len" "false" "$server_args" "$err")"
    return 1
  fi

  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 "$target_len" \
      --random-output-len "$OUTPUT_LEN" \
      --random-range-ratio 1.0 \
      --num-prompts "$NUM_PROMPTS" \
      --max-concurrency "$MAX_CONCURRENCY" \
      --request-rate "$REQUEST_RATE" \
      --output-file "$container_output" \
      --output-details \
      > "$detail_log" 2>&1 || bench_rc=$?

  if [[ "$bench_rc" -ne 0 ]]; then
    local err="bench_serving exited with code ${bench_rc}; see ${detail_log}"
    log "ERROR: sglang target_len=${target_len} failed; ${err}"
    append_scenario "$(build_scenario_json "$target_len" "$max_context_len" "false" "$server_args" "$err")"
    stop_server
    return 1
  fi

  docker cp "${CONTAINER_NAME}:${container_output}" "$output_file" || {
    local err="failed to copy output from container"
    log "ERROR: ${err}"
    append_scenario "$(build_scenario_json "$target_len" "$max_context_len" "false" "$server_args" "$err")"
    stop_server
    return 1
  }

  # Verify the request actually completed.
  local completed
  completed="$(python3 -c "
import json
with open('${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
    local err="only ${completed}/${NUM_PROMPTS} requests completed"
    log "ERROR: sglang target_len=${target_len} ${err}"
    append_scenario "$(build_scenario_json "$target_len" "$max_context_len" "false" "$server_args" "$err")"
    stop_server
    return 1
  fi

  local metrics_json
  metrics_json="$(python3 "${SCRIPT_DIR}/extract_metrics.py" "$output_file")"

  append_scenario "$(build_scenario_json "$target_len" "$max_context_len" "true" "$server_args" "" "$metrics_json")"

  log "sglang target_len=${target_len} succeeded"
  stop_server
}

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

mkdir -p "$RAW_DIR" "$LOG_DIR"

RESULT_JSON="${RESULT_ROOT}/results.json"
write_metadata_json \
  "$RESULT_JSON" \
  "$EXPERIMENT_NAME" \
  "$RUN_ID" \
  "$MODEL_PATH" \
  "sglang" \
  "sglang-xpu" \
  "$HARDWARE" \
  "$ACCELERATOR" \
  "$CHIP" \
  "experiments/${EXPERIMENT_NAME}/run_bench.sh" \
  "" \
  "P800 max context length exploration for DeepSeek-V4-Flash-INT8"

# Embed config.
python3 - "$RESULT_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",
    "output_len": 1,
    "num_prompts": 1,
    "max_concurrency": 1,
    "request_rate": 10000,
    "context_pad": 1024,
}
with open(path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
PY

stop_server

for target_len in "${CANDIDATE_LENS[@]}"; do
  if ! run_length "$target_len"; then
    log "stopping exploration after first failure at target_len=${target_len}"
    break
  fi
done

log "parsing results"
python3 "${SCRIPT_DIR}/parse_results.py" "$RESULT_ROOT" >> "${LOG_DIR}/parse.log" 2>&1 || {
  log "WARNING: parse_results.py failed; see ${LOG_DIR}/parse.log"
}

log "all results saved to ${RESULT_ROOT}"