#!/usr/bin/env bash
# Eight real Ref2VA cases with Larry v4-600 LoRA: four TP2 replicas,
# two batches of four requests. Larry's 9 sigma grid points execute 8 NFE.
set -Eeuo pipefail

export PATH="/root/.miniconda3/envs/sglang-lora/bin:${PATH}"
export LD_LIBRARY_PATH="/root/.miniconda3/envs/sglang-lora/lib:/usr/local/cuda/lib64:${LD_LIBRARY_PATH:-}"

MODEL=${MODEL:-/data/hf_models/MiniMax-H3}
REPO_ROOT=${REPO_ROOT:-/data/wxy/sskj-h3}
INPUT_ROOT=${INPUT_ROOT:-$REPO_ROOT/throughput/sglang-base/inputs/ref2va-feishu-20260831}
RECORDS=${RECORDS:-$INPUT_ROOT/records.json}
ASSETS_ROOT=${ASSETS_ROOT:-$INPUT_ROOT/assets}
PYTHON=${PYTHON:-/root/.miniconda3/envs/sglang-lora/bin/python}
SGLANG_BIN=${SGLANG_BIN:-/root/.miniconda3/envs/sglang-lora/bin/sglang}
MEDIA_BIN_DIR=${MEDIA_BIN_DIR:-/root/.miniconda3/envs/deploy/bin}
CLIENT_SCRIPT=${CLIENT_SCRIPT:-$REPO_ROOT/throughput/sglang-base/scripts/ref2va_feishu_bench.py}
BASE_PORT=${BASE_PORT:-34010}
PORT_STRIDE=${PORT_STRIDE:-10}
MASTER_PORT_BASE=${MASTER_PORT_BASE:-35000}
SCHEDULER_PORT_BASE=${SCHEDULER_PORT_BASE:-36000}
HOST=${HOST:-127.0.0.1}
TP_SIZE=2
REPLICAS=4
NUM_INFERENCE_STEPS=${NUM_INFERENCE_STEPS:-9}
EXPECTED_DENOISER_EVALS=${EXPECTED_DENOISER_EVALS:-8}
SHORT_EDGE=${SHORT_EDGE:-768}
ASPECT_RATIO=${ASPECT_RATIO:-9:16}
DURATION_SECONDS=${DURATION_SECONDS:-15}
FLOW_SHIFT=${FLOW_SHIFT:-12.0}
AUDIO_FLOW_SHIFT=${AUDIO_FLOW_SHIFT:-3.0}
SERVER_START_TIMEOUT=${SERVER_START_TIMEOUT:-1800}
LORA_PATH=${LORA_PATH:-/data/hf_models/MiniMax-H3-Turbo-Lora}
LORA_WEIGHT_NAME=${LORA_WEIGHT_NAME:-minimax_h3_turbo_v4_step600_ema.safetensors}
LORA_NICKNAME=${LORA_NICKNAME:-h3-turbo-v4-600}
LORA_SCALE=${LORA_SCALE:-1.0}
LORA_MERGE_MODE=${LORA_MERGE_MODE:-auto}
REFERENCE_SUMMARY=${REFERENCE_SUMMARY:-$REPO_ROOT/throughput/sglang-base/results/ref2va-feishu-base-tp2x4-768p-15s-20steps-20260831-153946/summary.json}
SSIM_SCRIPT=${SSIM_SCRIPT:-$REPO_ROOT/throughput/common/ref2va_paired_ssim.py}
RUN_ID=${RUN_ID:-ref2va-feishu-larry-v4-600-tp2x4-768p-15s-8nfe-$(date '+%Y%m%d-%H%M%S')}
RESULT_ROOT=${RESULT_ROOT:-$REPO_ROOT/throughput/sglang-lora/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"
[[ -x "$MEDIA_BIN_DIR/ffprobe" ]] || die "ffprobe not executable: $MEDIA_BIN_DIR/ffprobe"
[[ -x "$MEDIA_BIN_DIR/ffmpeg" ]] || die "ffmpeg not executable: $MEDIA_BIN_DIR/ffmpeg"
[[ -f "$CLIENT_SCRIPT" ]] || die "client script missing: $CLIENT_SCRIPT"
[[ -f "$RECORDS" ]] || die "records missing: $RECORDS"
[[ -d "$ASSETS_ROOT" ]] || die "assets missing: $ASSETS_ROOT"
[[ -f "$LORA_PATH/$LORA_WEIGHT_NAME" ]] || die "LoRA weight missing: $LORA_PATH/$LORA_WEIGHT_NAME"
[[ -f "$REFERENCE_SUMMARY" ]] || die "Base summary missing: $REFERENCE_SUMMARY"
[[ -f "$SSIM_SCRIPT" ]] || die "SSIM script missing: $SSIM_SCRIPT"
[[ "$NUM_INFERENCE_STEPS" -eq 9 ]] || die "Larry v4-600 requires 9 sigma grid points (8 NFE)"
[[ "$EXPECTED_DENOISER_EVALS" -eq 8 ]] || die "expected 8 denoiser evaluations"
mkdir -p "$RESULT_ROOT"

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 "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 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
  nvidia-smi --query-gpu=index,pstate,memory.used,utilization.gpu --format=csv,noheader >"$RESULT_ROOT/gpu_after.csv" 2>&1 || true
  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 -120 "$log_file" >&2 || true; return 1; fi
    sleep 5
  done
  tail -120 "$log_file" >&2 || true
  return 1
}

