"""
Upload all .pth checkpoints to a Hugging Face Model repo.

The Model repo holds the weights so the Space stays under its 1 GB cap.
The Space's app.py downloads from this repo via hf_hub_download.

Usage:
    huggingface-cli login                  # one-time, paste your write token
    python tools/upload_weights.py phiniqs/seg-models-weights

Folder layout in the Model repo:
    baseline/{model}_{share}_{kind}.pth        from clean_data_scaling_study
    grid/{cfg_id}_best.pth                     from finetune_grid_search
"""
import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
BASELINE_DIR = ROOT / "experiments" / "clean_data_scaling_study" / "checkpoints"
GRID_DIR = ROOT / "experiments" / "finetune_grid_search" / "checkpoints"


def main():
    p = argparse.ArgumentParser()
    p.add_argument("repo_id", help="HF model repo id, e.g. phiniqs/seg-models-weights")
    p.add_argument("--private", action="store_true", help="create as private repo")
    p.add_argument("--dry-run", action="store_true", help="show plan, don't upload")
    args = p.parse_args()

    try:
        from huggingface_hub import HfApi, create_repo
    except ImportError:
        print("ERROR: install huggingface_hub first:")
        print("  pip install --user huggingface_hub")
        sys.exit(1)

    baseline_files = sorted(BASELINE_DIR.glob("*.pth"))
    grid_files = sorted(GRID_DIR.glob("*.pth"))

    total_bytes = sum(f.stat().st_size for f in baseline_files + grid_files)

    print(f"baseline checkpoints : {len(baseline_files):>3}  "
          f"({sum(f.stat().st_size for f in baseline_files) / 1e6:.1f} MB)")
    print(f"grid checkpoints     : {len(grid_files):>3}  "
          f"({sum(f.stat().st_size for f in grid_files) / 1e6:.1f} MB)")
    print(f"total                : {len(baseline_files) + len(grid_files)} files, "
          f"{total_bytes / 1e9:.2f} GB")
    print()

    if args.dry_run:
        print("[dry-run] would create Model repo:", args.repo_id)
        print(f"[dry-run] would upload to baseline/  ({len(baseline_files)} files)")
        for f in baseline_files[:3]:
            print(f"    {f.name}")
        if len(baseline_files) > 3:
            print(f"    ... +{len(baseline_files) - 3} more")
        print(f"[dry-run] would upload to grid/  ({len(grid_files)} files)")
        for f in grid_files[:3]:
            print(f"    {f.name}")
        if len(grid_files) > 3:
            print(f"    ... +{len(grid_files) - 3} more")
        return

    api = HfApi()

    print(f"creating model repo: {args.repo_id} (private={args.private})")
    create_repo(
        repo_id=args.repo_id,
        repo_type="model",
        private=args.private,
        exist_ok=True,
    )

    print(f"uploading baseline/  ({len(baseline_files)} files, "
          f"{sum(f.stat().st_size for f in baseline_files) / 1e6:.1f} MB)…")
    api.upload_folder(
        folder_path=str(BASELINE_DIR),
        path_in_repo="baseline",
        repo_id=args.repo_id,
        repo_type="model",
        commit_message="Add baseline checkpoints (4 architectures × 3 shares)",
        allow_patterns=["*.pth"],
    )

    print(f"uploading grid/  ({len(grid_files)} files, "
          f"{sum(f.stat().st_size for f in grid_files) / 1e6:.1f} MB)…")
    api.upload_folder(
        folder_path=str(GRID_DIR),
        path_in_repo="grid",
        repo_id=args.repo_id,
        repo_type="model",
        commit_message="Add fine-tune grid checkpoints (54 configs)",
        allow_patterns=["*.pth"],
    )

    print()
    print(f"done. weights live at: https://huggingface.co/{args.repo_id}")
    print(f"the Space's app.py reads from this repo via HF_WEIGHTS_REPO env var "
          f"(default: phiniqs/seg-models-weights).")


if __name__ == "__main__":
    main()