#!/usr/bin/env bash
# 64 requests per TP: 2 tasks x 4 short-edge resolutions x 8 prompts.
set -Eeuo pipefail

TOTAL_GPUS=${TOTAL_GPUS:-8}
NUM_INFERENCE_STEPS=${NUM_INFERENCE_STEPS:-20}
DURATION_SECONDS=${DURATION_SECONDS:-5}
TP_MATRIX=${TP_MATRIX:-"8 4 2"}
TASKS=${TASKS:-"fl2va ref2va"}
RESOLUTIONS=${RESOLUTIONS:-"480,720,768,1080"}
REQUESTS_PER_RESOLUTION=${REQUESTS_PER_RESOLUTION:-8}
REQUESTS_PER_TASK=$((REQUESTS_PER_RESOLUTION * 4))
TOTAL_REQUESTS_PER_TP=$((REQUESTS_PER_TASK * 2))

BASE_PORT=${BASE_PORT:-30010}
PORT_STRIDE=${PORT_STRIDE:-10}
MASTER_PORT_BASE=${MASTER_PORT_BASE:-31000}
SCHEDULER_PORT_BASE=${SCHEDULER_PORT_BASE:-32000}
HOST=${HOST:-127.0.0.1}
MODEL=${MODEL:-/data/hf_models/MiniMax-H3}
REPO_ROOT=${REPO_ROOT:-/data/wxy/sskj-h3}
REFERENCE_IMAGE=${REFERENCE_IMAGE:-/data/wxy/sskj-MiniMax-H3/assets/reference_images/landscape_mountain_lake.jpg}
PROMPT_FILE=${PROMPT_FILE:-/root/.cache/sglang/vbench_subject_consistency.txt}
PYTHON=${PYTHON:-/root/.miniconda3/envs/sglang/bin/python}
SGLANG_BIN=${SGLANG_BIN:-/root/.miniconda3/envs/sglang/bin/sglang}
CLIENT_SCRIPT=${CLIENT_SCRIPT:-$REPO_ROOT/throughput/sglang-base/scripts/minimax_h3_mixed_bench.py}
SSIM_SCRIPT=${SSIM_SCRIPT:-$REPO_ROOT/throughput/common/paired_video_ssim.py}
SSIM_REFERENCE_ROOT=${SSIM_REFERENCE_ROOT:-}
SSIM_THRESHOLD=${SSIM_THRESHOLD:-0.90}
SSIM_FAIL_BELOW_THRESHOLD=${SSIM_FAIL_BELOW_THRESHOLD:-false}
FFMPEG_BIN=${FFMPEG_BIN:-/root/.miniconda3/envs/deploy/bin/ffmpeg}
FFPROBE_BIN=${FFPROBE_BIN:-/root/.miniconda3/envs/deploy/bin/ffprobe}
SERVER_START_TIMEOUT=${SERVER_START_TIMEOUT:-1800}
RUN_ID=${RUN_ID:-mixed64-$(date '+%Y%m%d-%H%M%S')}
RESULT_ROOT=${RESULT_ROOT:-$REPO_ROOT/throughput/sglang-base/results/$RUN_ID}

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

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

[[ -x "$PYTHON" ]] || die "python not executable: $PYTHON"
[[ -x "$SGLANG_BIN" ]] || die "sglang not executable: $SGLANG_BIN"
[[ -f "$CLIENT_SCRIPT" ]] || die "client script missing: $CLIENT_SCRIPT"
[[ -f "$REFERENCE_IMAGE" ]] || die "reference image missing: $REFERENCE_IMAGE"
if [[ -n "$SSIM_REFERENCE_ROOT" ]]; then
  [[ -f "$SSIM_SCRIPT" ]] || die "SSIM script missing: $SSIM_SCRIPT"
  [[ -x "$FFMPEG_BIN" ]] || die "ffmpeg not executable: $FFMPEG_BIN"
  [[ -x "$FFPROBE_BIN" ]] || die "ffprobe not executable: $FFPROBE_BIN"
