#!/usr/bin/env python3
"""PUBLIC release: code + Argoverse-2-only data + checkpoints. Waymo is EXCLUDED
(its license forbids redistribution). To guarantee no Waymo is ever public (incl.
git/LFS history), the existing private repo is DELETED and recreated fresh as
public, then only AV2 clips are uploaded. Token from $HF_TOKEN."""
import os, glob, tarfile
from huggingface_hub import HfApi, create_repo, delete_repo

TOKEN = os.environ["HF_TOKEN"]
REPO = os.environ.get("HF_REPO", "ChenmingWu/mapgs-maptokengs")
ROOT = "/mnt/william"
STAGE = "/mnt/william/_upload"
SHARD_GB = 18.0
api = HfApi(token=TOKEN)


def log(*a): print(*a, flush=True)


def dir_size(d):
    s = 0
    for r, _, fs in os.walk(d):
        for f in fs:
            try: s += os.path.getsize(os.path.join(r, f))
            except OSError: pass
    return s


def main():
    # nuke the private repo (it contained Waymo) and recreate fresh + PUBLIC, so no
    # Waymo remains anywhere (no leftover LFS objects in history).
    try:
        delete_repo(REPO, repo_type="model", token=TOKEN); log("deleted old private repo")
    except Exception as e:
        log("delete (ok if absent):", repr(e)[:120])
    create_repo(REPO, private=False, repo_type="model", exist_ok=True, token=TOKEN)
    log(f"recreated PUBLIC repo: {REPO}")

    api.upload_file(path_or_fileobj=f"{ROOT}/README_HF.md", path_in_repo="README.md",
                    repo_id=REPO, repo_type="model")
    api.upload_folder(folder_path=ROOT, repo_id=REPO, repo_type="model", path_in_repo=".",
                      allow_patterns=["mapgs/**", "maptokengs/**", "scripts/**", "_tokengs_repo/**"],
                      ignore_patterns=["**/__pycache__/**", "**/*.pyc", "**/.git/**", "**/*.tar",
                                       "_tokengs_repo/assets/**", "**/*.gif"])
    log("uploaded README + code")

    for p in [q for q in sorted(glob.glob(f"{ROOT}/runs/*.safetensors"))
              if "maptokengs_fixed" not in os.path.basename(q)]:
        api.upload_file(path_or_fileobj=p, path_in_repo=f"runs/{os.path.basename(p)}",
                        repo_id=REPO, repo_type="model")
        log("uploaded ckpt", os.path.basename(p))

    # AV2 ONLY (no Waymo)
    os.makedirs(STAGE, exist_ok=True)
    clip_dirs = sorted(d for d in glob.glob(f"{ROOT}/data/unified/av2/*/*") if os.path.isdir(d))
    assert all("/waymo/" not in d for d in clip_dirs), "Waymo path leaked into AV2 shard set!"
    log(f"AV2 clip dirs: {len(clip_dirs)}")
    shards, cur, sz, cap = [], [], 0, SHARD_GB * 1e9
    for d in clip_dirs:
        s = dir_size(d)
        if cur and sz + s > cap: shards.append(cur); cur, sz = [], 0
        cur.append(d); sz += s
    if cur: shards.append(cur)
    log(f"{len(shards)} AV2 shards")
    for i, dirs in enumerate(shards):
        tarp = f"{STAGE}/av2_shard_{i:03d}.tar"
        with tarfile.open(tarp, "w") as tf:
            for d in dirs: tf.add(d, arcname=os.path.relpath(d, ROOT))
        gb = os.path.getsize(tarp) / 1e9
        api.upload_file(path_or_fileobj=tarp, path_in_repo=f"data/av2_shard_{i:03d}.tar",
                        repo_id=REPO, repo_type="model")
        os.remove(tarp)
        log(f"uploaded av2_shard_{i:03d}.tar ({gb:.1f} GB, {len(dirs)} clips) [{i+1}/{len(shards)}]")
    log("=== PUBLIC UPLOAD COMPLETE ===")


if __name__ == "__main__":
    main()