#!/usr/bin/env bash
# MiniMax-H3 dense baseline: 4 replicas x TP=2, VBench-v1 standard suite.
# Default paper-oriented protocol follows LoSA: full suite, one fixed seed.
set -Eeuo pipefail

ACTION=${ACTION:-all}                         # preflight | generate | evaluate | all
RUN_ID=${RUN_ID:-dense-tp2x4-$(date '+%Y%m%d-%H%M%S')}
RESULT_ROOT=${RESULT_ROOT:-/data/wxy/results/h3_vbench_base/$RUN_ID}
MODEL=${MODEL:-/data/hf_models/MiniMax-H3}
SGLANG_PYTHON=${SGLANG_PYTHON:-/root/.miniconda3/envs/sglang/bin/python}
SGLANG_BIN=${SGLANG_BIN:-/root/.miniconda3/envs/sglang/bin/sglang}
MEDIA_BIN_DIR=${MEDIA_BIN_DIR:-/root/.miniconda3/envs/deploy/bin}
CLIENT_SCRIPT=${CLIENT_SCRIPT:-/data/wxy/h3_vbench_generate.py}

VBENCH_REPO=${VBENCH_REPO:-/data/wxy/VBench}
VBENCH_PYTHON=${VBENCH_PYTHON:-/root/.miniconda3/envs/vbench/bin/python}
VBENCH_METADATA=${VBENCH_METADATA:-$VBENCH_REPO/vbench/VBench_full_info.json}
VBENCH_SAMPLE_DATA=${VBENCH_SAMPLE_DATA:-/data/datasets/vbench}
EVAL_GPUS=${EVAL_GPUS:-8}
RUN_STATIC_FILTER=${RUN_STATIC_FILTER:-1}       # Official temporal-flickering preprocessing.

HOST=${HOST:-127.0.0.1}
PORTS=${PORTS:-30010,30020,30030,30040}
PORT_STRIDE=${PORT_STRIDE:-10}
MASTER_PORT_BASE=${MASTER_PORT_BASE:-31000}
SCHEDULER_PORT_BASE=${SCHEDULER_PORT_BASE:-32000}
SERVER_START_TIMEOUT=${SERVER_START_TIMEOUT:-1800}
GPU_MEMORY_LIMIT_MB=${GPU_MEMORY_LIMIT_MB:-1000}

NUM_INFERENCE_STEPS=${NUM_INFERENCE_STEPS:-20}
DURATION_SECONDS=${DURATION_SECONDS:-5.0}
SHORT_EDGE=${SHORT_EDGE:-768}
ASPECT_RATIO=${ASPECT_RATIO:-16:9}
SAMPLES_PER_PROMPT=${SAMPLES_PER_PROMPT:-1}   # LoSA-style fixed seed; set 5 for strict VBench protocol.
MAX_PROMPTS=${MAX_PROMPTS:-0}                 # 0 = full standard suite; >0 = smoke only, not paper-comparable.
SEED=${SEED:-1101}
WARMUP_REQUESTS=${WARMUP_REQUESTS:-1}
WARMUP_INFERENCE_STEPS=${WARMUP_INFERENCE_STEPS:-5}
FLOW_SHIFT=${FLOW_SHIFT:-12.0}
AUDIO_FLOW_SHIFT=${AUDIO_FLOW_SHIFT:-3.0}

DIMENSIONS=(
  subject_consistency background_consistency temporal_flickering motion_smoothness
  dynamic_degree aesthetic_quality imaging_quality object_class multiple_objects
  human_action color spatial_relationship scene temporal_style appearance_style
  overall_consistency
)
DIMENSIONS_CSV=$(IFS=,; printf '%s' "${DIMENSIONS[*]}")
IFS=, read -r -a PORT_ARRAY <<< "$PORTS"
[[ "${#PORT_ARRAY[@]}" -eq 4 ]] || { printf 'PORTS must contain exactly four ports\n' >&2; exit 1; }

VIDEOS_DIR="$RESULT_ROOT/videos"
GENERATION_DIR="$RESULT_ROOT/generation"
EVAL_DIR="$RESULT_ROOT/evaluation"
SERVER_DIR="$RESULT_ROOT/servers"
MANIFEST="$GENERATION_DIR/manifest.jsonl"
RESULTS="$GENERATION_DIR/results.jsonl"
GENERATION_SUMMARY="$GENERATION_DIR/summary.json"
RUN_CONFIG="$RESULT_ROOT/run_config.env"

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() {
  "$SGLANG_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
}

