"""Gemma 4 E2B (QAT Mobile) multimodal chat demo for HF Spaces ZeroGPU."""

# `spaces` MUST be imported before torch / transformers: it monkey-patches
# torch.cuda.*, which must happen before CUDA is initialized in this process.
import spaces  # noqa: F401

import os
import threading
from collections.abc import Iterator
from typing import Any

import torch
import gradio as gr
from transformers import (
    AutoProcessor,
    AutoModelForMultimodalLM,
    TextIteratorStreamer,
)

MODEL_ID = os.environ.get("MODEL_ID", "unsloth/gemma-4-E2B-it-qat-mobile")
# HF_TOKEN is only required if MODEL_ID points at a gated repo. This model is
# public, so it can be left unset.
HF_TOKEN = os.environ.get("HF_TOKEN")

DEFAULT_SYSTEM = (
    "You are Gemma 4, a friendly, helpful multimodal assistant. "
    "Be concise and accurate."
)

# Load eagerly at module scope. On ZeroGPU this call is intercepted: weights are
# packed to disk here, then streamed into VRAM on the first @spaces.GPU entry.
print(f"[app] loading model: {MODEL_ID}", flush=True)
processor = AutoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = AutoModelForMultimodalLM.from_pretrained(
    MODEL_ID,
    dtype="auto",
    device_map="cuda",
    token=HF_TOKEN,
)
model.eval()
print("[app] model ready", flush=True)


def _user_content_blocks(value: Any) -> list[dict]:
    """Turn a Gradio multimodal user value into Gemma 4 content blocks.

    Gradio hands us either a plain string (text-only turn) or a dict with
    ``{"text": str, "files": [paths]}``. Gemma 4 expects images *before* text.
    """
    if value is None:
        return []
    if isinstance(value, str):
        text = value.strip()
        return [{"type": "text", "text": text}] if text else []
    if isinstance(value, dict):
        blocks: list[dict] = []
        for path in value.get("files") or []:
            blocks.append({"type": "image", "url": path})
        text = (value.get("text") or "").strip()
        if text:
            blocks.append({"type": "text", "text": text})
        return blocks
    return []


def _build_messages(
    history: list[dict],
    message: dict,
    system_prompt: str,
    enable_thinking: bool,
) -> list[dict]:
    messages: list[dict] = [
        {"role": "system", "content": system_prompt.strip() or DEFAULT_SYSTEM}
    ]
    for turn in history:
        role = turn.get("role")
        content = turn.get("content")
        if role == "user":
            blocks = _user_content_blocks(content)
            if blocks:
                messages.append({"role": "user", "content": blocks})
        elif role == "assistant" and isinstance(content, str) and content.strip():
            # Per Gemma 4 best practice: prior assistant turns carry only the
            # final answer, never the thinking trace.
            messages.append({"role": "assistant", "content": content})
    blocks = _user_content_blocks(message)
    messages.append({"role": "user", "content": blocks})
    return messages


@spaces.GPU(duration=120)
def respond(
    message: dict,
    history: list[dict],
    system_prompt: str,
    enable_thinking: bool,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
) -> Iterator[str]:
    """Stream a multimodal chat completion from Gemma 4 E2B.

    Yielded strings are the growing assistant reply (shown live in the UI).
    """
    user_blocks = _user_content_blocks(message)
    if not user_blocks:
        yield "Please enter a message or attach an image first."
        return

    messages = _build_messages(history, message, system_prompt, enable_thinking)

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        add_generation_prompt=True,
        enable_thinking=enable_thinking,
    ).to(model.device)

    streamer = TextIteratorStreamer(
        processor.tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    do_sample = temperature > 0
    gen_kwargs = dict(
        inputs,
        max_new_tokens=int(max_new_tokens),
        top_p=float(top_p),
        top_k=int(top_k),
        do_sample=do_sample,
        streamer=streamer,
    )
    if do_sample:
        gen_kwargs["temperature"] = float(temperature)

    worker = threading.Thread(target=model.generate, kwargs=gen_kwargs)
    worker.start()

    partial = ""
    try:
        for chunk in streamer:
            partial += chunk
            yield partial
    finally:
        worker.join()


demo = gr.ChatInterface(
    fn=respond,
    multimodal=True,
    title="Gemma 4 E2B (QAT Mobile) \u00b7 Multimodal Demo",
    description=(
        "Chat with `unsloth/gemma-4-E2B-it-qat-mobile` \u2014 a 2.3B-effective-parameter "
        "multimodal Gemma 4 model (text + image). Upload an image and ask about it, "
        "or just chat. The first reply after a cold start may be slow while ZeroGPU "
        "streams the weights into VRAM."
    ),
    additional_inputs=[
        gr.Textbox(
            value=DEFAULT_SYSTEM,
            label="System prompt",
            lines=2,
        ),
        gr.Checkbox(value=False, label="Enable thinking mode"),
        gr.Slider(64, 2048, value=512, step=64, label="Max new tokens"),
        gr.Slider(0.0, 1.5, value=1.0, step=0.05, label="Temperature (0 = greedy)"),
        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p"),
        gr.Slider(0, 128, value=64, step=1, label="Top-k"),
    ],
    examples=[
        [{"text": "Write a short poem about a curious cat.", "files": []}],
        [{"text": "Explain quantization-aware training in two sentences.", "files": []}],
        [{"text": "Give me three name ideas for a friendly code-review bot.", "files": []}],
    ],
    cache_examples=False,
    fill_height=True,
)

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