import gradio as gr
from huggingface_hub import HfApi
import os
import shutil

# Initialize Hugging Face API
api = HfApi()

# Function to clean directory if exists and recreate it
def clean_directory(directory):
    if os.path.exists(directory):
        shutil.rmtree(directory)  # Remove the entire directory
    os.makedirs(directory)  # Recreate an empty directory

# Function to upload and push to Hugging Face
def push_to_huggingface(repo_name, file_or_directory, hf_token):
    # Ensure correct repo name format
    if not "/" in repo_name:
        repo_name = f"{hf_token.split(':')[0]}/{repo_name}"  # Assume the namespace is the user's namespace

    # Step 1: Handle directories or files
    target_dir = f"./{repo_name}"
    if os.path.isdir(file_or_directory):
        shutil.copytree(file_or_directory, target_dir, dirs_exist_ok=True)
    else:
        os.makedirs(target_dir, exist_ok=True)
        shutil.copy(file_or_directory.name, os.path.join(target_dir, os.path.basename(file_or_directory.name)))

    # Step 2: Check if the repository already exists, create if not
    try:
        api.create_repo(repo_name, token=hf_token)
    except Exception as e:
        if "already created" in str(e):
            print(f"Repository '{repo_name}' already exists, skipping creation.")
        else:
            raise e

    # Step 3: Push files using HTTP-based upload
    if os.path.isdir(target_dir):
        api.upload_folder(folder_path=target_dir, path_in_repo="", repo_id=repo_name, token=hf_token)
    else:
        api.upload_file(path_or_fileobj=target_dir, path_in_repo=os.path.basename(file_or_directory.name), repo_id=repo_name, token=hf_token)

    return f"Repository '{repo_name}' has been pushed to Hugging Face Hub successfully."

# Gradio Interface
interface = gr.Interface(
    fn=push_to_huggingface,
    inputs=[
        "text",    # Repository name
        "file",    # File or directory
        "text"     # Hugging Face Token
    ],
    outputs="text",
    title="Push File or Directory to Hugging Face",
    description="Upload any file or directory and push it to the Hugging Face Hub."
)

# Launch the interface
if __name__ == "__main__":
    interface.launch()