import spaces  # must be imported before torch / transformers

import re
import threading

import torch
import gradio as gr
from transformers import AutoProcessor, BitsAndBytesConfig, TextIteratorStreamer

try:
    from transformers import AutoModelForMultimodalLM as AutoVLM
except ImportError:
    from transformers import AutoModelForImageTextToText as AutoVLM

MODEL_ID = "meta-models/Muse-Glimmer-30B"

IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp")

# NF4 keeps the 30B model ~17 GB (model card: ~1% degradation), so it fits
# ZeroGPU's default 48 GB slice with ample KV-cache headroom.
bnb = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoVLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb,
    device_map="cuda",  # bitsandbytes' loader is ZeroGPU-aware
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).eval()


def _content_parts(text, files):
    parts = [{"type": "image", "path": f} for f in files if f.lower().endswith(IMAGE_EXTS)]
    if text:
        parts.append({"type": "text", "text": text})
    return parts


def _history_to_messages(history):
    messages = []
    for m in history:
        role, content = m["role"], m["content"]
        meta = m.get("metadata") or {}
        if role == "assistant" and meta.get("title"):
            continue  # drop surfaced reasoning blocks from the context
        if isinstance(content, str):
            if content:
                messages.append({"role": role, "content": [{"type": "text", "text": content}]})
        elif isinstance(content, (tuple, list)):
            parts = _content_parts("", [str(f) for f in content])
            if parts:
                messages.append({"role": role, "content": parts})
        elif isinstance(content, dict) and content.get("path"):
            parts = _content_parts("", [content["path"]])
            if parts:
                messages.append({"role": role, "content": parts})
    return messages


def _parse_channels(raw):
    """Split raw generation into (reasoning, answer).

    The model emits channelled turns: `to=self<|message|>…<|eom|>` for its
    reasoning, then `<|start|>assistant to=user<|message|>…<|eot|>`.
    """
    thinking, answer = [], []
    segments = raw.split("<|start|>")
    for i, seg in enumerate(segments):
        if "<|message|>" not in seg:
            continue
        header, content = seg.split("<|message|>", 1)
        content = content.replace("<|eot|>", "").replace("<|eom|>", "")
        if i == len(segments) - 1:
            # hide a marker that is still streaming in, e.g. a trailing "<|eo"
            content = re.sub(r"<\|?[^|>]*$", "", content)
        if "to=self" in header:
            thinking.append(content)
        else:
            answer.append(content)
    return "".join(thinking).strip(), "".join(answer).strip()


def _duration(message, history, reasoning_strength="high", max_new_tokens=1024, *args, **kwargs):
    return min(300, 60 + int(max_new_tokens) // 8)


@spaces.GPU(duration=_duration)
def chat(message, history, reasoning_strength="high", max_new_tokens=1024, temperature=1.0):
    """Chat with Muse Glimmer-30B, a multimodal (image + text) agentic model.

    Args:
        message: Dict with 'text' (the user prompt) and 'files' (optional image paths).
        history: Prior chat messages.
        reasoning_strength: How hard the model thinks first: low, medium, high, or xhigh.
        max_new_tokens: Upper bound on generated tokens (reasoning + answer).
        temperature: Sampling temperature (model card recommends 1.0).
    """
    messages = _history_to_messages(history)
    parts = _content_parts(message.get("text", ""), message.get("files") or [])
    if not parts:
        raise gr.Error("Please enter a message or attach an image.")
    messages.append({"role": "user", "content": parts})

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        reasoning_strength=reasoning_strength,
    ).to("cuda")

    streamer = TextIteratorStreamer(
        processor.tokenizer, skip_prompt=True, skip_special_tokens=False
    )
    kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=int(max_new_tokens),
        do_sample=temperature > 0,
        temperature=float(temperature) if temperature > 0 else None,
        top_p=0.95,
        top_k=64,
    )
    thread = threading.Thread(target=model.generate, kwargs=kwargs, daemon=True)
    thread.start()

    raw = ""
    for chunk in streamer:
        raw += chunk
        thinking, answer = _parse_channels(raw)
        out = []
        if thinking:
            out.append(
                gr.ChatMessage(
                    role="assistant",
                    content=thinking,
                    metadata={"title": "🧠 Reasoning", "status": "done" if answer else "pending"},
                )
            )
        if answer or not thinking:
            out.append(gr.ChatMessage(role="assistant", content=answer))
        yield out
    thread.join()


demo = gr.ChatInterface(
    fn=chat,
    multimodal=True,
    type="messages",
    title="Muse Glimmer-30B",
    description=(
        "Chat with [meta-models/Muse-Glimmer-30B](https://huggingface.co/meta-models/Muse-Glimmer-30B) — "
        "a 30B multimodal agentic model (Apache 2.0). Attach images (screenshots, charts, documents) "
        "or just talk to it. Running in 4-bit NF4 on ZeroGPU; its reasoning streams into the "
        "collapsible 🧠 block before the answer."
    ),
    textbox=gr.MultimodalTextbox(
        file_types=["image"], file_count="multiple", placeholder="Ask something or drop an image…"
    ),
    additional_inputs=[
        gr.Dropdown(
            ["low", "medium", "high", "xhigh"],
            value="high",
            label="Reasoning strength",
            info="How hard the model thinks before answering",
        ),
        gr.Slider(256, 4096, value=1024, step=64, label="Max new tokens"),
        gr.Slider(0.0, 1.5, value=1.0, step=0.05, label="Temperature"),
    ],
    examples=[
        [{"text": "Plan a 3-step approach to debug a web server that returns 502 errors intermittently.", "files": []}],
        [{"text": "What makes a good agentic model different from a chat model?", "files": []}],
        [{"text": "Écris un haïku sur les agents autonomes.", "files": []}],
    ],
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch(mcp_server=True)