import gradio as gr
from huggingface_hub import HfApi
from git import Repo
import uuid 
from slugify import slugify
import os
import shutil

def clone_and_upload(
    profile: gr.OAuthProfile,
    oauth_token: gr.OAuthToken, 
    repo_git: str,
    repo_type: str,
    repo_hf: str,
    sdk_type: str = None
):
    """
    Clone a Git repository and upload it to Hugging Face as a specified type.
    
    Args:
        profile: Gradio OAuth profile with user info
        oauth_token: OAuth token for authentication
        repo_git: Git repository URL to clone
        repo_type: Type of repository to create on HF (space, model, dataset)
        repo_hf: Name for the Hugging Face repository
        sdk_type: SDK type for spaces (gradio, streamlit, docker, static)
    
    Returns:
        URL of the created Hugging Face repository
    """
    folder = str(uuid.uuid4())
    
    try:
        # Clone the Git repository
        cloned_repo = Repo.clone_from(repo_git, folder)
        
        # Initialize API with token
        api = HfApi(token=oauth_token.token)
        
        # Create slugified repository name
        repo_name = slugify(repo_hf)
        repo_id = f"{profile.username}/{repo_name}"
        
        # Prepare repository creation parameters
        create_kwargs = {
            "repo_id": repo_id,
            "repo_type": repo_type,
        }
        
        # Add space_sdk only for spaces
        if repo_type == "space" and sdk_type:
            create_kwargs["space_sdk"] = sdk_type
        
        # Create repository on Hugging Face
        api.create_repo(**create_kwargs)
        
        # Upload the folder content
        if repo_type == "dataset":
            # For datasets, we typically want to preserve the folder structure
            api.upload_folder(
                folder_path=folder,
                repo_id=repo_id,
                repo_type=repo_type,
                commit_message="Initial upload from Git repository"
            )
        else:
            # For models and spaces, upload the entire folder
            api.upload_folder(
                folder_path=folder,
                repo_id=repo_id,
                repo_type=repo_type,
                commit_message="Initial upload from Git repository"
            )
        
        # Return appropriate URL based on repo type
        if repo_type == "space":
            return f"https://huggingface.co/spaces/{profile.username}/{repo_name}"
        elif repo_type == "model":
            return f"https://huggingface.co/{profile.username}/{repo_name}"
        elif repo_type == "dataset":
            return f"https://huggingface.co/datasets/{profile.username}/{repo_name}"
        else:
            return f"https://huggingface.co/{profile.username}/{repo_name}"
            
    except Exception as e:
        return f"Error: {str(e)}"
    finally:
        # Clean up the cloned directory
        if os.path.exists(folder):
            shutil.rmtree(folder)

def toggle_sdk_visibility(repo_type):
    """
    Toggle visibility of SDK choices based on repo type.
    
    Args:
        repo_type: Selected repository type
    
    Returns:
        Gradio update object for SDK choices visibility
    """
    if repo_type == "space":
        return gr.Radio(visible=True, interactive=True)
    else:
        return gr.Radio(visible=False, interactive=False)

