import gradio as gr

MODELS = [
    {"Model": "SmolLM2-135M-Instruct-mobile", "Params": "135M", "Size_MB": 270, "RAM_MB": 400, "Task": "Chat", "Quant": "FP16", "Speed_tps": 25.5},
    {"Model": "SmolLM2-360M-Instruct-mobile", "Params": "360M", "Size_MB": 720, "RAM_MB": 700, "Task": "Chat", "Quant": "FP16", "Speed_tps": 21.0},
    {"Model": "Qwen2.5-0.5B-Instruct-mobile-int4", "Params": "500M", "Size_MB": 350, "RAM_MB": 550, "Task": "Chat", "Quant": "INT4", "Speed_tps": 20.0},
    {"Model": "Llama-3.2-1B-Instruct-Q4-mobile", "Params": "1B", "Size_MB": 700, "RAM_MB": 1100, "Task": "Chat", "Quant": "Q4", "Speed_tps": 18.2},
    {"Model": "Llama-3.2-1B-Instruct-Q6-mobile", "Params": "1B", "Size_MB": 1100, "RAM_MB": 1300, "Task": "Chat", "Quant": "Q6", "Speed_tps": 16.8},
    {"Model": "TinyLlama-1.1B-Chat-Q5-mobile", "Params": "1.1B", "Size_MB": 800, "RAM_MB": 1200, "Task": "Chat", "Quant": "Q5", "Speed_tps": 17.5},
    {"Model": "Qwen2.5-0.5B-Coder-mobile", "Params": "500M", "Size_MB": 1000, "RAM_MB": 1500, "Task": "Code", "Quant": "FP16", "Speed_tps": 20.0},
    {"Model": "Qwen2.5-Coder-1.5B-mobile", "Params": "1.5B", "Size_MB": 3000, "RAM_MB": 4000, "Task": "Code", "Quant": "FP16", "Speed_tps": 10.5},
    {"Model": "Qwen2.5-Math-1.5B-mobile", "Params": "1.5B", "Size_MB": 3000, "RAM_MB": 4000, "Task": "Math", "Quant": "FP16", "Speed_tps": 10.5},
    {"Model": "Gemma-2B-Arabic-mobile", "Params": "2B", "Size_MB": 5000, "RAM_MB": 5500, "Task": "Arabic", "Quant": "FP16", "Speed_tps": 8.0},
    {"Model": "Gemma-2-2B-IT-Q5-mobile", "Params": "2B", "Size_MB": 1500, "RAM_MB": 2200, "Task": "Chat", "Quant": "Q5", "Speed_tps": 12.0},
    {"Model": "Llama-3.2-3B-Instruct-Q5-mobile", "Params": "3B", "Size_MB": 2100, "RAM_MB": 2700, "Task": "Chat", "Quant": "Q5", "Speed_tps": 8.5},
    {"Model": "Llama-3.2-1B-FunctionCall-mobile", "Params": "1B", "Size_MB": 2500, "RAM_MB": 3000, "Task": "Function Call", "Quant": "FP16", "Speed_tps": 12.0},
    {"Model": "Moondream2-Vision-Q5-mobile", "Params": "1.9B", "Size_MB": 1400, "RAM_MB": 2000, "Task": "Vision", "Quant": "Q5", "Speed_tps": 8.5},
    {"Model": "EmbeddingGemma-300M-Q8-mobile", "Params": "300M", "Size_MB": 300, "RAM_MB": 500, "Task": "Embedding", "Quant": "Q8", "Speed_tps": 22.0},
]

def recommend_model(ram_mb: int, task: str = "Any", priority: str = "Smallest size") -> str:
    """Recommend the best mobile-optimized LLM for a given RAM budget and task.
    
    Use this tool when a user asks "which model should I use on my phone" or 
    "find a small model for my device". Returns model name, size, and speed.
    
    Args:
        ram_mb: Available RAM in megabytes (e.g., 2048 for 2GB, 4096 for 4GB)
        task: The primary task - one of: Chat, Code, Math, Arabic, Function Call, Vision, Embedding, Any
        priority: What to optimize for - "Smallest size", "Fastest", or "Best quality"
    
    Returns:
        JSON string with recommended model details
    """
    import json
    filtered = [m for m in MODELS if m["RAM_MB"] <= ram_mb]
    if task != "Any":
        filtered = [m for m in filtered if m["Task"] == task]
    
    if not filtered:
        return json.dumps({"error": f"No models fit in {ram_mb}MB RAM for task '{task}'"})
    
    if priority == "Smallest size":
        filtered.sort(key=lambda x: x["Size_MB"])
    elif priority == "Fastest":
        filtered.sort(key=lambda x: x["Speed_tps"], reverse=True)
    elif priority == "Best quality":
        filtered.sort(key=lambda x: x["Params"], reverse=True)
    
    best = filtered[0]
    return json.dumps({
        "recommended_model": best["Model"],
        "huggingface_url": f"https://huggingface.co/dispatchAI/{best['Model']}",
        "params": best["Params"],
        "file_size_mb": best["Size_MB"],
        "ram_required_mb": best["RAM_MB"],
        "quantization": best["Quant"],
        "expected_speed_tps": best["Speed_tps"],
        "task": best["Task"],
        "alternatives": [{"model": m["Model"], "size_mb": m["Size_MB"]} for m in filtered[1:4]]
    }, indent=2)

def list_all_models() -> str:
    """List all 39 dispatchAI mobile-optimized models with their specs.
    
    Returns a catalog of all available models with sizes, RAM requirements, and speeds.
    Use when a user wants to browse all available mobile models.
    
    Returns:
        JSON string with all model details
    """
    import json
    return json.dumps({
        "total_models": len(MODELS),
        "models": MODELS,
        "org_url": "https://huggingface.co/dispatchAI"
    }, indent=2)

with gr.Blocks(title="dispatchAI Model Recommender MCP") as demo:
    gr.Markdown("## 📱 dispatchAI Model Recommender (MCP Tool)")
    gr.Markdown("This Space is an MCP server. Add it to Claude Desktop, Cursor, or any MCP client.")
    
    with gr.Row():
        ram = gr.Slider(512, 8192, value=2048, step=256, label="RAM (MB)")
        task = gr.Dropdown(["Any", "Chat", "Code", "Math", "Arabic", "Function Call", "Vision", "Embedding"], value="Any", label="Task")
        prio = gr.Radio(["Smallest size", "Fastest", "Best quality"], value="Smallest size", label="Priority")
    
    btn = gr.Button("Recommend", variant="primary")
    out = gr.Textbox(label="Recommendation (JSON)", lines=12)
    btn.click(fn=recommend_model, inputs=[ram, task, prio], outputs=out)
    
    list_btn = gr.Button("List All Models")
    list_out = gr.Textbox(label="All Models (JSON)", lines=15)
    list_btn.click(fn=list_all_models, outputs=list_out)

demo.launch(mcp_server=True)