import gradio as gr
from huggingface_hub import HfApi
import os

api = HfApi()

def list_models(search_query):
    models = api.list_models(search=search_query, limit=10)
    return [model.modelId for model in models]

def list_datasets(search_query):
    datasets = api.list_datasets(search=search_query, limit=10)
    return [dataset.id for dataset in datasets]

def download_model(model_id):
    try:
        api.snapshot_download(repo_id=model_id, repo_type="model")
        return f"Modell {model_id} erfolgreich heruntergeladen."
    except Exception as e:
        return f"Fehler beim Herunterladen des Modells {model_id}: {str(e)}"

def download_dataset(dataset_id):
    try:
        api.snapshot_download(repo_id=dataset_id, repo_type="dataset")
        return f"Dataset {dataset_id} erfolgreich heruntergeladen."
    except Exception as e:
        return f"Fehler beim Herunterladen des Datasets {dataset_id}: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# Modell- und Dataset-Manager")
    
    with gr.Tab("Modelle"):
        model_search = gr.Textbox(label="Modell-Suche")
        model_list = gr.Dropdown(label="Verfügbare Modelle")
        model_search_btn = gr.Button("Modelle suchen")
        model_download_btn = gr.Button("Ausgewähltes Modell herunterladen")
        model_output = gr.Textbox(label="Ausgabe")

        model_search_btn.click(list_models, inputs=model_search, outputs=model_list)
        model_download_btn.click(download_model, inputs=model_list, outputs=model_output)

    with gr.Tab("Datasets"):
        dataset_search = gr.Textbox(label="Dataset-Suche")
        dataset_list = gr.Dropdown(label="Verfügbare Datasets")
        dataset_search_btn = gr.Button("Datasets suchen")
        dataset_download_btn = gr.Button("Ausgewähltes Dataset herunterladen")
        dataset_output = gr.Textbox(label="Ausgabe")

        dataset_search_btn.click(list_datasets, inputs=dataset_search, outputs=dataset_list)
        dataset_download_btn.click(download_dataset, inputs=dataset_list, outputs=dataset_output)

demo.launch()