#!/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}"
WAIT_FOR_GPUS="${WAIT_FOR_GPUS:-1}"
GPU_MEMORY_THRESHOLD_MIB="${GPU_MEMORY_THRESHOLD_MIB:-2048}"
POLL_SECONDS="${POLL_SECONDS:-30}"
RUN_SMOKE_TEST="${RUN_SMOKE_TEST:-1}"
START_TRAINING="${START_TRAINING:-0}"
DATASET_ROOT="${DATASET_ROOT:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_1000_seed0}"
SMOKE_ROOT="${SMOKE_ROOT:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_smoke_seed0}"
NVME_TRAIN_ROOT="${NVME_TRAIN_ROOT:-/mnt/local_nvme/zoubin/cz/self_forcing_predictor_v4_training_1000_seed0}"
RUN_DIR="${RUN_DIR:-$ROOT_DIR/checkpoints/predictor_v4_1000_seed0}"

mkdir -p "$RUN_DIR"
exec 9>"$RUN_DIR/end_to_end.lock"
if ! flock -n 9; then
  echo "错误：已有 Predictor-v4 end-to-end supervisor 正在运行。" >&2
  exit 1
fi
echo "$$" >"$RUN_DIR/end_to_end.pid"
trap 'rm -f "$RUN_DIR/end_to_end.pid"' EXIT

IFS=',' read -r -a GPU_ARRAY <<< "$GPU_IDS"
if (( ${#GPU_ARRAY[@]} != 8 )); then
  echo "错误：正式构建与训练需要 8 张 GPU，当前 GPU_IDS=$GPU_IDS。" >&2
  exit 1
fi
if [[ "$WAIT_FOR_GPUS" != "0" && "$WAIT_FOR_GPUS" != "1" ]]; then
  echo "错误：WAIT_FOR_GPUS 必须为 0 或 1。" >&2
  exit 1
fi
if [[ "$START_TRAINING" != "0" && "$START_TRAINING" != "1" ]]; then
  echo "错误：START_TRAINING 必须为 0 或 1。" >&2
  exit 1
fi

gpus_are_free() {
  local used
  mapfile -t used < <(
    nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits
  )
  if (( ${#used[@]} < 8 )); then
    return 1
  fi
  local gpu_id
  for gpu_id in "${GPU_ARRAY[@]}"; do
    if (( used[gpu_id] > GPU_MEMORY_THRESHOLD_MIB )); then
      return 1
    fi
  done
}

if [[ "$WAIT_FOR_GPUS" == "1" ]]; then
  while ! gpus_are_free; do
    echo "[$(date -u +%FT%TZ)] GPU 尚忙，${POLL_SECONDS}s 后重试。" >&2
    sleep "$POLL_SECONDS"
  done
elif ! gpus_are_free; then
  echo "错误：GPU 显存占用超过 ${GPU_MEMORY_THRESHOLD_MIB} MiB。" >&2
  exit 1
fi

if [[ "$RUN_SMOKE_TEST" == "1" ]]; then
  echo "[$(date -u +%FT%TZ)] 开始单 case GPU 冒烟测试。"
  "$PYTHON_BIN" scripts/prepare_predictor_v4_prompts.py \
    --output_root "$SMOKE_ROOT" \
    --num_prompts 1 \
    --sampling_seed 0 \
    --inference_seed 0
  CUDA_VISIBLE_DEVICES="${GPU_ARRAY[0]}" \
    "$PYTHON_BIN" scripts/build_predictor_v4_dataset.py \
      --dataset_root "$SMOKE_ROOT" \
      --worker_id 0 \
      --num_workers 1 \
      --case_ids 0 \
      --seed 0 \
      --blocks 0 1 28 29 \
      --validate_kv_rebuild
  "$PYTHON_BIN" scripts/merge_predictor_v4_manifests.py \
    --dataset_root "$SMOKE_ROOT" \
    --num_workers 1
  "$PYTHON_BIN" scripts/validate_predictor_v4_dataset.py \
    --dataset_root "$SMOKE_ROOT" \
    --tensor_records 0 \
    --require_kv_metrics 7
  echo "[$(date -u +%FT%TZ)] 单 case GPU 冒烟测试通过。"
fi

echo "[$(date -u +%FT%TZ)] 开始 1000 prompts / 8 GPU 离线数据构建。"
GPU_IDS="$GPU_IDS" \
DATASET_ROOT="$DATASET_ROOT" \
PREPARE_PROMPTS=1 \
VALIDATE_KV_REBUILD=0 \
bash scripts/launch_predictor_v4_build.sh

if [[ "$START_TRAINING" == "0" ]]; then
  echo "[$(date -u +%FT%TZ)] 数据构建完成；START_TRAINING=0，不启动训练。"
  exit 0
fi

echo "[$(date -u +%FT%TZ)] 数据构建完成，开始 8 GPU Predictor-v4 训练。"
NPROC_PER_NODE=8 \
LOG_DIR="$NVME_TRAIN_ROOT" \
bash scripts/run_predictor_v4_training.sh

echo "[$(date -u +%FT%TZ)] 训练完成，开始把 NVMe 产物同步回项目目录。"
DATASET_SOURCE="$DATASET_ROOT" \
TRAIN_SOURCE="$NVME_TRAIN_ROOT" \
bash scripts/sync_predictor_v4_artifacts.sh

echo "[$(date -u +%FT%TZ)] Predictor-v4 数据构建与训练全部完成。"