#!/usr/bin/env bash
set -Eeuo pipefail

PYTHON=${PYTHON:-/root/.miniconda3/envs/sglang/bin/python}
SGLANG_BIN=${SGLANG_BIN:-/root/.miniconda3/envs/sglang/bin/sglang}
MODEL=${MODEL:-/data/hf_models/MiniMax-H3}
PROMPT_FILE=${PROMPT_FILE:-/root/.cache/sglang/vbench_subject_consistency.txt}
CLIENT_SCRIPT=${CLIENT_SCRIPT:-/data/wxy/sskj-h3/throughput/sglang-base/scripts/minimax_h3_t2va_bench.py}
BASE_PORT=${BASE_PORT:-35110}
MASTER_PORT_BASE=${MASTER_PORT_BASE:-36110}
SCHEDULER_PORT_BASE=${SCHEDULER_PORT_BASE:-37110}
RUN_ID=${RUN_ID:-t2va-base-tp2x4-allres-20steps-5s-$(date '+%Y%m%d-%H%M%S')}
RESULT_ROOT=${RESULT_ROOT:-/data/wxy/sskj-h3/throughput/sglang-base/results/$RUN_ID}
SERVER_START_TIMEOUT=${SERVER_START_TIMEOUT:-1800}

declare -a SERVER_PIDS=()
declare -a CLIENT_PIDS=()

log() { printf '[%s] %s\n' "$(date '+%F %T')" "$*"; }
die() { log "ERROR: $*" >&2; exit 1; }

port_is_open() {
  "$PYTHON" - "$1" <<'PY'
import socket, sys
s = socket.socket(); s.settimeout(0.5)
try: s.connect(("127.0.0.1", int(sys.argv[1])))
except OSError: raise SystemExit(1)
else: raise SystemExit(0)
finally: s.close()
PY
}

