import os
import shutil
import zipfile
from huggingface_hub import snapshot_download
import gradio as gr

def download_huggingface_repo(repo_url):
    # Extract the repo ID from the URL
    try:
        if "huggingface.co/" not in repo_url:
            raise ValueError("Invalid URL: The provided URL does not belong to Hugging Face.")
        
        repo_id = repo_url.replace("https://huggingface.co/", "").strip()
        
        # Detect repo type (spaces, models, or datasets)
        if repo_id.startswith("spaces/"):
            repo_type = "space"
            repo_id = repo_id.replace("spaces/", "")
        else:
            repo_type = None  # Default type (models or datasets)

        output_dir = "repo"
        zip_file = "repo.zip"
        
        # Clean up previous files
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        if os.path.exists(zip_file):
            os.remove(zip_file)
        
        # Download the repository using Hugging Face Hub
        snapshot_download(repo_id, local_dir=output_dir, repo_type=repo_type)
        
        # Create a zip file of the downloaded repository
        with zipfile.ZipFile(zip_file, "w", zipfile.ZIP_DEFLATED) as zipf:
            for root, _, files in os.walk(output_dir):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, output_dir)
                    zipf.write(file_path, arcname)
        
        return zip_file, "Repository downloaded and zipped successfully."
    except Exception as e:
        return None, f"Error: {str(e)}"

# Gradio interface
def handle_download(repo_url):
    file_path, message = download_huggingface_repo(repo_url)
    return file_path, message

with gr.Blocks() as demo:
    gr.Markdown("## Clone and Download Hugging Face Repository")
    
    with gr.Row():
        repo_url_input = gr.Textbox(label="Hugging Face Repo URL", placeholder="Enter the repository URL")
    
    with gr.Row():
        download_button = gr.Button("Clone and Download")
    
    with gr.Row():
        output_file = gr.File(label="Download Zip File")
        output_message = gr.Textbox(label="Status", interactive=False)

    download_button.click(handle_download, inputs=[repo_url_input], outputs=[output_file, output_message])

demo.launch()