#!/usr/bin/env python3
import os
import sys
import argparse
from huggingface_hub import HfApi, login, create_repo
from pathlib import Path

DEFAULT_IGNORE = [
    "**/.git/**", "**/__pycache__/**", "**/.DS_Store",
    "**/node_modules/**", ".gitignore", "*.log", "**/.venv/**", "**/env/**"
]

def main():
    p = argparse.ArgumentParser(
        description="Upload a local folder (code + models) to Hugging Face Hub."
    )
    p.add_argument("--token", help="HF write token. If omitted, reads from HF_TOKEN env.")
    p.add_argument("--repo-id", required=True, help="e.g. username/comfyui-bundle")
    p.add_argument("--repo-type", default="model", choices=["model","dataset","space"])
    p.add_argument("--private", action="store_true", help="Create repo as private")
    p.add_argument("--folder", required=True, help="Local folder to upload")
    p.add_argument("--path-in-repo", default="", help="Subfolder inside repo (default root)")
    p.add_argument("--commit-message", default="Upload folder")
    p.add_argument("--allow", nargs="*", default=None,
                   help="Glob patterns to ALLOW only (e.g. *.py *.safetensors). If omitted, upload all except ignored.")
    p.add_argument("--ignore", nargs="*", default=None,
                   help="Extra ignore glob patterns (merged with defaults).")
    p.add_argument("--max-workers", type=int, default=8,
                   help="Parallel workers for upload (default 8)")
    args = p.parse_args()

    token = args.token or os.getenv("HF_TOKEN")
    if not token:
        print("ERROR: Provide --token or set HF_TOKEN env.", file=sys.stderr)
        sys.exit(1)

    folder = Path(args.folder).expanduser().resolve()
    if not folder.exists() or not folder.is_dir():
        print(f"ERROR: Folder not found: {folder}", file=sys.stderr)
        sys.exit(1)

    # Login and ensure repo exists
    login(token=token)
    create_repo(
        repo_id=args.repo_id,
        repo_type=args.repo_type,
        private=args.private,
        exist_ok=True
    )

    # Build ignore list
    ignore_patterns = list(DEFAULT_IGNORE)
    if args.ignore:
        ignore_patterns.extend(args.ignore)

    # Safety: if uploading một repo git khác, chắc chắn bỏ .git
    if ".git" not in " ".join(ignore_patterns):
        ignore_patterns.append("**/.git/**")

    print("=== Upload config ===")
    print(f"Folder        : {folder}")
    print(f"Repo ID       : {args.repo_id} ({args.repo_type})")
    print(f"Private       : {args.private}")
    print(f"Path in repo  : {args.path_in_repo or '/'}")
    print(f"Allow patterns: {args.allow}")
    print(f"Ignore patterns: {ignore_patterns}")
    print(f"Max workers   : {args.max_workers}")
    print("=====================")

    api = HfApi()

    # Note:
    # - upload_folder tự chia nhỏ & resume file >10MB
    # - allow_patterns: chỉ upload đúng các glob này nếu được set
    # - ignore_patterns: bỏ qua những thứ không muốn đẩy
    api.upload_folder(
        folder_path=str(folder),
        repo_id=args.repo_id,
        repo_type=args.repo_type,
        path_in_repo=args.path_in_repo,
        commit_message=args.commit_message,
        allow_patterns=args.allow,          # None = all files
        ignore_patterns=ignore_patterns,
    )

    print("✅ Done. Check:", f"https://huggingface.co/{args.repo_id}")

if __name__ == "__main__":
    main()