#!/usr/bin/env python3
"""Загрузить локальный датасет в приватный репозиторий HF (большие объёмы, с возобновлением).

Зачем по коммиту на под-датасет: 37 ГБ и 165 файлов одним коммитом — это один длинный
запрос, который падает на таймауте и начинает всё заново. Здесь каждый `sub_*` уходит своим
коммитом через `run_as_future` (4 параллельно), а повторный запуск пропускает уже
залитые файлы — Xet (hf-xet) видит их на сервере и не передаёт заново.

Приватность: `--private` (по умолчанию True) — репозиторий создаётся приватным. ВАЖНО:
если репозиторий уже существует и он публичный, этот флаг его НЕ сделает приватным —
проверьте и пересоздайте вручную.

Запуск:
    python3 src/upload_hf.py --src src/work/ds_testg_filtered --repo recoilme/ds_testg_filtered
    nohup python3 src/upload_hf.py ... > /tmp/upload.log 2>&1 &
"""
import argparse
import os
import sys


def main():
    p = argparse.ArgumentParser(description="Залить датасет в HF как приватный")
    p.add_argument("--src", required=True, help="локальная папка с датасетом")
    p.add_argument("--repo", required=True, help="репозиторий, напр. recoilme/ds_testg_filtered")
    p.add_argument("--private", action="store_true", default=True)
    p.add_argument("--workers", type=int, default=8, help="параллельных загрузок")
    args = p.parse_args()

    from huggingface_hub import HfApi
    if not os.path.isdir(args.src):
        raise SystemExit(f"нет папки {args.src}")
    n_files = sum(len(f) for _, _, f in os.walk(args.src))
    size = sum(os.path.getsize(os.path.join(r, f))
               for r, _, fs in os.walk(args.src) for f in fs)
    print(f"заливаю {args.src}: {n_files} файлов, {size / 2**30:.1f} ГиБ -> {args.repo}",
          flush=True)

    api = HfApi()
    api.create_repo(repo_id=args.repo, repo_type="dataset", private=args.private,
                    exist_ok=True)
    print(f"репозиторий готов: https://huggingface.co/datasets/{args.repo}", flush=True)

    # README, чтобы у датасета было описание: как собран и чем отфильтрован
    readme = os.path.join(args.src, "README.md")
    if not os.path.exists(readme):
        with open(readme, "w", encoding="utf-8") as f:
            f.write(_readme(args.repo))
        print("добавлен README.md", flush=True)

    # по коммиту на под-датасет: обрыв не обнуляет прогресс
    subs = sorted(os.path.join(r, d) for r, ds_, _ in os.walk(args.src)
                  for d in ds_ if d.startswith("sub_"))
    groups = sorted({os.path.dirname(s) for s in subs})
    units = groups if groups else [args.src]
    print(f"коммитов будет {len(units)}", flush=True)
    futures = []
    for i, path in enumerate(units, 1):
        rel = os.path.relpath(path, args.src)
        print(f"[{i}/{len(units)}] -> {rel}", flush=True)
        futures.append((rel, api.upload_folder(
            folder_path=path, path_in_repo=rel if groups else "",
            repo_id=args.repo, repo_type="dataset", run_as_future=True)))
        if len(futures) >= args.workers:  # дождаться пачки, чтобы не копить запросы
            for name, f in futures:
                f.result()
                print(f"  готово: {name}", flush=True)
            futures = []
    for name, f in futures:
        f.result()
        print(f"  готово: {name}", flush=True)
    print(f"ГОТОВО: https://huggingface.co/datasets/{args.repo}", flush=True)


def _readme(repo):
    return f"""---
license: other
task_categories:
- text-to-image
tags:
- moss-image
- diffusers
- latents
- anime
private: true
---

# {repo.split('/')[-1]}

Закешированные латенты VAE и капшены для файнтюна Mossling (MOSS-Image, 30-блочный DiT).

Собран `src/make_dataset.py` из отфильтрованного набора, записан скриптом `src/filter_dataset.py`.

## Формат строки

| поле | тип | что это |
|---|---|---|
| `vae` | float16 `[128, H/16, W/16]` | латент VAE, упакованный 2x2 и bn-нормированный |
| `text` | str | капшен: теги gelbooru через запятую, без `*tags*` |
| `width`, `height` | int | размер кадра, из которого сделан латент |
| `aspect` | str | `1:1`, `3:4`, `4:3` ... |
| `source` | str | путь исходного файла |

Колонок `text_embeds`/`text_mask` нет: условие считается на лету энкодером Mossling
(`src/train.py --te`).

## Фильтрация

Из исходного набора отброшено 43.5% строк по тегам капшена: листы персонажа и ракурсы,
надписи и речь, монохром и скетч, несколько персонажей, низкое качество, обложки,
лишние конечности, обрезанный кадр, сцены без людей, капшены длиннее 30 тегов.

## Лицензия

Изображения — из открытого набора с тегами; права принадлежат их авторам.
Датасет собран для обучения и не является перераспространением оригиналов в исходном виде
(хранятся только латенты).
"""


if __name__ == "__main__":
    main()