import gradio as gr
import json

MODELS = [
    {"name": "SmolLM2-135M-Instruct-mobile", "size_mb": 101, "ram_mb": 400, "tps": 59.7, "phone_tps": 46.0},
    {"name": "Qwen2.5-0.5B-Instruct-mobile-int4", "size_mb": 469, "ram_mb": 550, "tps": 12.5, "phone_tps": 23.2},
    {"name": "Llama-3.2-1B-Instruct-Q4-mobile", "size_mb": 770, "ram_mb": 1100, "tps": 11.3, "phone_tps": 5.4},
    {"name": "TinyLlama-1.1B-Chat-Q5-mobile", "size_mb": 800, "ram_mb": 1200, "tps": 10.9, "phone_tps": ""},
    {"name": "Gemma-2-2B-IT-Q5-mobile", "size_mb": 1500, "ram_mb": 2200, "tps": 5.7, "phone_tps": ""},
    {"name": "Llama-3.2-3B-Instruct-Q5-mobile", "size_mb": 2100, "ram_mb": 2700, "tps": 2.6, "phone_tps": ""},
]

def check_ram(ram_mb: int) -> str:
    """Check which dispatchAI models fit in a given RAM budget.
    
    Args:
        ram_mb: Available RAM in MB
    
    Returns:
        JSON with models that fit and their specs
    """
    fitting = [m for m in MODELS if m["ram_mb"] <= ram_mb]
    not_fitting = [m for m in MODELS if m["ram_mb"] > ram_mb]
    
    return json.dumps({
        "ram_budget_mb": ram_mb,
        "models_that_fit": len(fitting),
        "best_option": fitting[0] if fitting else None,
        "all_fitting": fitting,
        "too_large": [{"name": m["name"], "needs_mb": m["ram_mb"]} for m in not_fitting],
    }, indent=2)

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="RAM Calculator") as demo:
    gr.Markdown("# 💾 Phone RAM Calculator (MCP)")
    ram = gr.Slider(512, 8192, value=2048, step=256, label="Available RAM (MB)")
    btn = gr.Button("Check", variant="primary")
    out = gr.Textbox(label="Results (JSON)", lines=15)
    btn.click(fn=check_ram, inputs=ram, outputs=out)
    gr.Markdown("---\n🚀 [dispatchAI](https://huggingface.co/dispatchAI)")
    demo.launch(mcp_server=True)