import gradio as gr
import json

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

def compare_models(model_a: str, model_b: str) -> str:
    """Compare two dispatchAI models side by side.
    
    Args:
        model_a: First model name
        model_b: Second model name
    
    Returns:
        JSON with comparison
    """
    a = next((m for m in MODELS if model_a.lower() in m["name"].lower()), None)
    b = next((m for m in MODELS if model_b.lower() in m["name"].lower()), None)
    
    if not a or not b:
        return json.dumps({"error": "Model not found", "available": [m["name"] for m in MODELS]})
    
    return json.dumps({
        "comparison": {
            "model_a": a,
            "model_b": b,
            "speed_winner": a["name"] if a.get("cpu_tps", 0) > b.get("cpu_tps", 0) else b["name"],
            "size_winner": a["name"] if a.get("size_mb", 0) < b.get("size_mb", 0) else b["name"],
            "speed_diff_pct": round(abs(a.get("cpu_tps",0) - b.get("cpu_tps",0)) / min(a.get("cpu_tps",1), b.get("cpu_tps",1)) * 100, 1),
        }
    }, indent=2)

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="Model Compare") as demo:
    gr.Markdown("# ⚖️ Model Comparison (MCP)")
    with gr.Row():
        a = gr.Textbox(label="Model A", placeholder="SmolLM2-135M")
        b = gr.Textbox(label="Model B", placeholder="Llama-3.2-1B")
    btn = gr.Button("Compare", variant="primary")
    out = gr.Textbox(label="Comparison (JSON)", lines=15)
    btn.click(fn=compare_models, inputs=[a, b], outputs=out)
    gr.Markdown("---\n🚀 [dispatchAI](https://huggingface.co/dispatchAI)")
    demo.launch(mcp_server=True)