stop_servers() {
  local pid alive deadline
  ((${#SERVER_PIDS[@]})) || return 0
  log "stopping ${#SERVER_PIDS[@]} server(s)"
  for pid in "${SERVER_PIDS[@]}"; do kill -INT "$pid" 2>/dev/null || true; done
  deadline=$((SECONDS + 120))
  while ((SECONDS < deadline)); do
    alive=0
    for pid in "${SERVER_PIDS[@]}"; do kill -0 "$pid" 2>/dev/null && alive=1; done
    ((alive == 0)) && break
    sleep 2
  done
  for pid in "${SERVER_PIDS[@]}"; do
    if kill -0 "$pid" 2>/dev/null; then
      kill -TERM -- "-$pid" 2>/dev/null || kill -TERM "$pid" 2>/dev/null || true
      sleep 3
      kill -KILL -- "-$pid" 2>/dev/null || kill -KILL "$pid" 2>/dev/null || true
    fi
    wait "$pid" 2>/dev/null || true
  done
  SERVER_PIDS=()
}

cleanup() {
  local rc=$? pid
  trap - EXIT INT TERM
  for pid in "${CLIENT_PIDS[@]}"; do kill -TERM "$pid" 2>/dev/null || true; done
  stop_servers
  exit "$rc"
}
trap cleanup EXIT INT TERM

wait_healthy() {
  local port=$1 pid=$2 server_log=$3 deadline=$((SECONDS + SERVER_START_TIMEOUT))
  while ((SECONDS < deadline)); do
    curl -fsS --max-time 5 "http://127.0.0.1:${port}/health" >/dev/null 2>&1 && return 0
    if ! kill -0 "$pid" 2>/dev/null; then tail -100 "$server_log" >&2 || true; return 1; fi
    sleep 5
  done
  tail -100 "$server_log" >&2 || true
  return 1
}

[[ -x "$PYTHON" ]] || die "missing python: $PYTHON"
[[ -x "$SGLANG_BIN" ]] || die "missing sglang: $SGLANG_BIN"
[[ -f "$CLIENT_SCRIPT" ]] || die "missing client: $CLIENT_SCRIPT"
[[ -f "$PROMPT_FILE" ]] || die "missing prompts: $PROMPT_FILE"
mkdir -p "$RESULT_ROOT"

cat >"$RESULT_ROOT/run_config.env" <<EOF
method=base
task=t2va
tp_size=2
replicas=4
requests_per_resolution=4
resolutions=480,720,768,1080
total_requests=16
num_inference_steps=20
duration_seconds=5
aspect_ratio=16:9
conditions=[]
model=$MODEL
sglang_bin=$SGLANG_BIN
EOF
"$PYTHON" -m pip show sglang >"$RESULT_ROOT/sglang_pip_show.txt"
"$PYTHON" -m pip freeze >"$RESULT_ROOT/pip_freeze.txt"
sha256sum "$CLIENT_SCRIPT" "$0" >"$RESULT_ROOT/source_sha256.txt"
nvidia-smi >"$RESULT_ROOT/nvidia_smi_before.txt"

for replica in 0 1 2 3; do
  port=$((BASE_PORT + replica * 10))
  master_port=$((MASTER_PORT_BASE + replica * 10))
  scheduler_port=$((SCHEDULER_PORT_BASE + replica * 10))
  for candidate in "$port" "$((port + 1))" "$master_port" "$scheduler_port"; do
    port_is_open "$candidate" && die "port already in use: $candidate"
  done
  first_gpu=$((replica * 2)); gpu_csv="$first_gpu,$((first_gpu + 1))"
  server_dir="$RESULT_ROOT/server_${replica}_port${port}"
  mkdir -p "$server_dir/outputs"
  printf '%s\n' "$gpu_csv" >"$server_dir/cuda_visible_devices.txt"
  log "starting FL2VA partition for T2VA replica=$replica GPUs=$gpu_csv port=$port"
  env -u SGLANG_CACHE_DIT_ENABLED -u SGLANG_CACHE_DIT_FN -u SGLANG_CACHE_DIT_BN \
    -u SGLANG_CACHE_DIT_RDT -u SGLANG_CACHE_DIT_MC -u SGLANG_CACHE_DIT_WARMUP \
    -u SGLANG_CACHE_DIT_TAYLORSEER -u SGLANG_CACHE_DIT_SCM_PRESET \
    -u SGLANG_CACHE_DIT_SCM_POLICY CUDA_VISIBLE_DEVICES="$gpu_csv" \
    PATH="/root/.miniconda3/envs/sglang/bin:$PATH" PYTHONUNBUFFERED=1 \
    TOKENIZERS_PARALLELISM=false SGLANG_USE_RUNAI_MODEL_STREAMER=false \
    setsid "$SGLANG_BIN" serve --model-path "$MODEL" --model-variant FL2VA \
    --backend sglang --performance-mode speed --num-gpus 2 --tp-size 2 \
    --ulysses-degree 1 --use-fsdp-inference false --enable-torch-compile false \
    --batching-max-size 1 --batching-delay-ms 0 --host 0.0.0.0 --port "$port" \
    --master-port "$master_port" --scheduler-port "$scheduler_port" \
    --output-path "$server_dir/outputs" >"$server_dir/server.log" 2>&1 &
  SERVER_PIDS+=("$!")
done

for replica in 0 1 2 3; do
  port=$((BASE_PORT + replica * 10))
  log "waiting for replica=$replica port=$port"
  wait_healthy "$port" "${SERVER_PIDS[$replica]}" "$RESULT_ROOT/server_${replica}_port${port}/server.log" \
    || die "replica=$replica failed startup"
  log "replica=$replica healthy"
done

for replica in 0 1 2 3; do
  port=$((BASE_PORT + replica * 10))
  client_dir="$RESULT_ROOT/client_${replica}_port${port}"
  mkdir -p "$client_dir"
  "$PYTHON" "$CLIENT_SCRIPT" run --host 127.0.0.1 --port "$port" --model "$MODEL" \
    --replica-index "$replica" --num-replicas 4 --prompt-file "$PROMPT_FILE" \
    --resolutions 480,720,768,1080 --requests-per-resolution 4 \
    --num-inference-steps 20 --warmup-inference-steps 5 --duration-seconds 5 \
    --aspect-ratio 16:9 --flow-shift 12.0 --audio-flow-shift 3.0 --seed 1101 \
    --output "$client_dir/results.jsonl" >"$client_dir/client.log" 2>&1 &
  CLIENT_PIDS+=("$!")
  log "started client replica=$replica port=$port requests=4"
done

failed=0
for replica in 0 1 2 3; do wait "${CLIENT_PIDS[$replica]}" || failed=1; done
CLIENT_PIDS=()
"$PYTHON" "$CLIENT_SCRIPT" summarize --input-dir "$RESULT_ROOT" --expected-requests 16 \
  --output "$RESULT_ROOT/summary.json" >"$RESULT_ROOT/summary.log" || failed=1
nvidia-smi >"$RESULT_ROOT/nvidia_smi_after.txt"
stop_servers
((failed == 0)) || die "T2VA benchmark had failed requests"
log "T2VA benchmark complete: $RESULT_ROOT"
trap - EXIT INT TERM