fi
mkdir -p "$RESULT_ROOT"
SUMMARY_TSV="$RESULT_ROOT/summary.tsv"
printf 'tp\treplicas\ttask\texpected\trecorded\tcompleted\tfailed\tmachine_qps\tlatency_mean_s\tlatency_p95_s\tmachine_wall_s\n' > "$SUMMARY_TSV"

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

stop_servers() {
  local pid alive deadline
  ((${#SERVER_PIDS[@]})) || return 0
  log "gracefully 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
      log "server pid=$pid did not exit after SIGINT; terminating process group"
      kill -TERM -- "-$pid" 2>/dev/null || kill -TERM "$pid" 2>/dev/null || true
      sleep 5
      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 log_file=$3 deadline=$((SECONDS + SERVER_START_TIMEOUT))
  while ((SECONDS < deadline)); do
    curl -fsS --max-time 5 "http://${HOST}:${port}/health" >/dev/null 2>&1 && return 0
    if ! kill -0 "$pid" 2>/dev/null; then tail -100 "$log_file" >&2 || true; return 1; fi
    sleep 5
  done
  tail -100 "$log_file" >&2 || true
  return 1
}

start_servers() {
  local tp=$1 replicas=$2 variant=$3 phase_dir=$4
  local replica port master_port scheduler_port first_gpu gpu_csv offset gpu server_dir server_log candidate
  SERVER_PIDS=()
  for ((replica=0; replica<replicas; replica++)); do
    port=$((BASE_PORT + replica * PORT_STRIDE))
    master_port=$((MASTER_PORT_BASE + replica * PORT_STRIDE))
    scheduler_port=$((SCHEDULER_PORT_BASE + replica * PORT_STRIDE))
    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 * tp)); gpu_csv=""
    for ((offset=0; offset<tp; offset++)); do
      gpu=$((first_gpu + offset)); [[ -z "$gpu_csv" ]] && gpu_csv="$gpu" || gpu_csv+=",$gpu"
    done
    server_dir="$phase_dir/server_${replica}_port${port}"; mkdir -p "$server_dir/outputs"
    server_log="$server_dir/server.log"; printf '%s\n' "$gpu_csv" > "$server_dir/cuda_visible_devices.txt"
    log "starting variant=$variant replica=$replica GPUs=$gpu_csv port=$port"
    CUDA_VISIBLE_DEVICES="$gpu_csv" PYTHONUNBUFFERED=1 TOKENIZERS_PARALLELISM=false \
      SGLANG_USE_RUNAI_MODEL_STREAMER=false setsid "$SGLANG_BIN" serve \
      --model-path "$MODEL" --model-variant "$variant" --backend sglang --performance-mode speed \
      --num-gpus "$tp" --tp-size "$tp" --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_log" 2>&1 &
    SERVER_PIDS+=("$!")
  done
  for ((replica=0; replica<replicas; replica++)); do
    port=$((BASE_PORT + replica * PORT_STRIDE))
    server_log="$phase_dir/server_${replica}_port${port}/server.log"
    wait_healthy "$port" "${SERVER_PIDS[$replica]}" "$server_log" || die "$variant replica=$replica failed startup"
    log "variant=$variant replica=$replica healthy port=$port"
  done
}

run_clients() {
  local tp=$1 replicas=$2 task=$3 phase_dir=$4
  local replica port client_dir failed=0
  CLIENT_PIDS=()
  for ((replica=0; replica<replicas; replica++)); do
    port=$((BASE_PORT + replica * PORT_STRIDE)); client_dir="$phase_dir/client_${replica}_port${port}"
    mkdir -p "$client_dir"
    "$PYTHON" "$CLIENT_SCRIPT" run --host "$HOST" --port "$port" --model "$MODEL" --task "$task" \
      --reference-image "$REFERENCE_IMAGE" --prompt-file "$PROMPT_FILE" --resolutions "$RESOLUTIONS" \
      --requests-per-resolution "$REQUESTS_PER_RESOLUTION" --replica-index "$replica" --num-replicas "$replicas" \
      --num-inference-steps "$NUM_INFERENCE_STEPS" --warmup-requests 1 --warmup-inference-steps 5 \
      --duration-seconds "$DURATION_SECONDS" --aspect-ratio 16:9 --output "$client_dir/results.jsonl" \
      >"$client_dir/client.log" 2>&1 &
    CLIENT_PIDS+=("$!")
    log "started task=$task client=$replica port=$port requests=$((REQUESTS_PER_TASK / replicas))"
  done
  for ((replica=0; replica<replicas; replica++)); do
    wait "${CLIENT_PIDS[$replica]}" || failed=1
  done
  CLIENT_PIDS=()
  "$PYTHON" "$CLIENT_SCRIPT" summarize --input-dir "$phase_dir" --output "$phase_dir/summary.json" \
    --task "$task" --tp "$tp" --replicas "$replicas" --expected-requests "$REQUESTS_PER_TASK" \
    >> "$SUMMARY_TSV" || failed=1
  return "$failed"
}

score_ssim() {
  local tp=$1 replicas=$2 task=$3 phase_dir=$4
  [[ -n "$SSIM_REFERENCE_ROOT" ]] || return 0
  local reference_phase="$SSIM_REFERENCE_ROOT/tp${tp}_replicas${replicas}/${task}"
  local quality_dir="$phase_dir/quality"
  [[ -d "$reference_phase" ]] || {
    log "ERROR: SSIM reference phase missing: $reference_phase" >&2
    return 1
  }
  mkdir -p "$quality_dir"
  local -a command=(
    "$PYTHON" "$SSIM_SCRIPT" compare
    --candidate-dir "$phase_dir"
    --reference-dir "$reference_phase"
    --output-dir "$quality_dir"
    --threshold "$SSIM_THRESHOLD"
    --ffmpeg "$FFMPEG_BIN"
    --ffprobe "$FFPROBE_BIN"
  )
  [[ "$SSIM_FAIL_BELOW_THRESHOLD" == true ]] && command+=(--fail-below-threshold)
  log "computing paired SSIM task=$task reference=$reference_phase"
  "${command[@]}" >"$quality_dir/paired_ssim.log" 2>&1
}

read -r -a TP_VALUES <<< "$TP_MATRIX"
read -r -a TASK_VALUES <<< "$TASKS"
for tp in "${TP_VALUES[@]}"; do
  ((TOTAL_GPUS % tp == 0)) || die "TOTAL_GPUS=$TOTAL_GPUS not divisible by TP=$tp"
  replicas=$((TOTAL_GPUS / tp))
  ((REQUESTS_PER_TASK % replicas == 0)) || die "requests/task not divisible by replicas"
  log "===== TP=$tp replicas=$replicas total_requests=$TOTAL_REQUESTS_PER_TP ====="
  for task in "${TASK_VALUES[@]}"; do
    [[ "$task" == fl2va ]] && variant=FL2VA || variant=Ref2VA
    phase_dir="$RESULT_ROOT/tp${tp}_replicas${replicas}/${task}"; mkdir -p "$phase_dir"
    start_servers "$tp" "$replicas" "$variant" "$phase_dir"
    phase_failed=0; run_clients "$tp" "$replicas" "$task" "$phase_dir" || phase_failed=1
    stop_servers
    ((phase_failed == 0)) && score_ssim "$tp" "$replicas" "$task" "$phase_dir" || phase_failed=1
    ((phase_failed == 0)) || die "TP=$tp task=$task had failed requests; inspect $phase_dir"
  done
done

trap - EXIT INT TERM
log "mixed matrix complete: $SUMMARY_TSV"