import gradio as gr
import random
from huggingface_hub import HfApi

def get_models(author: str, library: str, model_name: str):
    # filter out based on above
    api = HfApi()
    filtered_models = api.list_models(
        author=None if author == "" else author,
        library=None if library == "" else library,
        model_name=None if model_name == "" else model_name,
    )

    filtered_models_list = list(filtered_models)
    random_model = filtered_models_list[random.randrange(0, len(filtered_models_list))].id
    return gr.Textbox(value=random_model, info=f"Chosen from {len(filtered_models_list)} possible models based on your selection", label="Random Model", interactive=False)


def get_pipeline_template(model: str) -> str:
    return f"""
# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline(model="{model}")
"""


def create_code_block(chosen: str):
    return gr.Code(
        label="Python",
        value=get_pipeline_template(chosen),
        language="python",
        visible=True
    )



with gr.Blocks() as demo:
    author = gr.Textbox(label="Author", info="A string which identify the author (user or organization) of the returned models", interactive=True)
    library = gr.Textbox(label="Library", info="A string or list of strings of foundational libraries models were originally trained from", interactive=True)
    model_name = gr.Textbox(label="Model Name", info="A string that contain complete or partial names for models on the Hub", interactive=True)

    gr.Examples(
        examples=[
            ["google", "pytorch", "flan"],
            ["meta-llama", "transformers", "Llama-2"]
        ], 
        inputs=[author, library, model_name],
        run_on_click=True,
        fn=get_models
    )

    submit = gr.Button(value="Generate Random Model")
    
    with gr.Row(equal_height=True):
        chosen = gr.Textbox(label="Random Model", interactive=False)
        # open_model = gr.Button(value="Open Model Card", link="https://www.google.com")
        open_model = gr.Button(value="Open Model Card")
        generate_code = gr.Button(value="Generate code")


    submit.click(
        get_models,
        [author, library, model_name],
        chosen
    )

    open_model.click(
        fn=None,
        inputs=chosen,
        js=f"(chosen) => {{ window.open('https://huggingface.co/' + chosen, '_blank') }}",
    )

    code_block = gr.Code(language="python", visible=False)

    generate_code.click(
        create_code_block,
        inputs=chosen,
        outputs=code_block
    )

demo.launch()