#!/usr/bin/env python3
"""Upload Waymo-derived unified clips to a PRIVATE HF repo (owner's licensed-use
backup; NOT redistribution). Waymo-only; AV2/code/checkpoints live in the public
companion repo. Token from $HF_TOKEN."""
import os, glob, tarfile
from huggingface_hub import HfApi, create_repo

TOKEN = os.environ["HF_TOKEN"]
REPO = os.environ.get("HF_REPO", "ChenmingWu/mapgs-waymo")
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():
    create_repo(REPO, private=True, repo_type="model", exist_ok=True, token=TOKEN)
    log(f"PRIVATE repo ready: {REPO}")
    api.upload_file(path_or_fileobj=f"{ROOT}/README_WAYMO.md", path_in_repo="README.md",
                    repo_id=REPO, repo_type="model")
    log("uploaded README")

    os.makedirs(STAGE, exist_ok=True)
    clip_dirs = sorted(d for d in glob.glob(f"{ROOT}/data/unified/waymo/*/*") if os.path.isdir(d))
    assert clip_dirs and all("/av2/" not in d for d in clip_dirs), "non-Waymo path in shard set!"
    log(f"Waymo 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)} Waymo shards")
    for i, dirs in enumerate(shards):
        tarp = f"{STAGE}/waymo_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/waymo_shard_{i:03d}.tar",
                        repo_id=REPO, repo_type="model")
        os.remove(tarp)
        log(f"uploaded waymo_shard_{i:03d}.tar ({gb:.1f} GB, {len(dirs)} clips) [{i+1}/{len(shards)}]")
    log("=== WAYMO PRIVATE UPLOAD COMPLETE ===")


if __name__ == "__main__":
    main()