import os
import shutil
import tempfile
import gradio as gr
from huggingface_hub import snapshot_download, model_info, dataset_info, space_info
from huggingface_hub.utils import RepositoryNotFoundError, HFValidationError

def detect_repo_type(repo_id):
    """Automatically detect repo type: model, dataset, space"""
    try:
        model_info(repo_id)
        return "model"
    except (RepositoryNotFoundError, HFValidationError):
        pass
    try:
        dataset_info(repo_id)
        return "dataset"
    except (RepositoryNotFoundError, HFValidationError):
        pass
    try:
        space_info(repo_id)
        return "space"
    except (RepositoryNotFoundError, HFValidationError):
        pass
    return None

def zip_folder(source_folder):
    """Zip the folder into a temporary ZIP file"""
    tmp_zip_fd, tmp_zip_path = tempfile.mkstemp(suffix=".zip")
    os.close(tmp_zip_fd)  # Close fd, just need path
    shutil.make_archive(tmp_zip_path.replace(".zip", ""), 'zip', source_folder)
    return tmp_zip_path

def download_hf_repo_temp(repo_id, repo_type_override=None):
    repo_id = repo_id.strip()
    if not repo_id or "/" not in repo_id:
        return "❌ Invalid repository ID. Format should be: username/repo_name", None

    # Use override if provided, otherwise auto-detect
    if repo_type_override and repo_type_override != "auto":
        repo_type = repo_type_override
        log = f"🔍 Using specified repo type: `{repo_type}`\n"
    else:
        repo_type = detect_repo_type(repo_id)
        if not repo_type:
            return f"❌ Repository '{repo_id}' not found or inaccessible", None
        log = f"🔍 Auto-detected repo type: `{repo_type}`\n"

    # Create temporary folder
    temp_dir = tempfile.mkdtemp()
    log += f"⬇️ Downloading '{repo_id}' to temporary folder...\n"

    try:
        # Download repo
        snapshot_download(
            repo_id=repo_id,
            repo_type=repo_type,
            local_dir=temp_dir,
            local_dir_use_symlinks=False
        )

        # Zip temp folder
        zip_path = zip_folder(temp_dir)
        log += "✅ Download complete and zipped successfully!\n"
        log += f"📦 ZIP file created: {os.path.basename(zip_path)}\n"
        log += "🗑️ Temporary files will be automatically cleaned up."

        # Cleanup temp folder after creating zip
        shutil.rmtree(temp_dir)

        return log, zip_path
    except Exception as e:
        # Cleanup on error
        shutil.rmtree(temp_dir)
        return f"❌ Error during download: {str(e)}", None

# Custom CSS for better styling
custom_css = """
#main-header {
    text-align: center;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    padding: 2rem;
    border-radius: 10px;
    margin-bottom: 1rem;
    color: white;
}
#main-header h1 {
    margin: 0;
    font-size: 2.5rem;
    font-weight: 700;
}
#main-header p {
    margin: 0.5rem 0 0 0;
    opacity: 0.9;
    font-size: 1.1rem;
}
.feature-grid {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
    gap: 1rem;
    margin: 1.5rem 0;
}
.feature-card {
    background: white;
    padding: 1rem;
    border-radius: 8px;
    border-left: 4px solid #667eea;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.download-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    font-weight: 600 !important;
    border: none !important;
}
.download-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 4px 8px rgba(0,0,0,0.2);
}
.log-output {
    font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
    font-size: 0.9em;
}
"""

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    
    # Header Section
    with gr.Column(elem_id="main-header"):
        gr.HTML("""
        <div style="text-align: center;">
            <h1>🤗 HF Downloader</h1>
            <p>Download & Package Hugging Face Repositories</p>
        </div>
        """)
    
    # Main Description
    gr.Markdown("""
    **Easily download any Hugging Face model, dataset, or space repository and get it as a ZIP file.** 
    Perfect for backups, offline use, or sharing repositories.
    """)
    
    # Feature Cards
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🚀 How to Use")
            gr.Markdown("""
            1. Enter repository ID (e.g., `Jonny001/HF-Downloader`)
            2. Choose repo type or leave as 'Auto-detect'
            3. Click **Download & ZIP**
            4. Wait for processing and download your ZIP file
            """)
        
        with gr.Column():
            gr.Markdown("### 📚 Supported Types")
            gr.Markdown("""
            - **Models**: Machine learning models
            - **Datasets**: Training/evaluation datasets  
            - **Spaces**: Gradio/demo applications
            - **Auto-detection**: Automatic type recognition
            """)

    gr.Markdown("---")
    
    # Input Section
    with gr.Row():
        with gr.Column(scale=3):
            repo_input = gr.Textbox(
                label="Hugging Face Repository ID",
                placeholder="e.g. Jonny001/HF-Downloader, microsoft/DialoGPT-medium, huggingface/datasets",
                info="Enter the full repository ID including username"
            )
        with gr.Column(scale=1):
            repo_type = gr.Dropdown(
                choices=["auto", "model", "dataset", "space"],
                value="auto",
                label="Repository Type",
                info="Choose type or let us auto-detect"
            )
    
    # Action Section
    with gr.Row():
        download_btn = gr.Button(
            "🚀 Download & Create ZIP", 
            variant="primary",
            size="lg",
            elem_classes="download-btn"
        )
        clear_btn = gr.Button("🗑️ Clear", size="lg")
    
    # Output Section
    with gr.Row():
        with gr.Column():
            log_output = gr.Textbox(
                label="📋 Process Log",
                lines=10,
                interactive=False,
                elem_classes="log-output",
                show_copy_button=True
            )
    
    with gr.Row():
        file_output = gr.File(
            label="📦 Download ZIP File",
            file_types=[".zip"],
            height=100
        )
    
    # Examples Section
    gr.Markdown("### 🎯 Quick Examples")
    examples = gr.Examples(
        examples=[
            ["Jonny001/HF-Downloader", "auto"],
            ["Jonny001/R.I.O.W", "auto"],
            ["huggingface/datasets", "auto"],
            ["Jonny001/HF-Downloader", "space"]
        ],
        inputs=[repo_input, repo_type],
        label="Click any example to try it out!"
    )
    
    # Event Handlers
    download_btn.click(
        fn=download_hf_repo_temp,
        inputs=[repo_input, repo_type],
        outputs=[log_output, file_output]
    )
    
    def clear_all():
        return "", None, "auto"
    
    clear_btn.click(
        fn=clear_all,
        outputs=[repo_input, file_output, repo_type]
    )
    
    # Footer
    gr.Markdown("---")
    gr.Markdown("""
    <div style="text-align: center; color: #666;">
        <p>💡 <strong>Tip:</strong> Large repositories may take several minutes to download and process.</p>
        <p>🔒 All temporary files are automatically cleaned up after download.</p>
    </div>
    """)

if __name__ == "__main__":
    demo.launch(
        share=False,
        show_error=True
    )