validate_vbench_layout() {
  [[ -f "$VBENCH_METADATA" ]] || {
    if [[ -d "$VBENCH_SAMPLE_DATA" ]]; then
      die "$VBENCH_METADATA is missing. $VBENCH_SAMPLE_DATA is the VBench-2.0 sample-video dataset, not the VBench-v1 evaluator repository. Clone/install official VBench and set VBENCH_REPO."
    fi
    die "VBench-v1 metadata missing: $VBENCH_METADATA"
  }
}

check_gpu_idle() {
  command -v nvidia-smi >/dev/null || die "nvidia-smi not found"
  local gpu_count used
  gpu_count=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)
  [[ "$gpu_count" -eq 8 ]] || die "expected 8 GPUs, found $gpu_count"
  while IFS= read -r used; do
    used=${used// /}
    ((used <= GPU_MEMORY_LIMIT_MB)) || die "GPU memory is already in use (${used} MiB > ${GPU_MEMORY_LIMIT_MB} MiB); refusing to touch unrelated processes"
  done < <(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits)
}

preflight_generation() {
  [[ -x "$SGLANG_PYTHON" ]] || die "SGLang python not executable: $SGLANG_PYTHON"
  [[ -x "$SGLANG_BIN" ]] || die "SGLang executable missing: $SGLANG_BIN"
  [[ -x "$MEDIA_BIN_DIR/ffmpeg" ]] || die "ffmpeg missing: $MEDIA_BIN_DIR/ffmpeg"
  [[ -x "$MEDIA_BIN_DIR/ffprobe" ]] || die "ffprobe missing: $MEDIA_BIN_DIR/ffprobe"
  [[ -f "$CLIENT_SCRIPT" ]] || die "generation client missing: $CLIENT_SCRIPT"
  [[ -d "$MODEL" ]] || die "model missing: $MODEL"
  validate_vbench_layout
  check_gpu_idle
  local replica candidate port
  for ((replica=0; replica<4; ++replica)); do
    port=${PORT_ARRAY[$replica]}
    for candidate in "$port" "$((port + 1))" "$((MASTER_PORT_BASE + replica * PORT_STRIDE))" "$((SCHEDULER_PORT_BASE + replica * PORT_STRIDE))"; do
      port_is_open "$candidate" && die "port already in use: $candidate"
    done
  done
  return 0
}

preflight_evaluation() {
  validate_vbench_layout
  [[ -f "$VBENCH_REPO/evaluate.py" ]] || die "official evaluate.py missing under VBENCH_REPO=$VBENCH_REPO"
  [[ -x "$VBENCH_PYTHON" ]] || die "VBench python not executable: $VBENCH_PYTHON"
  (
    cd "$VBENCH_REPO"
    "$VBENCH_PYTHON" -c 'import torch, vbench'
  ) || die "VBENCH_PYTHON cannot import torch and vbench from $VBENCH_REPO"
  ((EVAL_GPUS >= 1 && EVAL_GPUS <= 8)) || die "EVAL_GPUS must be in [1,8]"
  if ((RUN_STATIC_FILTER)); then
    [[ -f "$VBENCH_REPO/static_filter.py" ]] || die "official static_filter.py missing under VBENCH_REPO=$VBENCH_REPO"
  fi
}

