import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "SupraLabs/Supra-Router-51M"

# Load model + tokenizer at startup on CPU.
# The model is only ~52M parameters, so float32 is tiny (~200MB) and avoids any
# bfloat16-CPU compatibility issues. use_cache is enabled at generation time for
# fast autoregressive decoding.
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)
model.eval()
DEVICE = model.device


def _parse_fields(text: str) -> dict:
    """Parse the pipe-separated routing output into an ordered dict."""
    fields = {}
    for part in text.split("|"):
        part = part.strip()
        if not part:
            continue
        key, sep, val = part.partition(":")
        if sep:
            fields[key.strip().lower()] = val.strip()
    return fields


def _format_result(raw_text: str) -> str:
    """Render the model's structured output as a readable Markdown card."""
    fields = _parse_fields(raw_text)
    route = fields.get("route", "")
    low = route.lower()
    if "big" in low:
        emoji, tag = "☁️", "Cloud / Frontier model"
    elif "small" in low:
        emoji, tag = "🖥️", "Local / Edge SLM"
    else:
        emoji, tag = "🛤️", (route or "—")

    lines = ["## " + emoji + " Routing Decision", ""]
    lines.append(f"**Route:** `{route or 'unknown'}` — {tag}")
    lines.append("")

    rows = []
    for key, label in [
        ("domain", "Domain"),
        ("complexity", "Complexity"),
        ("math", "Math"),
        ("code", "Code"),
    ]:
        if key in fields:
            rows.append(f"| {label} | {fields[key]} |")
    if rows:
        lines.append("| Field | Value |")
        lines.append("| :-- | :-- |")
        lines.extend(rows)
        lines.append("")

    justification = fields.get("justification")
    if justification:
        lines.append(f"**Justification:** {justification}")

    if rows or justification or route:
        return "\n".join(lines)
    # Fallback: show whatever the model produced.
    return "```\n" + raw_text + "\n```"


def route_prompt(user_prompt: str, max_new_tokens: int):
    """Run the router on a user prompt and return (formatted card, raw output)."""
    prompt = (user_prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a prompt to route.")

    # Match the structural framing tokens used during training.
    formatted = f"Task: {prompt}\nAnalysis: "
    inputs = tokenizer(formatted, return_tensors="pt")
    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
            use_cache=True,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
    raw = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
    return _format_result(raw), raw


EXAMPLES = [
    ["Write a movie script about a chef who gets lost at sea.", 128],
    ["Explain how to implement a balanced AVL tree in Python with rotations.", 128],
    ["What is the derivative of x^3 * sin(x) using the product rule?", 128],
    ["Suggest a quick weeknight dinner using chicken and rice.", 128],
]

with gr.Blocks(title="Supra-Router-51M · Infrastructure Routing") as demo:
    gr.Markdown("# 🛤️ Supra-Router-51M")
    gr.Markdown(
        "An ultra-lightweight (51.7M parameter) infrastructure routing model that analyzes a prompt "
        "and decides whether it should be handled by a **local edge SLM** or offloaded to a "
        "**cloud frontier model**. It emits a structured multi-task analysis: domain, complexity, "
        "math/code flags, the routing decision, and a justification."
    )
    gr.Markdown("👉 [Model on Hugging Face](https://huggingface.co/SupraLabs/Supra-Router-51M)")

    with gr.Row():
        with gr.Column():
            prompt_in = gr.Textbox(
                label="Prompt to route",
                lines=4,
                placeholder="e.g. Write a movie script about a chef who gets lost at sea.",
            )
            with gr.Accordion("Advanced settings", open=False):
                max_new = gr.Slider(
                    minimum=16,
                    maximum=256,
                    value=128,
                    step=16,
                    label="Max new tokens",
                    info="Maximum number of tokens the router may generate.",
                )
            route_btn = gr.Button("Route prompt", variant="primary")

        with gr.Column():
            result_md = gr.Markdown(
                "_Enter a prompt and click **Route prompt** to see the analysis._"
            )
            with gr.Accordion("Raw model output", open=False):
                raw_out = gr.Textbox(label="Raw output", lines=6, interactive=False)

    gr.Examples(
        examples=EXAMPLES,
        inputs=[prompt_in, max_new],
        fn=route_prompt,
        outputs=[result_md, raw_out],
        cache_examples=False,
    )

    gr.Markdown(
        "Built with [SpaceFactory](https://huggingface.co/spaces/mrfakename/spacefactory2) · "
        "Model: [SupraLabs/Supra-Router-51M](https://huggingface.co/SupraLabs/Supra-Router-51M)"
    )

    route_btn.click(route_prompt, inputs=[prompt_in, max_new], outputs=[result_md, raw_out])
    prompt_in.submit(route_prompt, inputs=[prompt_in, max_new], outputs=[result_md, raw_out])

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Soft())