app.py

import os
import gradio as gr
from huggingface_hub import HfApi, upload_folder, create_repo, login, list_repo_files
from rich.console import Console
from rich.logging import RichHandler
import logging
import time
import shutil
import threading
from pathlib import Path

# Default values
DEFAULT_REPO = os.getenv("DEFAULT_REPO", "luigi12345/megacursos-master-999")
DEFAULT_LOCAL_PATH = os.getenv(
    "DEFAULT_LOCAL_PATH", "/Users/samihalawa/Documents/Megacursos/MEGACURSOS_S3_MASTER"
)
token = os.getenv("HF_TOKEN", "")

# Initialize Rich console for logging
console = Console()
logging.basicConfig(
    level=logging.INFO,
    format="%(message)s",
    datefmt="[%Y-%m-%d %H:%M:%S]",
    handlers=[RichHandler(console=console, rich_tracebacks=True)],
)
logger = logging.getLogger("InfiniteStorageFace")

# Initialize Hugging Face API client
api = HfApi()

# Centralized ignore patterns mapping
IGNORE_PATTERNS_MAP = {
    "Ignore __pycache__": "**/__pycache__/**",
    "Ignore .git": ".git/**",
    "Ignore .venv": "venv/**",
    "Ignore *.pyc": "*.pyc",
    "Ignore *.log": "*.log",
    "Ignore *.tmp": "*.tmp",
    "Ignore *.DS_Store": "*.DS_Store",
}

# Shared logs list
shared_logs = []

# Event to cancel upload
cancel_event = False


# Function to log messages
def log(message):
    timestamp = time.strftime("[%Y-%m-%d %H:%M:%S]")
    full_message = f"{timestamp} {message}"
    shared_logs.append(full_message)
    logger.info(message)
    return full_message


# Function to authenticate user with Hugging Face token
def authenticate(token):
    if not token:
        return False, log("❌ Hugging Face Token is required.")
    try:
        login(token)
        return True, log("✅ Authenticated successfully!")
    except Exception as e:
        return False, log(f"❌ Authentication failed: {e}")


# Function to create repository if it doesn't exist
def create_repo_if_not_exists(repo_id, token, repo_type, private):
    try:
        api.list_repo_files(repo_id=repo_id, repo_type=repo_type, token=token)
        return True, log(f"✅ Repository '{repo_id}' exists. Proceeding with upload...")
    except Exception:
        log(f"❌ Repository '{repo_id}' does not exist. Creating it...")
        try:
            create_repo(
                repo_id=repo_id,
                token=token,
                private=private,
                repo_type=repo_type,
                exist_ok=True,
                space_sdk="static" if repo_type == "space" else None,
            )
            return True, log(f"✅ Created new repository: '{repo_id}'.")
        except Exception as create_err:
            return False, log(
                f"❌ Failed to create repository '{repo_id}': {create_err}"
            )


# Function to clean up ignored folders
def cleanup_before_upload(folder_path, ignore_patterns):
    for pattern in ignore_patterns:
        for path in Path(folder_path).rglob(pattern):
            if path.is_dir():
                shutil.rmtree(path)
                log(f"🗑️ Removed ignored folder: {path}")


# Function to upload a folder
def upload_folder_structure(
    folder_path, repo_id, token, repo_type, target_path, ignore_patterns
):
    # Clean up ignored folders
    cleanup_before_upload(folder_path, ignore_patterns)

    upload_params = {
        "folder_path": folder_path,
        "repo_id": repo_id,
        "repo_type": repo_type,
        "token": token,
        "path_in_repo": target_path,
        "multi_commits": True,
        "multi_commits_verbose": True,
    }
    log(
        f"🚀 Uploading folder '{folder_path}' to '{target_path}' in repository '{repo_id}'..."
    )
    try:
        upload_folder(**upload_params)
        log(f"✅ Upload completed for '{folder_path}'!")
    except Exception as upload_err:
        log(f"❌ Upload failed for '{folder_path}': {upload_err}")


