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

REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
SOURCE_DATA_DIRS="${SOURCE_DATA_DIRS:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/logic_topup_proof_pile_17b_v1_20260701}"
WORK_DIR="${WORK_DIR:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_direct_smoke_v1}"
TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}"
OUTPUT_PREFIX_PREFIX="${OUTPUT_PREFIX_PREFIX:-laoyao_2b_moe_8192_direct}"
TEXT_KEY="${TEXT_KEY:-text}"
PARALLEL_FILES="${PARALLEL_FILES:-1}"
WORKERS_PER_FILE="${WORKERS_PER_FILE:-8}"
BATCH_SIZE="${BATCH_SIZE:-8192}"
CHUNKSIZE="${CHUNKSIZE:-128}"
MAX_FILES="${MAX_FILES:-1}"
MAX_DOCS="${MAX_DOCS:-100000}"
MAX_SEQ_LEN="${MAX_SEQ_LEN:-65536}"
MIN_FREE_GB="${MIN_FREE_GB:-20}"
OVERWRITE="${OVERWRITE:-1}"

mkdir -p "$WORK_DIR"

available_gb="$(df -BG "$WORK_DIR" | awk 'NR==2 {gsub("G", "", $4); print $4}')"
if [[ "$MAX_FILES" == "0" && "$MAX_DOCS" == "0" && "$available_gb" -lt "$MIN_FREE_GB" ]]; then
  cat >&2 <<EOF
Refusing full direct preprocessing: only ${available_gb}GB free at $WORK_DIR.
Direct Megatron indexed output for 200B GLM5.2 tokens is expected to require
roughly 800GB plus index/cache overhead. Increase free space or lower MAX_FILES/MAX_DOCS.
EOF
  exit 2
fi

input_args=()
IFS=':' read -r -a source_dirs <<< "$SOURCE_DATA_DIRS"
for source_dir in "${source_dirs[@]}"; do
  input_args+=(--input "${source_dir}/*.parquet")
done

max_file_args=()
if [[ "$MAX_FILES" != "0" ]]; then
  max_file_args+=(--max-files "$MAX_FILES")
fi

overwrite_args=()
if [[ "$OVERWRITE" == "1" ]]; then
  overwrite_args+=(--overwrite)
fi

DOCKER_MOUNTS=()
if [[ -d /mnt/beegfs ]]; then
  DOCKER_MOUNTS+=(-v /mnt/beegfs:/mnt/beegfs)
fi
if [[ -d /ssd ]]; then
  DOCKER_MOUNTS+=(-v /ssd:/ssd)
fi

docker run --rm --ipc=host --network=host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  "${DOCKER_MOUNTS[@]}" \
  -w "$REPO_ROOT" \
  "$IMAGE" \
  python3 dataset/pretrain/scripts/convert_pretrain_parquet_to_megatron.py \
    "${input_args[@]}" \
    --output-dir "$WORK_DIR" \
    --manifest "$WORK_DIR/manifest.json" \
    --megatron-dir /opt/Megatron-Bridge/3rdparty/Megatron-LM \
    --tokenizer-type HuggingFaceTokenizer \
    --tokenizer-model "$TOKENIZER_MODEL" \
    --text-key "$TEXT_KEY" \
    --output-prefix-prefix "$OUTPUT_PREFIX_PREFIX" \
    --parallel-files "$PARALLEL_FILES" \
    --workers-per-file "$WORKERS_PER_FILE" \
    --batch-size "$BATCH_SIZE" \
    --chunksize "$CHUNKSIZE" \
    --max-docs "$MAX_DOCS" \
    --max-seq-len "$MAX_SEQ_LEN" \
    "${max_file_args[@]}" \
    "${overwrite_args[@]}"

echo "Direct Megatron indexed dataset manifest: $WORK_DIR/manifest.json"
