import gradio as gr
import logging
from optimum.neuron import utils
import os

# Set up logging
logging.basicConfig(level=logging.INFO)

def get_model_info(model_id):
    """Retrieve configuration details for a specific model."""
    if not model_id:
        return "No model selected"
    config_list = utils.get_hub_cached_entries(model_id=model_id, mode="inference")
    return config_list

def get_model_list():
    """Retrieve list of cached models."""
    return utils.get_hub_cached_models(mode="inference")

def create_model_list_interface():
    """Create a Gradio interface with model list as clickable buttons."""
    # Get the model list and sort alphabetically by full model ID
    models = sorted(get_model_list(), key=lambda x: (x[1] + x[2]).lower())

    print(models)
    
    # Prepare model IDs for buttons
    model_ids = [f"{org}/{name}" for (arch, org, name) in models]

    # Create the interface
    with gr.Blocks() as iface:
        with gr.Tab("Model List"):
            # Buttons
            with gr.Row():
                buttons = gr.Radio(
                    choices=model_ids, 
                    label="Select a Model"
                )
            
            # Output for model details
            output = gr.Textbox(label="Model Configuration", lines=20)

            # Connect model selection to output
            buttons.change(
                fn=get_model_info, 
                inputs=buttons, 
                outputs=output
            )

    return iface

# Launch the interface
if __name__ == "__main__":
    iface = create_model_list_interface()
    iface.launch()