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

REPO_ROOT="${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe}"
DATA_ROOT="${DATA_ROOT:-/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset}"
LOG_DIR="${LOG_DIR:-$DATA_ROOT/logs}"
SLEEP_SECONDS="${SLEEP_SECONDS:-300}"
EXPECTED_MIN_PARQUETS="${EXPECTED_MIN_PARQUETS:-1600}"

mkdir -p "$LOG_DIR"

echo "[wait] started at $(date -Is)"
echo "[wait] repo=$REPO_ROOT"
echo "[wait] data_root=$DATA_ROOT"

while true; do
  if ! docker ps --format '{{.ID}} {{.Image}} {{.Command}} {{.Names}}' \
      | grep -q 'modelscope.*download.*laoyao_2b_moe_pretrain_parquet_20260702'; then
    if ! pgrep -af 'modelscope download --dataset eigentom/laoyao_2b_moe_pretrain_parquet_20260702' >/dev/null; then
      break
    fi
  fi
  echo "[wait] modelscope download still running at $(date -Is)"
  du -sh "$DATA_ROOT" || true
  find "$DATA_ROOT" -name '*.parquet' | wc -l || true
  sleep "$SLEEP_SECONDS"
done

echo "[wait] download appears complete at $(date -Is)"
du -sh "$DATA_ROOT" || true
parquet_count="$(find "$DATA_ROOT" -name '*.parquet' | wc -l | tr -d ' ')"
echo "[wait] parquet_count=$parquet_count expected_min=$EXPECTED_MIN_PARQUETS"
if [[ "$parquet_count" -lt "$EXPECTED_MIN_PARQUETS" ]]; then
  echo "[wait] refusing to tokenize: parquet_count is below expected minimum" >&2
  exit 3
fi

cd "$REPO_ROOT"

export REPO_ROOT
export IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
export SOURCE_DATA_DIRS="${SOURCE_DATA_DIRS:-$DATA_ROOT/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:$DATA_ROOT/train/logic_topup_proof_pile_17b_v1_20260701}"
export WORK_DIR="${WORK_DIR:-$DATA_ROOT/megatron_bridge/pretrain_8192_glm52_direct_v1}"
export TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}"
export OUTPUT_PREFIX_PREFIX="${OUTPUT_PREFIX_PREFIX:-laoyao_2b_moe_glm52_8192}"
export TEXT_KEY="${TEXT_KEY:-text}"
export PARALLEL_FILES="${PARALLEL_FILES:-4}"
export WORKERS_PER_FILE="${WORKERS_PER_FILE:-8}"
export BATCH_SIZE="${BATCH_SIZE:-8192}"
export CHUNKSIZE="${CHUNKSIZE:-128}"
export MAX_FILES="${MAX_FILES:-0}"
export MAX_DOCS="${MAX_DOCS:-0}"
export MAX_SEQ_LEN="${MAX_SEQ_LEN:-8192}"
export MIN_FREE_GB="${MIN_FREE_GB:-800}"
export OVERWRITE="${OVERWRITE:-0}"

echo "[tokenize] started at $(date -Is)"
echo "[tokenize] work_dir=$WORK_DIR"
echo "[tokenize] max_seq_len=$MAX_SEQ_LEN parallel_files=$PARALLEL_FILES workers_per_file=$WORKERS_PER_FILE"

bash scripts/preprocess_megatron_bridge_pretrain_direct.sh

echo "[tokenize] finished at $(date -Is)"
du -sh "$WORK_DIR" || true
find "$WORK_DIR" -maxdepth 1 -type f | wc -l || true
ls -lh "$WORK_DIR/manifest.json" || true
