from __future__ import annotations

import argparse
import os
from pathlib import Path


def _resolve_token() -> str | None:
    """
    Prefer environment variables to avoid accidentally writing tokens into shell history.
    """

    for key in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN"):
        token = os.environ.get(key)
        if token:
            return token
    return None


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Publish this inference bundle (folder) to Hugging Face Hub as a model repo."
    )
    parser.add_argument(
        "--repo-id",
        type=str,
        required=True,
        help="e.g. USERNAME/judgment-partition-infer",
    )
    parser.add_argument(
        "--skip-create",
        action="store_true",
        help="Skip create_repo (useful when the repo already exists).",
    )
    parser.add_argument(
        "--public",
        action="store_true",
        help="Make the repo public (default: private).",
    )
    parser.add_argument(
        "--commit-message",
        type=str,
        default="Upload inference bundle",
        help="Hub commit message.",
    )
    args = parser.parse_args()

    try:
        from huggingface_hub import HfApi
    except Exception as exc:  # pragma: no cover
        raise SystemExit(
            "Missing dependency: huggingface_hub. Install with:\n"
            "  pip install -r requirements-publish.txt\n"
        ) from exc

    bundle_root = Path(__file__).resolve().parent
    token = _resolve_token()

    api = HfApi(token=token)
    if not args.skip_create:
        api.create_repo(
            repo_id=args.repo_id,
            repo_type="model",
            private=not args.public,
            exist_ok=True,
        )
    api.upload_folder(
        repo_id=args.repo_id,
        repo_type="model",
        folder_path=str(bundle_root),
        path_in_repo="",
        commit_message=args.commit_message,
        ignore_patterns=[
            "**/__pycache__/**",
            "**/*.pyc",
            "assets/best_model.pt",
            "assets/best_model.pt.part-*",
            "output/**",
            ".pytest_cache/**",
            ".mypy_cache/**",
            "dist/**",
            "build/**",
            "*.egg-info/**",
        ],
    )

    print(f"[DONE] https://huggingface.co/{args.repo_id}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())