#!/usr/bin/env python3
import argparse
import os
from huggingface_hub import HfApi, create_repo

def push_to_hf(repo_id, token=None, private=False):
    api = HfApi(token=token)
    
    # Create repo if it doesn't exist
    try:
        create_repo(repo_id, repo_type="model", private=private, exist_ok=True)
        print(f"Repository {repo_id} ready")
    except Exception as e:
        print(f"Repo creation note: {e}")
    
    # Upload all files in current directory
    cwd = os.getcwd()
    print(f"Uploading from {cwd} to {repo_id}...")
    
    try:
        api.upload_folder(
            folder_path=cwd,
            repo_id=repo_id,
            repo_type="model",
            commit_message="Upload Vortex model"
        )
        print(f"Successfully uploaded to {repo_id}")
    except Exception as e:
        print(f"Upload failed: {e}")
        raise

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo_id", type=str, required=True, help="HuggingFace repo ID")
    parser.add_argument("--token", type=str, help="HuggingFace token (optional if logged in)")
    parser.add_argument("--private", action="store_true", help="Make repository private")
    args = parser.parse_args()
    
    push_to_hf(args.repo_id, args.token, args.private)