stop_servers() {
  local pid alive deadline
  ((${#SERVER_PIDS[@]})) || return 0
  log "stopping only the ${#SERVER_PIDS[@]} server process groups started by this run"
  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=$?
  trap - EXIT INT TERM
  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 replica port master_port scheduler_port gpu_csv server_log
  mkdir -p "$SERVER_DIR"
  for ((replica=0; replica<4; ++replica)); do
    port=${PORT_ARRAY[$replica]}
    master_port=$((MASTER_PORT_BASE + replica * PORT_STRIDE))
    scheduler_port=$((SCHEDULER_PORT_BASE + replica * PORT_STRIDE))
    gpu_csv="$((replica * 2)),$((replica * 2 + 1))"
    mkdir -p "$SERVER_DIR/replica${replica}_port${port}/outputs"
    server_log="$SERVER_DIR/replica${replica}_port${port}/server.log"
    printf '%s\n' "$gpu_csv" > "$SERVER_DIR/replica${replica}_port${port}/cuda_visible_devices.txt"
    log "starting FL2VA replica=$replica TP=2 GPUs=$gpu_csv port=$port"
    PATH="$MEDIA_BIN_DIR:$PATH" 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 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/replica${replica}_port${port}/outputs" >"$server_log" 2>&1 &
    SERVER_PIDS+=("$!")
    printf '%s\n' "$!" > "$SERVER_DIR/replica${replica}_port${port}/server.pid"
  done
  for ((replica=0; replica<4; ++replica)); do
    port=${PORT_ARRAY[$replica]}
    server_log="$SERVER_DIR/replica${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
  return 0
}

write_config() {
  mkdir -p "$RESULT_ROOT" "$GENERATION_DIR" "$VIDEOS_DIR" "$EVAL_DIR"
  {
    printf 'RUN_ID=%q\n' "$RUN_ID"
    printf 'MODEL=%q\n' "$MODEL"
    printf 'TOPOLOGY=%q\n' 'TP2x4'
    printf 'PORTS=%q\n' "$PORTS"
    printf 'VBENCH_REPO=%q\n' "$VBENCH_REPO"
    printf 'VBENCH_METADATA=%q\n' "$VBENCH_METADATA"
    printf 'NUM_INFERENCE_STEPS=%q\n' "$NUM_INFERENCE_STEPS"
    printf 'DURATION_SECONDS=%q\n' "$DURATION_SECONDS"
    printf 'SHORT_EDGE=%q\n' "$SHORT_EDGE"
    printf 'ASPECT_RATIO=%q\n' "$ASPECT_RATIO"
    printf 'SAMPLES_PER_PROMPT=%q\n' "$SAMPLES_PER_PROMPT"
    printf 'MAX_PROMPTS=%q\n' "$MAX_PROMPTS"
    printf 'SEED=%q\n' "$SEED"
    printf 'RUN_STATIC_FILTER=%q\n' "$RUN_STATIC_FILTER"
    printf 'DIMENSIONS=%q\n' "$DIMENSIONS_CSV"
  } > "$RUN_CONFIG"
}

generate() {
  preflight_generation
  write_config
  start_servers
  log "generation begins: full_suite=$([[ "$MAX_PROMPTS" -eq 0 ]] && printf true || printf false) samples_per_prompt=$SAMPLES_PER_PROMPT"
  "$SGLANG_PYTHON" "$CLIENT_SCRIPT" \
    --metadata "$VBENCH_METADATA" --videos-dir "$VIDEOS_DIR" --results "$RESULTS" \
    --manifest "$MANIFEST" --summary "$GENERATION_SUMMARY" --host "$HOST" --ports "$PORTS" \
    --model "$MODEL" --dimensions "$DIMENSIONS_CSV" --samples-per-prompt "$SAMPLES_PER_PROMPT" \
    --max-prompts "$MAX_PROMPTS" --num-inference-steps "$NUM_INFERENCE_STEPS" \
    --short-edge "$SHORT_EDGE" --duration-seconds "$DURATION_SECONDS" --aspect-ratio "$ASPECT_RATIO" \
    --seed "$SEED" --warmup-requests "$WARMUP_REQUESTS" --warmup-inference-steps "$WARMUP_INFERENCE_STEPS" \
    --flow-shift "$FLOW_SHIFT" --audio-flow-shift "$AUDIO_FLOW_SHIFT" \
    > "$GENERATION_DIR/client.log" 2>&1
  stop_servers
  log "generation complete: $GENERATION_SUMMARY"
}

evaluate() {
  preflight_evaluation
  [[ -f "$GENERATION_SUMMARY" ]] || die "generation summary missing: $GENERATION_SUMMARY"
  "$SGLANG_PYTHON" - "$GENERATION_SUMMARY" <<'PY'
import json, sys
summary = json.load(open(sys.argv[1]))
if summary["failed_or_missing"]:
    raise SystemExit(f"generation is incomplete: {summary}")
PY
  if ((MAX_PROMPTS > 0)); then
    die "MAX_PROMPTS=$MAX_PROMPTS is a smoke subset; VBench standard scores would be incomplete and misleading"
  fi
  mkdir -p "$EVAL_DIR/raw_main"
  local result_json flicker_json score_zip
  if ((RUN_STATIC_FILTER)); then
    local main_dimensions=(
      subject_consistency background_consistency motion_smoothness dynamic_degree
      aesthetic_quality imaging_quality object_class multiple_objects human_action color
      spatial_relationship scene temporal_style appearance_style overall_consistency
    )
    log "VBench-v1 evaluates 15 dimensions on $EVAL_GPUS GPUs"
    (
      cd "$VBENCH_REPO"
      CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 "$VBENCH_PYTHON" -m torch.distributed.run \
        --nproc_per_node="$EVAL_GPUS" --standalone evaluate.py \
        --videos_path "$VIDEOS_DIR" --output_path "$EVAL_DIR/raw_main" \
        --full_json_dir "$VBENCH_METADATA" --mode vbench_standard \
        --dimension "${main_dimensions[@]}"
    ) > "$EVAL_DIR/evaluate_main.log" 2>&1
    log "running official static filter before temporal_flickering"
    (
      cd "$VBENCH_REPO"
      CUDA_VISIBLE_DEVICES=0 "$VBENCH_PYTHON" static_filter.py \
        --videos_path "$VIDEOS_DIR" --result_path "$EVAL_DIR/static_filter"
    ) > "$EVAL_DIR/static_filter.log" 2>&1
    [[ -n "$(find "$EVAL_DIR/static_filter/filtered_videos" -maxdepth 1 -type f -name '*.mp4' -print -quit 2>/dev/null)" ]] || \
      die "static filter produced no videos; inspect $EVAL_DIR/static_filter.log"
    mkdir -p "$EVAL_DIR/raw_flicker"
    (
      cd "$VBENCH_REPO"
      CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 "$VBENCH_PYTHON" -m torch.distributed.run \
        --nproc_per_node="$EVAL_GPUS" --standalone evaluate.py \
        --videos_path "$EVAL_DIR/static_filter/filtered_videos" --output_path "$EVAL_DIR/raw_flicker" \
        --full_json_dir "$VBENCH_METADATA" --mode vbench_standard --dimension temporal_flickering
    ) > "$EVAL_DIR/evaluate_flicker.log" 2>&1
    result_json=$(find "$EVAL_DIR/raw_main" -maxdepth 1 -type f -name '*_eval_results.json' -print -quit)
    flicker_json=$(find "$EVAL_DIR/raw_flicker" -maxdepth 1 -type f -name '*_eval_results.json' -print -quit)
    [[ -n "$result_json" && -n "$flicker_json" ]] || die "VBench result JSON missing; inspect evaluation logs"
    "$VBENCH_PYTHON" - "$result_json" "$flicker_json" "$EVAL_DIR/vbench_16_dimensions.json" <<'PY'
import json, sys
with open(sys.argv[1], encoding="utf-8") as handle:
    result = json.load(handle)
with open(sys.argv[2], encoding="utf-8") as handle:
    result.update(json.load(handle))
with open(sys.argv[3], "w", encoding="utf-8") as handle:
    json.dump(result, handle, ensure_ascii=False, indent=2)
PY
  else
    log "VBench-v1 evaluates all 16 dimensions without static filtering (non-standard flicker preprocessing)"
    (
      cd "$VBENCH_REPO"
      CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 "$VBENCH_PYTHON" -m torch.distributed.run \
        --nproc_per_node="$EVAL_GPUS" --standalone evaluate.py \
        --videos_path "$VIDEOS_DIR" --output_path "$EVAL_DIR/raw_main" \
        --full_json_dir "$VBENCH_METADATA" --mode vbench_standard \
        --dimension "${DIMENSIONS[@]}"
    ) > "$EVAL_DIR/evaluate_main.log" 2>&1
    result_json=$(find "$EVAL_DIR/raw_main" -maxdepth 1 -type f -name '*_eval_results.json' -print -quit)
    [[ -n "$result_json" ]] || die "VBench produced no *_eval_results.json; inspect $EVAL_DIR/evaluate_main.log"
    cp "$result_json" "$EVAL_DIR/vbench_16_dimensions.json"
  fi
  score_zip="$EVAL_DIR/evaluation_results.zip"
  "$VBENCH_PYTHON" - "$EVAL_DIR/vbench_16_dimensions.json" "$score_zip" <<'PY'
import os, sys, zipfile
source, destination = sys.argv[1:]
with zipfile.ZipFile(destination, "w", compression=zipfile.ZIP_DEFLATED) as archive:
    archive.write(source, arcname=os.path.basename(source))
PY
  (
    cd "$EVAL_DIR"
    "$VBENCH_PYTHON" "$VBENCH_REPO/scripts/cal_final_score.py" \
      --zip_file "$score_zip" --model_name aggregate_input
  ) > "$EVAL_DIR/final_score.txt" 2>&1
  log "evaluation complete: $EVAL_DIR/vbench_16_dimensions.json and $EVAL_DIR/final_score.txt"
}

case "$ACTION" in
  preflight)
    preflight_generation
    log "generation preflight passed"
    if [[ -x "$VBENCH_PYTHON" && -f "$VBENCH_REPO/evaluate.py" ]]; then
      preflight_evaluation
      log "evaluation preflight passed"
    else
      log "evaluation preflight skipped: VBench code/env is not ready"
    fi
    ;;
  generate) generate ;;
  evaluate) write_config; evaluate ;;
  all) generate; evaluate ;;
  *) die "unknown ACTION=$ACTION (use preflight, generate, evaluate, or all)" ;;
esac

trap - EXIT INT TERM
log "done: $RESULT_ROOT"