import gradio as gr
import pandas as pd
import json

with open("models.json") as f:
    MODELS = json.load(f)

df = pd.DataFrame(MODELS)

def filter_models(family: str, min_speed: float, max_size: int):
    filtered = df.copy()
    if family != "All":
        filtered = filtered[filtered["name"].str.contains(family, case=False)]
    if min_speed > 0:
        filtered = filtered[filtered["cpu_tps"] >= min_speed]
    if max_size < 99999:
        filtered = filtered[filtered["size_mb"] <= max_size]
    return filtered.sort_values("cpu_tps", ascending=False)

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="Model Browser") as demo:
    gr.Markdown("# 🗂️ dispatchAI Model Browser")
    gr.Markdown(f"Browse all {len(MODELS)} verified working models.")
    
    with gr.Row():
        family = gr.Dropdown(["All", "SmolLM", "Llama", "Qwen", "Gemma", "Phi", "TinyLlama", "MiniCPM"], value="All", label="Family")
        min_speed = gr.Slider(0, 60, value=0, step=1, label="Min Speed (t/s)")
        max_size = gr.Slider(100, 10000, value=10000, step=100, label="Max Size (MB)")
    
    table = gr.Dataframe(df.sort_values("cpu_tps", ascending=False))
    
    for inp in [family, min_speed, max_size]:
        inp.change(fn=filter_models, inputs=[family, min_speed, max_size], outputs=table)
    
    gr.Markdown("---\n🚀 [dispatchAI](https://huggingface.co/dispatchAI)")
    demo.launch()