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

# Prepare g0050 for Laoyao 2B MoE training and download the pre-tokenization
# parquet dataset from ModelScope. Keep credentials out of this file:
#
#   MODELSCOPE_API_TOKEN=ms-... bash scripts/g0050_download_and_setup_from_modelscope.sh
#
# Optional private git URL:
#
#   GIT_REPO_URL=https://user:token@git.deeepseek.net/yi_lu/laoyao_2b_moe.git ...

REPO_ROOT="${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe}"
DATA_ROOT="${DATA_ROOT:-/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset}"
DATASET_ID="${DATASET_ID:-eigentom/laoyao_2b_moe_pretrain_parquet_20260702}"
GIT_REPO_URL="${GIT_REPO_URL:-https://git.deeepseek.net/yi_lu/laoyao_2b_moe.git}"
GIT_BRANCH="${GIT_BRANCH:-main}"

NEMO_IMAGE="${NEMO_IMAGE:-nvcr.io/nvidia/nemo:26.06}"
MODELSCOPE_IMAGE="${MODELSCOPE_IMAGE:-modelscope-registry.cn-hangzhou.cr.aliyuncs.com/modelscope-repo/modelscope:ubuntu22.04-cuda13.0.3-py312-torch2.11.0-vllm0.23.0-modelscope1.37.1-swift4.3.1}"
MODELSCOPE_ENDPOINT="${MODELSCOPE_ENDPOINT:-https://www.modelscope.cn}"
MAX_WORKERS="${MAX_WORKERS:-8}"

HTTP_PROXY_DEFAULT="${HTTP_PROXY_DEFAULT:-http://100.72.0.101:8888}"
HTTPS_PROXY_DEFAULT="${HTTPS_PROXY_DEFAULT:-http://100.72.0.101:8888}"
export http_proxy="${http_proxy:-$HTTP_PROXY_DEFAULT}"
export https_proxy="${https_proxy:-$HTTPS_PROXY_DEFAULT}"
export HTTP_PROXY="${HTTP_PROXY:-$http_proxy}"
export HTTPS_PROXY="${HTTPS_PROXY:-$https_proxy}"

if [[ -z "${MODELSCOPE_API_TOKEN:-}" ]]; then
  echo "ERROR: set MODELSCOPE_API_TOKEN before running this script." >&2
  exit 2
fi

mkdir -p "$(dirname "$REPO_ROOT")" "$DATA_ROOT"

echo "[setup] repo_root=$REPO_ROOT"
if [[ -d "$REPO_ROOT/.git" ]]; then
  git -C "$REPO_ROOT" fetch origin "$GIT_BRANCH"
  git -C "$REPO_ROOT" checkout "$GIT_BRANCH"
  git -C "$REPO_ROOT" pull --ff-only origin "$GIT_BRANCH"
elif [[ -d "$REPO_ROOT" ]] && [[ -z "$(find "$REPO_ROOT" -mindepth 1 -maxdepth 1 -print -quit)" ]]; then
  rmdir "$REPO_ROOT"
  git clone --branch "$GIT_BRANCH" "$GIT_REPO_URL" "$REPO_ROOT"
else
  git clone --branch "$GIT_BRANCH" "$GIT_REPO_URL" "$REPO_ROOT"
fi

echo "[setup] checking docker images"
if ! docker image inspect "$NEMO_IMAGE" >/dev/null 2>&1; then
  docker pull "$NEMO_IMAGE"
fi
if ! docker image inspect "$MODELSCOPE_IMAGE" >/dev/null 2>&1; then
  docker pull "$MODELSCOPE_IMAGE"
fi

echo "[download] dataset=$DATASET_ID"
echo "[download] data_root=$DATA_ROOT"
docker run --rm \
  --network=host \
  -e http_proxy="$http_proxy" \
  -e https_proxy="$https_proxy" \
  -e HTTP_PROXY="$HTTP_PROXY" \
  -e HTTPS_PROXY="$HTTPS_PROXY" \
  -e MODELSCOPE_API_TOKEN="$MODELSCOPE_API_TOKEN" \
  -v "$DATA_ROOT:$DATA_ROOT" \
  "$MODELSCOPE_IMAGE" \
  modelscope download \
    --dataset "$DATASET_ID" \
    --token "$MODELSCOPE_API_TOKEN" \
    --endpoint "$MODELSCOPE_ENDPOINT" \
    --local_dir "$DATA_ROOT" \
    --max-workers "$MAX_WORKERS"

echo "[verify] downloaded tree"
find "$DATA_ROOT" -maxdepth 3 -type f | sed -n '1,20p'
du -sh "$DATA_ROOT" || true

cat <<EOF

Done.

Expected source parquet roots:
  $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

Next steps:
  cd $REPO_ROOT
  bash scripts/preprocess_megatron_bridge_pretrain_direct.sh
  bash scripts/train_megatron_bridge_2b_moe.sh
EOF
