#!/usr/bin/env bash
# Параллельная загрузка шардов одного чекпоинта на Hugging Face.
# Имеет смысл при очень хорошем uplink; на обычном канале прироста почти нет.
#
# Примеры:
#   HF_REPO_ID=user/model ./scripts/push_checkpoint_parallel.sh runs/checkpoint-5000-sharded fanta/fanta-1b/checkpoint-5000
#   JOBS=4 HF_REPO_ID=user/model ./scripts/push_checkpoint_parallel.sh runs/checkpoint-5000-sharded fanta/fanta-1b/checkpoint-5000
#
# Ускорение Rust-клиентом (опционально, широкий канал):
#   pip install 'huggingface_hub[hf_transfer]'
#   export HF_HUB_ENABLE_HF_TRANSFER=1
#
set -euo pipefail

REPO="${HF_REPO_ID:?Задайте HF_REPO_ID=user/model}"
JOBS="${JOBS:-4}"

if [[ $# -lt 1 ]]; then
  echo "Использование: HF_REPO_ID=user/model $0 <local_dir> [path_in_repo]" >&2
  exit 1
fi

LOCAL="$1"
REMOTE="${2:-$(basename "$LOCAL")}"

if [[ ! -d "$LOCAL" ]]; then
  echo "Нет каталога: $LOCAL" >&2
  exit 1
fi

if ! command -v hf >/dev/null 2>&1; then
  echo "hf CLI не найден" >&2
  exit 1
fi

echo "Repo: $REPO"
echo "Local: $LOCAL"
echo "Remote: $REMOTE"
echo "JOBS=$JOBS  HF_HUB_ENABLE_HF_TRANSFER=${HF_HUB_ENABLE_HF_TRANSFER:-0}"
echo ""

export REPO REMOTE LOCAL
upload_one() {
  local file="$1"
  local rel
  rel="$(basename "$file")"
  echo ">> $rel"
  hf upload --repo-type model "$REPO" "$file" "${REMOTE}/${rel}" >/dev/null
  echo "<< $rel OK"
}
export -f upload_one

# Соберём список файлов (шарды, index, config, и всё прочее)
mapfile -d '' FILES < <(find "$LOCAL" -maxdepth 1 -type f -print0)
if (( ${#FILES[@]} == 0 )); then
  echo "В $LOCAL нет файлов" >&2
  exit 1
fi

printf '%s\n' "${FILES[@]}" | xargs -I{} -P "$JOBS" bash -c 'upload_one "$@"' _ {}

echo ""
echo "Готово: $((${#FILES[@]})) файлов → ${REPO}:${REMOTE}"