#!/usr/bin/env python3
"""Upload an export bundle (+ model card) to the Hugging Face Hub.

python3 scripts/push_hf.py --export exports/jev-judge-qwen35-9b --repo autotrust/JEV --dry-run
python3 scripts/push_hf.py --export exports/jev-judge-qwen35-9b --repo autotrust/JEV --private
"""

from __future__ import annotations

import argparse
import os

import yaml


def validate_card(path: str) -> dict:
    text = open(path, encoding="utf-8").read()
    if not text.startswith("---\n"):
        raise SystemExit("README.md has no YAML front matter")
    fm = text.split("---\n", 2)[1]
    meta = yaml.safe_load(fm)
    for k in ("license", "base_model", "library_name", "pipeline_tag", "tags"):
        if k not in meta:
            raise SystemExit(f"model card front matter missing `{k}`")
    return meta


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--export", required=True)
    ap.add_argument("--repo", required=True, help="e.g. autotrust/JEV")
    ap.add_argument("--private", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--message", default="Upload JEV v0.7.0 (Qwen3.5-9B distilled from Jev 1.13)")
    args = ap.parse_args()

    card = os.path.join(args.export, "README.md")
    meta = validate_card(card)
    files = sorted(os.listdir(args.export))
    total = sum(os.path.getsize(os.path.join(args.export, f)) for f in files)
    print(f"repo: {args.repo} ({'private' if args.private else 'public'})")
    print(f"card front matter: license={meta['license']} base_model={meta['base_model']} pipeline_tag={meta['pipeline_tag']} tags={len(meta['tags'])}")
    print(f"files ({total/1e9:.2f} GB):")
    for f in files:
        print(f"  {os.path.getsize(os.path.join(args.export, f))/1e9:8.3f} GB  {f}")
    if args.dry_run:
        print("dry run — nothing uploaded")
        return
    from huggingface_hub import HfApi

    api = HfApi()
    who = api.whoami()
    print("authenticated as", who.get("name"), "orgs:", [o.get("name") for o in who.get("orgs", [])])
    api.create_repo(args.repo, repo_type="model", private=args.private, exist_ok=True)
    url = api.upload_folder(repo_id=args.repo, repo_type="model", folder_path=args.export, commit_message=args.message,
                            ignore_patterns=["*.log", "__pycache__"],
                            delete_patterns=["model-*.safetensors", "model.safetensors.index.json", "code/**"])
    print("uploaded ->", url)
    print(f"https://huggingface.co/{args.repo}")


if __name__ == "__main__":
    main()