with gr.Blocks(theme="NeoPy/Soft") as demo:
    gr.Markdown("# 🚀 Multi-Type Repository Importer")
    gr.Markdown("Clone Git repositories and upload them to Hugging Face as Spaces, Models, or Datasets")
    
    # Authentication
    gr.LoginButton()
    
    with gr.Row():
        with gr.Column(scale=2):
            # Input section
            gr.Markdown("## Input Configuration")
            
            repo_git = gr.Textbox(
                label="Git Repository URL",
                placeholder="https://github.com/username/repository.git",
                info="Enter the Git repository URL to clone"
            )
            
            repo_type = gr.Dropdown(
                choices=["space", "model", "dataset"],
                value="space",
                label="Hugging Face Repository Type",
                info="Select the type of repository to create on Hugging Face"
            )
            
            repo_hf = gr.Textbox(
                label="Hugging Face Repository Name",
                placeholder="my-awesome-project",
                info="Choose a name for your Hugging Face repository"
            )
            
            # Conditional SDK selection (only for spaces)
            sdk_choices = gr.Radio(
                choices=["gradio", "streamlit", "docker", "static"],
                value="gradio",
                label="Space SDK",
                info="Select the SDK for your Space (only applicable for Spaces)",
                visible=True
            )
            
            # Update SDK visibility based on repo type
            repo_type.change(
                fn=toggle_sdk_visibility,
                inputs=repo_type,
                outputs=sdk_choices
            )
            
            # Submit button
            submit_btn = gr.Button(
                "🚀 Clone & Upload to Hugging Face",
                variant="primary",
                size="lg"
            )
        
        with gr.Column(scale=1):
            # Info section
            gr.Markdown("## 📚 Repository Type Guide")
            
            with gr.Accordion("Spaces", open=False):
                gr.Markdown("""
                **Spaces** are for hosting ML demos and apps.
                - Requires SDK (Gradio, Streamlit, Docker, or static)
                - Automatically deploys your application
                - Public URL for sharing demos
                """)
            
            with gr.Accordion("Models", open=False):
                gr.Markdown("""
                **Models** are for hosting ML models.
                - Use for model weights and configurations
                - Compatible with Hugging Face Transformers
                - Version control for model files
                """)
            
            with gr.Accordion("Datasets", open=False):
                gr.Markdown("""
                **Datasets** are for hosting data collections.
                - Store and version datasets
                - Use with Hugging Face Datasets library
                - Share datasets with the community
                """)
    
    # Output section
    gr.Markdown("## Output")
    output = gr.Textbox(
        label="Result",
        placeholder="Your Hugging Face repository URL will appear here...",
        interactive=False
    )
    
    # Status tracking
    status = gr.Textbox(
        label="Status",
        value="Ready to import",
        interactive=False,
        visible=True
    )
    
    def process_upload(profile, oauth_token, repo_git, repo_type, repo_hf, sdk_type):
        if not all([repo_git, repo_hf]):
            return "Please fill in all required fields", "Error: Missing required fields"
        
        try:
            status_update = f"Cloning {repo_git}..."
            yield status_update, ""
            
            result = clone_and_upload(
                profile=profile,
                oauth_token=oauth_token,
                repo_git=repo_git,
                repo_type=repo_type,
                repo_hf=repo_hf,
                sdk_type=sdk_type
            )
            
            if result.startswith("Error:"):
                return result, f"Failed: {result}"
            else:
                return result, f"✅ Successfully created! Repository is being processed..."
                
        except Exception as e:
            return f"Error: {str(e)}", f"Failed: {str(e)}"
    
    # Connect the button click event
    submit_btn.click(
        fn=process_upload,
        inputs=[gr.State(), gr.State(), repo_git, repo_type, repo_hf, sdk_choices],
        outputs=[output, status]
    )
    
    # Examples section
    gr.Markdown("## 💡 Examples")
    
    with gr.Row():
        with gr.Column():
            gr.Examples(
                examples=[
                    ["https://github.com/gradio-app/gradio.git", "space", "gradio-demo", "gradio"],
                    ["https://github.com/huggingface/transformers.git", "model", "my-transformers", ""],
                    ["https://github.com/huggingface/datasets.git", "dataset", "my-dataset", ""]
                ],
                inputs=[repo_git, repo_type, repo_hf, sdk_choices],
                label="Quick Examples"
            )
    
    # Footer
    gr.Markdown("---")
    gr.Markdown(
        """
        <div style='text-align: center'>
        <p>This tool helps you import Git repositories to Hugging Face Hub</p>
        <p>Supports: Spaces (with various SDKs), Models, and Datasets</p>
        </div>
        """
    )

if __name__ == "__main__":
    demo.launch()