from __future__ import annotations

import argparse
from pathlib import Path

from huggingface_hub import HfApi, upload_file


def main() -> None:
    parser = argparse.ArgumentParser(description="Upload custom model checkpoint to Hugging Face model repo.")
    parser.add_argument("--repo-id", required=True, help="e.g., username/pokemon-transfer-resnet18")
    parser.add_argument("--model-path", default="models/custom_resnet18.pth")
    parser.add_argument("--private", action="store_true")
    args = parser.parse_args()

    model_path = Path(args.model_path)
    if not model_path.exists():
        raise FileNotFoundError(f"Model file not found: {model_path}")

    api = HfApi()
    api.create_repo(repo_id=args.repo_id, repo_type="model", private=args.private, exist_ok=True)

    upload_file(
        path_or_fileobj=str(model_path),
        path_in_repo=model_path.name,
        repo_id=args.repo_id,
        repo_type="model",
    )

    print(f"Uploaded {model_path} to https://huggingface.co/{args.repo_id}")


if __name__ == "__main__":
    main()