#!/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}"

TARGET_TOKENS="${TARGET_TOKENS:-200000000000}"
WORKERS_NOTE="${WORKERS_NOTE:-single-process streaming builder; parallelize by splitting inputs if needed}"

echo "[rebalanced-data] root=$ROOT_DIR"
echo "[rebalanced-data] default_target_tokens=$TARGET_TOKENS (CLI --target-tokens overrides this)"
echo "[rebalanced-data] note=$WORKERS_NOTE"

python dataset/pretrain/scripts/build_rebalanced_pretrain_dataset.py \
  --target-tokens "$TARGET_TOKENS" \
  --output-dir "${OUTPUT_DIR:-data/train/pretrain_rebalanced_knowledge80_logic10_math5_code5_v1_20260630}" \
  --seed "${SEED:-20260630}" \
  --output-shard-rows "${OUTPUT_SHARD_ROWS:-100000}" \
  --progress-every-files "${PROGRESS_EVERY_FILES:-25}" \
  "$@"
