#!/usr/bin/env python3
"""Push the research code to the HF Space. Never uploads data, checkpoints or secrets."""
import os
import sys
from pathlib import Path

from huggingface_hub import HfApi

ROOT = Path(__file__).parent.parent.resolve()
REPO = os.environ.get("HF_SPACE", "cazyundee/training")

INCLUDE = ["app.py", "train_selfplay.py", "inspect_trace.py", "requirements.txt",
           "README.md", "tinychess", "tests", "experiments", "configs", "scripts",
           "docs"]
SKIP_SUFFIX = (".pt", ".pkl", ".pgn", ".zst", ".parquet", ".log", ".pyc")
SKIP_PARTS = {"__pycache__", ".git", "runs", "data", "bin", ".venv"}


def collect():
    files = []
    for item in INCLUDE:
        p = ROOT / item
        if p.is_file():
            files.append(p)
        elif p.is_dir():
            for f in p.rglob("*"):
                if f.is_file() and not any(s in f.parts for s in SKIP_PARTS) \
                        and not f.name.endswith(SKIP_SUFFIX):
                    files.append(f)
    return files


def main():
    token = os.environ.get("HF_TOKEN")
    if not token:
        sys.exit("HF_TOKEN not set")
    api = HfApi(token=token)
    files = collect()
    total = sum(f.stat().st_size for f in files)
    print(f"uploading {len(files)} files ({total/1024:.0f} KB) -> {REPO}")
    ops = []
    from huggingface_hub import CommitOperationAdd
    for f in files:
        ops.append(CommitOperationAdd(path_in_repo=str(f.relative_to(ROOT)),
                                      path_or_fileobj=str(f)))
    api.create_commit(repo_id=REPO, repo_type="space", operations=ops,
                      commit_message=sys.argv[1] if len(sys.argv) > 1
                      else "update tinychess research code")
    print(f"done: https://huggingface.co/spaces/{REPO}")


if __name__ == "__main__":
    main()