start_servers() {
  local replica port master_port scheduler_port first_gpu gpu_csv offset gpu server_dir server_log candidate
  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_SIZE)); gpu_csv=""
    for ((offset=0; offset<TP_SIZE; offset++)); do
      gpu=$((first_gpu + offset)); [[ -z "$gpu_csv" ]] && gpu_csv="$gpu" || gpu_csv+=",$gpu"
    done
    server_dir="$RESULT_ROOT/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 Larry v4-600 Ref2VA replica=$replica GPUs=$gpu_csv port=$port"
    CUDA_VISIBLE_DEVICES="$gpu_csv" \
      PATH="$MEDIA_BIN_DIR:$PATH" \
      PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
      PYTHONUNBUFFERED=1 TOKENIZERS_PARALLELISM=false SGLANG_USE_RUNAI_MODEL_STREAMER=false \
      setsid "$SGLANG_BIN" serve --model-path "$MODEL" --model-variant Ref2VA \
      --backend sglang --performance-mode speed --num-gpus "$TP_SIZE" --tp-size "$TP_SIZE" \
      --ulysses-degree 1 --use-fsdp-inference false --enable-torch-compile false \
      --batching-max-size 1 --batching-delay-ms 0 \
      --lora-path "$LORA_PATH" --lora-weight-name "$LORA_WEIGHT_NAME" \
      --lora-nickname "$LORA_NICKNAME" --lora-scale "$LORA_SCALE" \
      --lora-merge-mode "$LORA_MERGE_MODE" \
      --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="$RESULT_ROOT/server_${replica}_port${port}/server.log"
    wait_healthy "$port" "${SERVER_PIDS[$replica]}" "$server_log" || die "replica=$replica failed startup"
    log "replica=$replica healthy port=$port"
  done
}

run_batch() {
  local batch=$1 first_case=$2 warmup_steps=$3 replica port case_number client_dir failed=0
  local batch_dir="$RESULT_ROOT/batch_$(printf '%02d' "$batch")"
  mkdir -p "$batch_dir"
  CLIENT_PIDS=()
  for ((replica=0; replica<REPLICAS; replica++)); do
    port=$((BASE_PORT + replica * PORT_STRIDE))
    case_number=$((first_case + replica))
    client_dir="$batch_dir/client_${replica}_case$(printf '%02d' "$case_number")_port${port}"
    mkdir -p "$client_dir"
    "$PYTHON" "$CLIENT_SCRIPT" run --host "$HOST" --port "$port" --replica-index "$replica" \
      --case-number "$case_number" --records "$RECORDS" --assets-root "$ASSETS_ROOT" \
      --model "$MODEL" --num-inference-steps "$NUM_INFERENCE_STEPS" --warmup-steps "$warmup_steps" \
      --short-edge "$SHORT_EDGE" --aspect-ratio "$ASPECT_RATIO" --duration-seconds "$DURATION_SECONDS" \
      --flow-shift "$FLOW_SHIFT" --audio-flow-shift "$AUDIO_FLOW_SHIFT" \
      --output "$client_dir/result.json" >"$client_dir/client.log" 2>&1 &
    CLIENT_PIDS+=("$!")
    log "batch=$batch started case=$case_number replica=$replica port=$port"
  done
  for pid in "${CLIENT_PIDS[@]}"; do wait "$pid" || failed=1; done
  CLIENT_PIDS=()
  ((failed == 0)) || die "batch=$batch had failed requests"
  log "batch=$batch complete"
}

