from huggingface_hub import HfApi
import os
import argparse

def upload_to_hub(
    local_path: str = ".",
    username: str = "puremood",
    model_name: str = "llama70b-MARTZ",
    token: str = None,
    private: bool = False
):
    """
    Upload a model to Hugging Face Hub
    """
    # Get token from environment if not provided
    if not token:
        token = os.getenv('HF_TOKEN')
        if not token:
            raise ValueError("No token provided. Set HF_TOKEN environment variable or pass token argument")
    
    # Initialize the API
    api = HfApi(token=token)
    
    # Construct the repo_id
    repo_id = f"{username}/{model_name}"
    
    print(f"Creating/accessing repository: {repo_id}")
    
    try:
        # Create or access the repository
        api.create_repo(
            repo_id=repo_id,
            private=private,
            exist_ok=True
        )
        
        print(f"Uploading files from {local_path} to {repo_id}")
        
        # Upload the folder
        api.upload_folder(
            folder_path=local_path,
            repo_id=repo_id,
            repo_type="model"
        )
        
        print(f"Upload complete! Your model is available at: https://huggingface.co/{repo_id}")
        
    except Exception as e:
        print(f"Error occurred: {str(e)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Upload model to Hugging Face Hub')
    parser.add_argument('--path', default=".", help='Path to model files')
    parser.add_argument('--username', default="puremood", help='Hugging Face username')
    parser.add_argument('--model', default="llama70b-MARTZ", help='Model name')
    parser.add_argument('--token', help='Hugging Face token (or set HF_TOKEN env variable)')
    parser.add_argument('--private', action='store_true', help='Make repository private')
    
    args = parser.parse_args()
    
    upload_to_hub(
        local_path=args.path,
        username=args.username,
        model_name=args.model,
        token=args.token,
        private=args.private
    )