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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT_DIR"

PYTHON_BIN="${PYTHON_BIN:-/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/self_forcing/bin/python}"
GPU_IDS="${GPU_IDS:-0,1,2,3,4,5,6,7}"
NUM_WORKERS="${NUM_WORKERS:-8}"
DATASET_ROOT="${DATASET_ROOT:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_1000_seed0}"
CONFIG_PATH="${CONFIG_PATH:-$ROOT_DIR/configs/self_forcing_dmd.yaml}"
CHECKPOINT_PATH="${CHECKPOINT_PATH:-$ROOT_DIR/checkpoints/self_forcing_dmd.pt}"
PREPARE_PROMPTS="${PREPARE_PROMPTS:-1}"
VALIDATE_KV_REBUILD="${VALIDATE_KV_REBUILD:-0}"
DRY_RUN="${DRY_RUN:-0}"

IFS=',' read -r -a GPU_ARRAY <<< "$GPU_IDS"
if (( ${#GPU_ARRAY[@]} != NUM_WORKERS )); then
  echo "错误：GPU_IDS 数量 ${#GPU_ARRAY[@]} 与 NUM_WORKERS=$NUM_WORKERS 不一致。" >&2
  exit 1
fi
if (( NUM_WORKERS != 8 )); then
  echo "警告：正式数据构建约定为 8 workers；当前 NUM_WORKERS=$NUM_WORKERS。" >&2
fi
if [[ "$PREPARE_PROMPTS" != "0" && "$PREPARE_PROMPTS" != "1" ]]; then
  echo "错误：PREPARE_PROMPTS 必须为 0 或 1。" >&2
  exit 1
fi
if [[ "$VALIDATE_KV_REBUILD" != "0" && "$VALIDATE_KV_REBUILD" != "1" ]]; then
  echo "错误：VALIDATE_KV_REBUILD 必须为 0 或 1。" >&2
  exit 1
fi
for required_path in "$PYTHON_BIN" "$CONFIG_PATH" "$CHECKPOINT_PATH"; do
  if [[ ! -e "$required_path" ]]; then
    echo "错误：找不到 $required_path" >&2
    exit 1
  fi
done

PREPARE_COMMAND=(
  "$PYTHON_BIN" scripts/prepare_predictor_v4_prompts.py
  --output_root "$DATASET_ROOT"
  --num_prompts 1000
  --sampling_seed 0
  --inference_seed 0
)

echo "Predictor v4 离线数据构建："
echo "  Dataset:   $DATASET_ROOT"
echo "  GPUs:      $GPU_IDS"
echo "  Workers:   $NUM_WORKERS"
echo "  Teacher:   $CHECKPOINT_PATH [generator_ema]"
echo "  Schedule:  F-F-F-F"
echo "  Seed:      每个 case 重置为 0"
echo "  Blocks:    0 1 28 29"
echo "  KV check:  $VALIDATE_KV_REBUILD"

if [[ "$DRY_RUN" == "1" ]]; then
  printf '  Prepare:   '
  printf '%q ' "${PREPARE_COMMAND[@]}"
  printf '\n'
  for worker_id in $(seq 0 $((NUM_WORKERS - 1))); do
    echo "  Worker $worker_id: CUDA_VISIBLE_DEVICES=${GPU_ARRAY[$worker_id]} ... --worker_id $worker_id"
  done
  echo "DRY_RUN=1：未创建 prompt split，也未启动 GPU 构建。"
  exit 0
fi

if [[ "$PREPARE_PROMPTS" == "1" && ! -f "$DATASET_ROOT/cases.jsonl" ]]; then
  "${PREPARE_COMMAND[@]}"
fi
if [[ ! -f "$DATASET_ROOT/cases.jsonl" ]]; then
  echo "错误：缺少 $DATASET_ROOT/cases.jsonl，请先运行 prompt 准备脚本。" >&2
  exit 1
fi
mkdir -p "$DATASET_ROOT/logs"

pids=()
for worker_id in $(seq 0 $((NUM_WORKERS - 1))); do
  WORKER_COMMAND=(
    "$PYTHON_BIN" scripts/build_predictor_v4_dataset.py
    --dataset_root "$DATASET_ROOT"
    --config "$CONFIG_PATH"
    --checkpoint "$CHECKPOINT_PATH"
    --worker_id "$worker_id"
    --num_workers "$NUM_WORKERS"
    --seed 0
    --blocks 0 1 28 29
  )
  if [[ "$VALIDATE_KV_REBUILD" == "1" ]]; then
    WORKER_COMMAND+=(--validate_kv_rebuild)
  fi
  CUDA_VISIBLE_DEVICES="${GPU_ARRAY[$worker_id]}" \
    "${WORKER_COMMAND[@]}" \
    >"$DATASET_ROOT/logs/worker_${worker_id}.log" 2>&1 &
  worker_pid="$!"
  pids+=("$worker_pid")
  echo "已启动 worker $worker_id，GPU=${GPU_ARRAY[$worker_id]}，PID=$worker_pid"
done

terminate_workers() {
  for pid in "${pids[@]}"; do
    kill "$pid" 2>/dev/null || true
  done
}
trap terminate_workers INT TERM

failed=0
for index in "${!pids[@]}"; do
  if ! wait "${pids[$index]}"; then
    echo "错误：worker $index 失败，见 $DATASET_ROOT/logs/worker_${index}.log" >&2
    failed=1
  fi
done
trap - INT TERM
if (( failed != 0 )); then
  exit 1
fi

"$PYTHON_BIN" scripts/merge_predictor_v4_manifests.py \
  --dataset_root "$DATASET_ROOT" \
  --num_workers "$NUM_WORKERS"

VALIDATE_ARGS=(
  scripts/validate_predictor_v4_dataset.py
  --dataset_root "$DATASET_ROOT"
  --tensor_records 16
)
if [[ "$VALIDATE_KV_REBUILD" == "1" ]]; then
  VALIDATE_ARGS+=(--require_kv_metrics 1)
fi
"$PYTHON_BIN" "${VALIDATE_ARGS[@]}"
echo "Predictor v4 离线数据构建和验证完成：$DATASET_ROOT"