{
  printf 'run_id=%s\n' "$RUN_ID"
  printf 'method=larry-v4-600-lora\ntask=ref2va\ntp_size=2\nreplicas=4\n'
  printf 'num_inference_steps=%s\nshort_edge=%s\naspect_ratio=%s\nduration_seconds=%s\n' \
    "$NUM_INFERENCE_STEPS" "$SHORT_EDGE" "$ASPECT_RATIO" "$DURATION_SECONDS"
  printf 'denoiser_evaluations=%s\nlora_path=%s\nlora_weight_name=%s\n' \
    "$EXPECTED_DENOISER_EVALS" "$LORA_PATH" "$LORA_WEIGHT_NAME"
  printf 'lora_nickname=%s\nlora_scale=%s\nlora_merge_mode=%s\n' \
    "$LORA_NICKNAME" "$LORA_SCALE" "$LORA_MERGE_MODE"
  printf 'flow_shift=%s\naudio_flow_shift=%s\nmodel=%s\nmedia_bin_dir=%s\n' \
    "$FLOW_SHIFT" "$AUDIO_FLOW_SHIFT" "$MODEL" "$MEDIA_BIN_DIR"
  printf 'cache_dit_enabled=false\nfsdp_inference=false\ntorch_compile=false\nbatching_max_size=1\n'
  printf 'pytorch_cuda_alloc_conf=expandable_segments:True\n'
} >"$RESULT_ROOT/run_config.env"
cp "$RECORDS" "$RESULT_ROOT/input_records.json"
sha256sum "$RECORDS" "$CLIENT_SCRIPT" "$0" >"$RESULT_ROOT/source_sha256.txt"
find "$ASSETS_ROOT" -type f -print0 | sort -z | xargs -0 sha256sum >"$RESULT_ROOT/input_assets_sha256.txt"
nvidia-smi >"$RESULT_ROOT/nvidia_smi_before.txt"
nvidia-smi --query-gpu=index,pstate,memory.used,utilization.gpu --format=csv,noheader >"$RESULT_ROOT/gpu_before.csv"
"$PYTHON" -V >"$RESULT_ROOT/python_version.txt" 2>&1
"$PYTHON" -m pip show sglang >"$RESULT_ROOT/sglang_pip_show.txt" 2>&1 || true
"$PYTHON" -m pip freeze >"$RESULT_ROOT/pip_freeze.txt" 2>&1

start_servers
run_batch 1 1 "$NUM_INFERENCE_STEPS"
run_batch 2 5 0
stop_servers
"$PYTHON" "$CLIENT_SCRIPT" summarize --result-root "$RESULT_ROOT" --output "$RESULT_ROOT/summary.json" \
  >"$RESULT_ROOT/summary.log" 2>&1
"$PYTHON" - "$RESULT_ROOT/summary.json" <<'PY'
import json, pathlib, sys
path = pathlib.Path(sys.argv[1])
data = json.loads(path.read_text(encoding="utf-8"))
data.update({
    "method": "larry-v4-600-lora",
    "num_inference_steps_grid_points": 9,
    "denoiser_evaluations": 8,
    "lora_scale": 1.0,
    "ref2va_support": "experimental",
})
path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
PY
mkdir -p "$RESULT_ROOT/ssim_vs_base"
"$PYTHON" "$SSIM_SCRIPT" \
  --reference-summary "$REFERENCE_SUMMARY" \
  --candidate-summary "$RESULT_ROOT/summary.json" \
  --output-dir "$RESULT_ROOT/ssim_vs_base" \
  --threshold 0.90 \
  --ffmpeg "$MEDIA_BIN_DIR/ffmpeg" \
  --ffprobe "$MEDIA_BIN_DIR/ffprobe" \
  >"$RESULT_ROOT/ssim_vs_base/score.log" 2>&1
nvidia-smi --query-gpu=index,pstate,memory.used,utilization.gpu --format=csv,noheader >"$RESULT_ROOT/gpu_after.csv"
touch "$RESULT_ROOT/DONE"
trap - EXIT INT TERM
log "Ref2VA Feishu Larry v4-600 TP2x4 complete: $RESULT_ROOT"