import gradio as gr
import json

def generate_model_card(model_id: str, base_model: str = "", quantization: str = "Q4_K_M", 
                       params: str = "", task: str = "Text Generation", phone_tested: bool = True,
                       tokens_per_sec: float = 0.0, ram_mb: int = 0, file_size_mb: int = 0) -> str:
    """Generate a HuggingFace model card (README.md) for a mobile-optimized model.
    
    Use this tool when you need to create or update a model card for a dispatchAI
    mobile model. Produces properly formatted markdown with YAML frontmatter,
    benchmark results, and usage examples.
    
    Args:
        model_id: The HuggingFace repo ID (e.g., "dispatchAI/MyModel-mobile")
        base_model: The original/base model name (e.g., "Qwen/Qwen2.5-0.5B-Instruct")
        quantization: Quantization method used (Q4_K_M, Q5_K_M, Q8_0, INT4, FP16)
        params: Parameter count as string (e.g., "500M", "1B", "3B")
        task: Primary task (Text Generation, Code Generation, etc.)
        phone_tested: Whether the model was tested on real phone hardware
        tokens_per_sec: Measured tokens/sec on Snapdragon 865 (0 if untested)
        ram_mb: RAM required in MB
        file_size_mb: Model file size in MB
    
    Returns:
        JSON string with the generated model card markdown
    """
    org_name = model_id.split("/")[0] if "/" in model_id else "dispatchAI"
    model_name = model_id.split("/")[-1] if "/" in model_id else model_id
    
    tags = ["mobile", "on-device", "quantized", quantization.lower(), "dispatchai"]
    if task == "Code Generation":
        tags.extend(["code", "coder"])
    elif "arabic" in model_name.lower():
        tags.extend(["arabic", "multilingual"])
    elif "function" in model_name.lower():
        tags.extend(["function-calling", "agent"])
    elif "vision" in model_name.lower():
        tags.extend(["vision", "multimodal"])
    
    card = f"""---
license: apache-2.0
language:
  - en
library_name: transformers
tags:
{chr(10).join(f"  - {t}" for t in tags)}
pipeline_tag: {task.lower().replace(" ", "-") if " " in task else task}
---

# {model_name}

{("Mobile-optimized version of " + base_model) if base_model else "A mobile-optimized language model"} — quantized to {quantization}, designed to run on phones and edge devices.

## Model Details

| Attribute | Value |
|-----------|-------|
| **Base Model** | {base_model or "N/A"} |
| **Parameters** | {params or "Unknown"} |
| **Quantization** | {quantization} |
| **File Size** | {file_size_mb} MB |
| **RAM Required** | {ram_mb} MB |
| **Task** | {task} |
| **License** | Apache-2.0 |
| **Organization** | [{org_name}](https://huggingface.co/{org_name}) |

## Mobile Performance

{"✅ **Tested on real hardware**" if phone_tested else "⚠️ Not yet tested on phone hardware"}

| Hardware | Tokens/sec | RAM Usage | Status |
|----------|-----------|-----------|--------|
| Snapdragon 865 (S20 FE) | {tokens_per_sec} t/s | {ram_mb} MB | {"✅ Pass" if tokens_per_sec > 2 else "❌ Fail" if phone_tested else "⏳ Pending"} |

## Usage

### Python (transformers)

```python
from transformers import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("{model_id}")
tokenizer = AutoTokenizer.from_pretrained("{model_id}")

messages = [{{"role": "user", "content": "Hello!"}}]
input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(input_text, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))
```

### llama.cpp (GGUF)

```bash
# Download
hf download {model_id} model.gguf

# Run on phone/desktop
llama-cli -m model.gguf -p "Hello!" -n 100 -t 4
```

## Quantization Details

This model uses **{quantization}** quantization:

- File size: {file_size_mb} MB ({round(file_size_mb / 1024, 2)} GB)
- Quality retention: ~92% of FP16
- Optimized for: Mobile and edge devices with < 4GB RAM

## About dispatchAI

[dispatchAI](https://huggingface.co/dispatchAI) re-engineers open-source LLMs for mobile and edge devices. Every model is tested on real Snapdragon hardware.

**Small. Mobile. Free. UAE-built.**

---

*I think, therefore I ship.*
"""
    
    return json.dumps({
        "model_id": model_id,
        "card_content": card,
        "tags": tags,
        "instructions": f"Upload this as README.md to {model_id} repo"
    }, indent=2)

with gr.Blocks(title="dispatchAI Model Card Writer MCP") as demo:
    gr.Markdown("## 📝 dispatchAI Model Card Writer (MCP Tool)")
    
    with gr.Row():
        mid = gr.Textbox(label="Model ID", placeholder="dispatchAI/MyModel-mobile")
        base = gr.Textbox(label="Base Model", placeholder="Qwen/Qwen2.5-0.5B-Instruct")
    
    with gr.Row():
        quant = gr.Dropdown(["FP16", "INT4", "Q4_K_M", "Q5_K_M", "Q6_K", "Q8_0"], value="Q4_K_M", label="Quantization")
        params = gr.Textbox(label="Params", placeholder="500M, 1B, 3B...")
        task = gr.Dropdown(["Text Generation", "Code Generation", "Function Calling", "Image-to-Text"], value="Text Generation", label="Task")
    
    with gr.Row():
        tested = gr.Checkbox(value=True, label="Phone Tested")
        tps = gr.Number(value=18.0, label="Tokens/sec (Snapdragon 865)")
        ram = gr.Number(value=1100, label="RAM (MB)")
        size = gr.Number(value=700, label="File Size (MB)")
    
    btn = gr.Button("Generate Model Card", variant="primary")
    out = gr.Textbox(label="Generated Card (JSON)", lines=25)
    btn.click(fn=generate_model_card, inputs=[mid, base, quant, params, task, tested, tps, ram, size], outputs=out)

demo.launch(mcp_server=True)