import gradio as gr
import json

def detect_chat_format(model_name: str) -> str:
    """Detect the correct chat format for a dispatchAI model.
    
    Args:
        model_name: Model name or repo ID
    
    Returns:
        JSON with the correct chat_format and usage code
    """
    lower = model_name.lower()
    
    if "smollm" in lower or "llama-3" in lower:
        fmt = "llama-3"
    elif "gemma" in lower:
        fmt = "gemma"
    elif "qwen" in lower or "phi" in lower or "tinyllama" in lower or "minicpm" in lower:
        fmt = "chatml"
    else:
        fmt = "chatml"
    
    code = f'''from llama_cpp import Llama

llm = Llama(
    model_path="model.gguf",
    chat_format="{fmt}",
    n_ctx=512,
    n_threads=4,
    verbose=False
)

response = llm.create_chat_completion(
    messages=[{{"role": "user", "content": "Hello!"}}],
    max_tokens=50
)
print(response["choices"][0]["message"]["content"])'''
    
    return json.dumps({
        "model_name": model_name,
        "chat_format": fmt,
        "code": code,
    }, indent=2)

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="Chat Format Detector") as demo:
    gr.Markdown("# 🎯 Chat Format Detector (MCP)")
    inp = gr.Textbox(label="Model Name", placeholder="SmolLM2-135M-Instruct-mobile")
    btn = gr.Button("Detect Format", variant="primary")
    out = gr.Textbox(label="Result (JSON)", lines=15)
    btn.click(fn=detect_chat_format, inputs=inp, outputs=out)
    gr.Markdown("---\n🚀 [dispatchAI](https://huggingface.co/dispatchAI)")
    demo.launch(mcp_server=True)