import gradio as gr
import os
import shutil
from huggingface_hub import snapshot_download, list_models

DOWNLOAD_DIR = "./models"
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

# 🔍 Search Hugging Face Models
def search_models(query, limit=5):
    try:
        models = list_models(search=query, limit=limit)
        results = []
        for m in models:
            results.append(f"📦 {m.modelId} | 🏷 {m.pipeline_tag} | ⏰ {m.lastModified}")
        return "\n".join(results)
    except Exception as e:
        return f"❌ Error: {str(e)}"

# 📥 Download Model
def download_model(model_name):
    try:
        path = snapshot_download(model_name, cache_dir=DOWNLOAD_DIR)
        return f"✅ Downloaded: {model_name}\n📂 Path: {path}"
    except Exception as e:
        return f"❌ Error: {str(e)}"

# 📂 List Downloaded Models
def list_downloaded_models():
    models_list = []
    for root, dirs, files in os.walk(DOWNLOAD_DIR):
        for d in dirs:
            models_list.append(os.path.relpath(os.path.join(root, d), DOWNLOAD_DIR))
    return models_list if models_list else ["No models downloaded yet."]

# 📑 Show files inside model folder
def show_model_files(model_folder):
    folder_path = os.path.join(DOWNLOAD_DIR, model_folder)
    if not os.path.exists(folder_path):
        return "⚠️ Folder not found"
    
    files = []
    for root, dirs, filenames in os.walk(folder_path):
        for f in filenames:
            rel_path = os.path.relpath(os.path.join(root, f), folder_path)
            files.append(rel_path)
    return "\n".join(files) if files else "📂 Empty folder"

# 🗑 Delete a downloaded model
def delete_model(model_folder):
    folder_path = os.path.join(DOWNLOAD_DIR, model_folder)
    if not os.path.exists(folder_path):
        return "⚠️ Model folder not found"
    
    try:
        shutil.rmtree(folder_path)
        return f"🗑 Deleted model: {model_folder}"
    except Exception as e:
        return f"❌ Error deleting model: {str(e)}"

# ⚡ Gradio UI
with gr.Blocks(theme="soft") as demo:
    gr.Markdown("# 🌟 Model Downloader Hub")
    gr.Markdown("Search, download, explore, and delete Hugging Face models locally.")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🔍 Search Models")
            query = gr.Textbox(label="Enter keyword or model name")
            limit = gr.Slider(1, 10, value=5, step=1, label="Number of Results")
            search_btn = gr.Button("Search")
            search_results = gr.Textbox(label="Results", interactive=False, lines=8)
            
        with gr.Column():
            gr.Markdown("### 📥 Download Models")
            model_name = gr.Textbox(label="Model Name (from Hugging Face)")
            download_btn = gr.Button("Download")
            status = gr.Textbox(label="Status", interactive=False)

    gr.Markdown("### 📂 Local Model Explorer")
    refresh_btn = gr.Button("🔄 Refresh Model List")
    model_list = gr.Dropdown(choices=[], label="Downloaded Models")
    show_files_btn = gr.Button("📑 Show Files")
    delete_btn = gr.Button("🗑 Delete Selected Model")
    file_list = gr.Textbox(label="Files in Model Folder", interactive=False, lines=12)
    delete_status = gr.Textbox(label="Delete Status", interactive=False)

    # 🔗 Button Actions
    search_btn.click(search_models, [query, limit], search_results)
    download_btn.click(download_model, [model_name], status)
    refresh_btn.click(list_downloaded_models, outputs=model_list)
    show_files_btn.click(show_model_files, [model_list], file_list)
    delete_btn.click(delete_model, [model_list], delete_status).then(list_downloaded_models, outputs=model_list)

demo.launch()