import gradio as gr
from huggingface_hub import model_info, hf_hub_url
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from netron import build_netron_iframe, is_netron_file

current_repo = None

def repo_change(repo: str, oauth_token: str) -> dict[str, any] | None:
    global current_repo 
    current_repo = repo

    choices = [] 

    if len(repo) == 0 or repo.find("/") == -1:
        return gr.update(choices=choices, value=None, interactive=False)

    try:
        # get the files from the model repo
        info = model_info(repo, token=oauth_token)

        # filter only netron capable files
        for sibling in info.siblings:
            if is_netron_file(sibling.rfilename):
                choices.append(sibling.rfilename)

        return gr.update(choices=choices, value=None, interactive=True)
    except Exception as e:
        print("Exception", e)
        return None

def viewer_change(file):
    if current_repo is None:
        return gr.update(value="<p>Repo not selected</p>")

    file_url = hf_hub_url(current_repo, file)
    if file_url is None: 
        return gr.update(value="<p>File not found</p>")

    return gr.update(value=build_netron_iframe(file_url))

with gr.Blocks() as viewer_app:
    with gr.Row():
        gr.Markdown("# Netron Model Viewer") 
        gr.LoginButton(min_width=250)

    with gr.Column():
        with gr.Row():
            model_in = HuggingfaceHubSearch(
                label="Model ID",
                placeholder="Search for models on Huggingface",
                search_type="model",
            )
        files = gr.Dropdown(choices=[], interactive=False, label="File")

    viewer = gr.HTML("""<p>Enter the model ID and select the compatible model file to view the model's graph view.</p>""")

    gr.on([model_in.submit, model_in.change], fn=repo_change, inputs=[model_in], outputs=[files], trigger_mode="always_last")
    files.select(fn=viewer_change, inputs=[files], outputs=[viewer])

if __name__ == "__main__":
    viewer_app.launch(share=True)