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

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

export http_proxy="${http_proxy:-http://10.20.34.2:3128}"
export https_proxy="${https_proxy:-http://10.20.34.2:3128}"
export HTTP_PROXY="${HTTP_PROXY:-$http_proxy}"
export HTTPS_PROXY="${HTTPS_PROXY:-$https_proxy}"
export no_proxy="${no_proxy:-localhost,127.0.0.1,10.20.0.0/16,internal-cr,internal-cr:3443,10.20.34.1:6100}"
export NO_PROXY="${NO_PROXY:-$no_proxy}"

DATA_ROOT="${DATA_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset}"
OUTPUT_DIR="${OUTPUT_DIR:-$DATA_ROOT/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701}"
TARGET_TOKENS="${TARGET_TOKENS:-200000000000}"
EXPECTED_ZH_SHARDS="${EXPECTED_ZH_SHARDS:-32}"
POLL_SECONDS="${POLL_SECONDS:-120}"

ENGLISH_WEB_INPUT="${ENGLISH_WEB_INPUT:-$DATA_ROOT/ms_ultra_fineweb_en/**/*.parquet}"
ENGLISH_EDU_INPUT="${ENGLISH_EDU_INPUT:-$DATA_ROOT/ms_fineweb_edu/**/*.parquet}"
CHINESE_CLEAN_INPUT="${CHINESE_CLEAN_INPUT:-$DATA_ROOT/ms_ultra_fineweb_zh/**/*.parquet}"
SCIENCE_INPUT="${SCIENCE_INPUT:-$DATA_ROOT/science_*/*/*.parquet}"
JIAYI_SCORED_INPUT="${JIAYI_SCORED_INPUT:-/mnt/beegfs/cjy/cjy-training/laoyao_model/data/scored/edu_fineweb_math_knowledge_cross_near_h1_v1_20260620/rank_*/*.jsonl}"
CODE_INPUT="${CODE_INPUT:-/mnt/beegfs/cjy/cjy-training/laoyao_model/data/dedup/pretrain_python_clean_code_multilang_cross_near_h1_min256_v1_20260623/*.parquet}"

echo "[wait-build] root=$ROOT_DIR"
echo "[wait-build] data_root=$DATA_ROOT"
echo "[wait-build] output_dir=$OUTPUT_DIR"
echo "[wait-build] target_tokens=$TARGET_TOKENS"
echo "[wait-build] waiting for zh shards: $EXPECTED_ZH_SHARDS"

while true; do
  zh_count="$(find "$DATA_ROOT/ms_ultra_fineweb_zh" -type f -name '*.parquet' 2>/dev/null | wc -l | tr -d ' ')"
  echo "[wait-build] $(date -u +%Y-%m-%d_%H:%M:%S_UTC) zh_shards=$zh_count/$EXPECTED_ZH_SHARDS"
  if [[ "$zh_count" -ge "$EXPECTED_ZH_SHARDS" ]]; then
    break
  fi
  sleep "$POLL_SECONDS"
done

echo "[wait-build] running full dry-run"
bash scripts/data/start_rebalanced_pretrain_dataset.sh \
  --dry-run \
  --english-web-input "$ENGLISH_WEB_INPUT" \
  --english-edu-input "$ENGLISH_EDU_INPUT" \
  --chinese-clean-input "$CHINESE_CLEAN_INPUT" \
  --science-input "$SCIENCE_INPUT" \
  --jiayi-scored-input "$JIAYI_SCORED_INPUT" \
  --code-input "$CODE_INPUT"

echo "[wait-build] starting full materialization"
OUTPUT_DIR="$OUTPUT_DIR" TARGET_TOKENS="$TARGET_TOKENS" \
bash scripts/data/start_rebalanced_pretrain_dataset.sh \
  --overwrite \
  --english-web-input "$ENGLISH_WEB_INPUT" \
  --english-edu-input "$ENGLISH_EDU_INPUT" \
  --chinese-clean-input "$CHINESE_CLEAN_INPUT" \
  --science-input "$SCIENCE_INPUT" \
  --jiayi-scored-input "$JIAYI_SCORED_INPUT" \
  --code-input "$CODE_INPUT"

echo "[wait-build] done output_dir=$OUTPUT_DIR"
