#!/usr/bin/env python
"""upload_task.py — upload one trained task's pretrained_model folder to a PUBLIC
model repo as a subfolder, then verify what landed.

Usage:
  python upload_task.py <PRETRAINED_MODEL_DIR> <PATH_IN_REPO> <REPO_ID> <EXPECT_ROOT_SUBSTR>

Uploads the WHOLE folder (handles sharded safetensors). Creates repo (private=False)
if missing.
"""
import json
import os
import sys

from huggingface_hub import HfApi, create_repo, upload_folder


def main():
    src, path_in_repo, repo_id, expect = sys.argv[1:5]
    if not os.path.isdir(src):
        print(f"[upload] SRC MISSING: {src} -> SKIP")
        sys.exit(1)

    create_repo(repo_id, repo_type="model", private=False, exist_ok=True)
    print(f"[upload] {src} -> {repo_id}:{path_in_repo} (public)")
    upload_folder(
        folder_path=src,
        path_in_repo=path_in_repo,
        repo_id=repo_id,
        repo_type="model",
        commit_message=f"Add {path_in_repo} GR00T N1.7 EE checkpoint",
    )

    print(f"[verify] files under {repo_id}:{path_in_repo}")
    api = HfApi()
    files = [f for f in api.list_repo_files(repo_id, repo_type="model") if f.startswith(path_in_repo + "/")]
    total = 0
    infos = api.model_info(repo_id, files_metadata=True).siblings
    size_by = {s.rfilename: (s.size or 0) for s in infos}
    for f in sorted(files):
        sz = size_by.get(f, 0)
        total += sz
        print(f"    {f}  ({sz/1e6:.1f} MB)")
    print(f"[verify] safetensors total ~ {total/1e9:.2f} GB")

    # local content checks
    tc = os.path.join(src, "train_config.json")
    cfg = os.path.join(src, "config.json")
    if os.path.exists(tc):
        t = json.load(open(tc))
        d = t.get("dataset", {})
        root = str(d.get("root")); rid = str(d.get("repo_id"))
        ok = expect in root or expect in rid
        print(f"[verify] train_config dataset.root={root} repo_id={rid} -> "
              f"{'OK points to ' + expect if ok else 'WARN does not mention ' + expect}")
    else:
        print("[verify] train_config.json MISSING locally")
    if os.path.exists(cfg):
        c = json.load(open(cfg))
        print(f"[verify] config.json: model_version={c.get('model_version')} "
              f"base_model_path={c.get('base_model_path')} embodiment_tag={c.get('embodiment_tag')}")
    has_index = any(f.endswith("model.safetensors.index.json") for f in files)
    print(f"[verify] sharded index present: {has_index}")
    print(f"[verify] DONE {path_in_repo}")


if __name__ == "__main__":
    main()