#!/usr/bin/env python3
"""
Script to upload all .pt and .onnx model files from models_mask directory to Hugging Face.
"""

import os
import glob
from pathlib import Path
from huggingface_hub import HfApi, login
from tqdm import tqdm

def upload_models_mask():
    """
    Upload all .pt and .onnx files from models_mask directory to Hugging Face
    """
    
    # Configuration
    repo_name = "PetchMa/BLADE_FRBNN"  # Change this to your desired repo name
    models_dir = "models_mask"
    
    # Initialize HF API
    api = HfApi()
    
    # Login to Hugging Face (you need to have your token set up)
    print("Logging into Hugging Face...")
    try:
        login()
        print("✓ Successfully logged into Hugging Face")
    except Exception as e:
        print(f"❌ Failed to login to Hugging Face: {e}")
        print("Please make sure you have your HF token set up:")
        print("1. Run: huggingface-cli login")
        print("2. Or set HF_TOKEN environment variable")
        return
    
    # Check if models directory exists
    if not os.path.exists(models_dir):
        print(f"❌ Directory {models_dir} does not exist!")
        return
    
    # Find all .pt and .onnx files
    pt_files = glob.glob(os.path.join(models_dir, "*.pt"))
    onnx_files = glob.glob(os.path.join(models_dir, "*.onnx"))
    
    all_files = pt_files + onnx_files
    
    if not all_files:
        print(f"❌ No .pt or .onnx files found in {models_dir}")
        return
    
    print(f"Found {len(all_files)} model files to upload:")
    print(f"  - {len(pt_files)} .pt files")
    print(f"  - {len(onnx_files)} .onnx files")
    
    # Calculate total size
    total_size = sum(os.path.getsize(f) for f in all_files)
    total_size_gb = total_size / (1024**3)
    print(f"  - Total size: {total_size_gb:.2f} GB")
    
    # Create repository if it doesn't exist
    try:
        api.create_repo(repo_name, exist_ok=True)
        print(f"✓ Repository {repo_name} is ready")
    except Exception as e:
        print(f"❌ Failed to create repository: {e}")
        return
    
    # Use upload_large_folder for better performance with many files
    print(f"\nUploading {models_dir} directory to {repo_name}...")
    print("This may take a while due to the large number of files...")
    
    try:
        api.upload_large_folder(
            folder_path=models_dir,
            repo_id=repo_name,
            repo_type="model",
            commit_message=f"Upload models_mask directory with {len(all_files)} model files",
            allow_patterns=["*.pt", "*.onnx"],  # Only upload .pt and .onnx files
            ignore_patterns=["*.png", "*.pdf", "*.pkl", "*.bin"],  # Ignore other file types
        )
        
        print(f"\n✓ Successfully uploaded {len(all_files)} model files!")
        print(f"✓ Repository: https://huggingface.co/{repo_name}")
        print(f"✓ Models available at: https://huggingface.co/{repo_name}/tree/main/{models_dir}")
        
    except Exception as e:
        print(f"❌ Failed to upload folder: {e}")
        print("You might want to try uploading in smaller batches or check your internet connection.")

def main():
    """
    Main function to run the upload process
    """
    print("BLADE FRBNN Model Upload Tool")
    print("=" * 40)
    print("This script will upload all .pt and .onnx files from models_mask/")
    print("to your Hugging Face repository using the efficient large folder upload.")
    print()
    
    # Ask for confirmation
    response = input("Do you want to proceed? (y/N): ")
    if response.lower() not in ['y', 'yes']:
        print("Upload cancelled.")
        return
    
    upload_models_mask()

if __name__ == "__main__":
    main()