"""Upload the code and the asset archives to a Hugging Face repo.

Requires a login first (this script will NOT prompt, so it is safe to run unattended):

    hf auth login            # or:  huggingface-cli login

Then:

    python scripts/push_to_hf.py --repo yqi19/mp_yam_code            # dry run, lists what goes
    python scripts/push_to_hf.py --repo yqi19/mp_yam_code --push     # actually upload

Code and assets go up as separate commits so a failed 3 GB asset upload does not lose the code.
"""
import argparse
import os
import sys

ap = argparse.ArgumentParser()
ap.add_argument("--repo", default="yqi19/mp_yam_code")
ap.add_argument("--repo-type", default="model", choices=["model", "dataset"])
ap.add_argument("--stage", default="/home/yu/internship_yu/hf_stage")
ap.add_argument("--push", action="store_true", help="without this, only report what would go")
ap.add_argument("--skip-assets", action="store_true")
a = ap.parse_args()

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# what counts as "the code": everything that is ours, none of the multi-GB outputs
CODE_ALLOW = ["source/**", "scripts/**", "README.md", "pyproject.toml", "setup.py",
              "outputs/frames/overview.png"]
CODE_IGNORE = ["**/__pycache__/**", "**/*.pyc", "outputs/tasks/**", "outputs/sheets/**",
               "**/.git/**", "**/*.mp4"]


def human(n):
    for u in ("B", "KB", "MB", "GB"):
        if n < 1024 or u == "GB":
            return f"{n:.1f} {u}"
        n /= 1024


def main():
    from huggingface_hub import HfApi
    api = HfApi()
    try:
        who = api.whoami()["name"]
    except Exception:
        sys.exit("[hf] not logged in. Run:  hf auth login      (then re-run this script)")
    print(f"[hf] authenticated as {who}, target {a.repo} ({a.repo_type})")

    assets = []
    if not a.skip_assets:
        for fn in ("robotwin_usd.zip", "robotwin_source_meshes.zip"):
            p = os.path.join(a.stage, fn)
            if os.path.exists(p):
                assets.append((p, f"assets/{fn}", os.path.getsize(p)))
            else:
                print(f"[hf] missing {p} -- skipping")

    print(f"[hf] code:   {REPO_ROOT}  (allow={CODE_ALLOW})")
    for _p, dest, size in assets:
        print(f"[hf] asset:  {dest}  {human(size)}")
    if not a.push:
        print("\n[hf] DRY RUN -- nothing uploaded. Re-run with --push to upload.")
        return 0

    api.create_repo(a.repo, repo_type=a.repo_type, exist_ok=True)
    print("[hf] uploading code ...")
    api.upload_folder(folder_path=REPO_ROOT, repo_id=a.repo, repo_type=a.repo_type,
                      allow_patterns=CODE_ALLOW, ignore_patterns=CODE_IGNORE,
                      commit_message="YAM bimanual task suite: env, solvers, tasks, converters")
    for p, dest, size in assets:
        print(f"[hf] uploading {dest} ({human(size)}) -- this is slow, it is LFS ...")
        api.upload_file(path_or_fileobj=p, path_in_repo=dest, repo_id=a.repo,
                        repo_type=a.repo_type,
                        commit_message=f"assets: {os.path.basename(p)}")
    print(f"[hf] done -> https://huggingface.co/{a.repo}")
    return 0


if __name__ == "__main__":
    sys.exit(main())