#!/usr/bin/env python3
"""
Upload MemGen to HuggingFace Hub

Usage:
    python scripts/upload_to_huggingface.py --repo_id LCZZZZ/model111
"""

import os
import argparse
from pathlib import Path
from huggingface_hub import HfApi, create_repo, upload_folder
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def create_gitignore(base_path: str):
    """Create a .gitignore file to exclude unnecessary files."""
    gitignore_content = """
# Python cache
__pycache__/
*.py[cod]
*$py.class
*.so
.Python

# Virtual environments
venv/
env/
ENV/

# IDE
.vscode/
.idea/
*.swp
*.swo

# Large model files (if any)
*.bin
*.safetensors
*.pth
*.pt
*.ckpt

# Logs and outputs (optional - comment out if you want to upload)
test_output/
wandb/
*.log

# Backup files
*_backup_original/
*.backup
*.backup_no_letters

# Data caches
dataset/*/cache/
.cache/

# System files
.DS_Store
Thumbs.db
"""
    
    gitignore_path = os.path.join(base_path, ".gitignore")
    if not os.path.exists(gitignore_path):
        with open(gitignore_path, 'w') as f:
            f.write(gitignore_content)
        logging.info(f"Created .gitignore at {gitignore_path}")
    else:
        logging.info(f".gitignore already exists at {gitignore_path}")


def upload_to_hf(repo_id: str, folder_path: str, create_if_not_exists: bool = True):
    """
    Upload a folder to HuggingFace Hub.
    
    Args:
        repo_id: HuggingFace repository ID (e.g., "username/repo_name")
        folder_path: Local folder path to upload
        create_if_not_exists: Whether to create the repo if it doesn't exist
    """
    api = HfApi()
    
    # Check if repo exists, create if needed
    try:
        api.repo_info(repo_id=repo_id, repo_type="model")
        logging.info(f"Repository {repo_id} already exists")
    except Exception as e:
        if create_if_not_exists:
            logging.info(f"Repository {repo_id} not found, creating...")
            create_repo(repo_id=repo_id, repo_type="model", exist_ok=True)
            logging.info(f"Repository {repo_id} created successfully")
        else:
            raise Exception(f"Repository {repo_id} does not exist and create_if_not_exists=False")
    
    # Create .gitignore to exclude unnecessary files
    create_gitignore(folder_path)
    
    # Upload folder
    logging.info(f"Uploading {folder_path} to {repo_id}...")
    logging.info("This may take a while depending on the size...")
    
    try:
        result = upload_folder(
            repo_id=repo_id,
            folder_path=folder_path,
            repo_type="model",
            commit_message="Upload MemGen code and data",
            ignore_patterns=[
                "*.pyc",
                "__pycache__",
                ".git",
                "*.log",
                "test_output/*",
                "wandb/*",
                "*_backup_original/*",
                "*.backup",
                "*.backup_no_letters",
                ".cache/*",
            ]
        )
        
        logging.info(f"✓ Upload completed successfully!")
        logging.info(f"Repository URL: https://huggingface.co/{repo_id}")
        return result
        
    except Exception as e:
        logging.error(f"Upload failed: {e}")
        raise


def main():
    parser = argparse.ArgumentParser(description="Upload MemGen to HuggingFace Hub")
    parser.add_argument(
        "--repo_id",
        type=str,
        default="LCZZZZ/model111",
        help="HuggingFace repository ID (e.g., 'username/repo_name')"
    )
    parser.add_argument(
        "--folder_path",
        type=str,
        default="/root/CVPR/MemGen",
        help="Local folder path to upload"
    )
    parser.add_argument(
        "--no-create",
        action="store_true",
        help="Don't create repository if it doesn't exist"
    )
    
    args = parser.parse_args()
    
    if not os.path.exists(args.folder_path):
        logging.error(f"Folder not found: {args.folder_path}")
        return
    
    print("=" * 80)
    print("Upload to HuggingFace Hub")
    print("=" * 80)
    print(f"Repository: {args.repo_id}")
    print(f"Folder: {args.folder_path}")
    print("=" * 80)
    print("\nStarting upload...")
    
    upload_to_hf(
        repo_id=args.repo_id,
        folder_path=args.folder_path,
        create_if_not_exists=not args.no_create
    )


if __name__ == "__main__":
    main()