# Function to handle uploads, allowing direct folder uploads
def upload_files(
    files,
    repo_id,
    token,
    private,
    threads,
    subfolder,
    repo_type,
    ignore_patterns_selected,
):
    global cancel_event
    cancel_event = False

    logs = []

    # Authenticate
    auth_success, auth_message = authenticate(token)
    logs.append(auth_message)
    if not auth_success:
        return "\n".join(logs)

    # Create repo if not exists
    repo_success, repo_message = create_repo_if_not_exists(
        repo_id, token, repo_type, private
    )
    logs.append(repo_message)
    if not repo_success:
        return "\n".join(logs)

    # Prepare target path
    target_path = subfolder.replace("\\", "/") if subfolder else ""

    # Map selected ignore patterns to actual patterns
    ignore_patterns = [
        IGNORE_PATTERNS_MAP[pattern] for pattern in ignore_patterns_selected
    ]

    # Upload folders directly if provided
    if files:
        for file in files:
            if cancel_event:
                logs.append(log("❌ Upload has been cancelled."))
                return "\n".join(logs)
            file_path = file.name
            if os.path.isdir(file_path):  # Check if the file is a directory
                upload_folder_structure(
                    file_path, repo_id, token, repo_type, target_path, ignore_patterns
                )
                logs.append(log(f"✅ Uploaded folder '{file_path}'."))
            else:
                logs.append(
                    log(
                        f"❌ '{file_path}' is not a folder. Only folders can be uploaded."
                    )
                )
    else:
        logs.append(log("❌ No files provided for upload."))

    if cancel_event:
        logs.append(log("❌ Upload has been cancelled."))
        return "\n".join(logs)

    logs.append(log("🚀 Upload completed. Check the logs for details."))
    return "\n".join(logs)


# Function to cancel upload
def cancel_upload():
    global cancel_event
    cancel_event = True
    return log("Upload has been cancelled.")


# Function to refresh logs
def refresh_logs():
    return "\n".join(shared_logs)


# Gradio Interface
def create_interface():
    with gr.Blocks() as app:
        gr.Markdown("# 🚀 InfiniteStorageFace", elem_id="main-title")
        gr.Markdown(
            "Effortlessly upload your files or folders to Hugging Face repositories with real-time feedback and progress tracking!",
            elem_id="sub-title"
        )

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## Upload Section", elem_id="upload-section")

                token = gr.Textbox(label="Hugging Face Token", type="password",
                    placeholder="Enter your Hugging Face API token",
                    value=os.getenv("HF_TOKEN", ""),
                    interactive=True
                )

                repo_type = gr.Radio(
                    label="Repository Type", choices=["space", "model", "dataset"],
                    value="space", interactive=True
                )

                repo_id = gr.Textbox(label="Repository ID",
                    placeholder="e.g., username/repo-name", value=os.getenv("DEFAULT_REPO", ""),
                    interactive=True
                )

                private = gr.Checkbox(label="Make Repository Private",
                    value=False, interactive=True
                )

                files = gr.Files(label="Upload Folders",
                    file_count="multiple", interactive=True, type="directory"
                )

                subfolder = gr.Textbox(label="Subfolder in Repository (Optional)",
                    placeholder="e.g., data/uploads", value="", interactive=True
                )

                threads = gr.Slider(label="Number of Threads", minimum=1, maximum=20,
                    step=1, value=5, interactive=True
                )

                ignore_patterns_selected = gr.CheckboxGroup(label="Select Patterns to Ignore",
                    choices=list(IGNORE_PATTERNS_MAP.keys()), value=["Ignore __pycache__", "Ignore .git", "Ignore *.pyc"],
                    interactive=True
                )

                upload_button = gr.Button("Start Upload", variant="primary", interactive=True)
                cancel_button = gr.Button("Cancel Upload", variant="secondary", interactive=True)

            with gr.Column(scale=1):
                gr.Markdown("## Status Section", elem_id="status-section")
                upload_status = gr.Textbox(label="Upload Status", lines=10,
                    interactive=False, value="Idle"
                )

                with gr.Tab("Logs"):
                    log_output = gr.Textbox(label="Upload Logs", lines=15,
                        interactive=False, placeholder="Logs will appear here...", value=""
                    )
                    log_refresh = gr.Button("Refresh Logs", interactive=True)

        # Define the upload button click event
        def handle_upload(*args):
            threading.Thread(target=upload_files, args=args).start()

        upload_button.click(fn=handle_upload, inputs=[
            files, repo_id, token, private, threads, subfolder, repo_type, ignore_patterns_selected
        ], outputs=upload_status)

        cancel_button.click(fn=cancel_upload, inputs=None, outputs=upload_status)
        
        log_refresh.click(fn=refresh_logs, inputs=None, outputs=log_output)

        return app


if __name__ == "__main__":
    app = create_interface()
    app.launch(debug=True, share=True, server_port=7860)