from __future__ import annotations

import argparse
from pathlib import Path

from huggingface_hub import HfApi, create_repo


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo-id", default=None, help="Target model repo, e.g. user/stochastic-tbrm-agent-benchmarks")
    parser.add_argument("--private", action="store_true", help="Create/upload to a private repo")
    parser.add_argument("--commit-message", default="Upload Stochastic-TBRM and DMPO benchmark checkpoints")
    args = parser.parse_args()

    api = HfApi()
    repo_id = args.repo_id
    if repo_id is None:
        user = api.whoami()["name"]
        repo_id = f"{user}/stochastic-tbrm-agent-benchmarks"

    create_repo(repo_id=repo_id, repo_type="model", private=args.private, exist_ok=True)
    api.upload_folder(
        repo_id=repo_id,
        repo_type="model",
        folder_path=str(Path(__file__).resolve().parent),
        path_in_repo=".",
        commit_message=args.commit_message,
    )
    print(f"Uploaded to https://huggingface.co/{repo_id}")


if __name__ == "__main__":
    main()