import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  -- must be before torch
import torch
import gradio as gr
from transformers import AutoModelForImageTextToText, AutoProcessor


MODEL_ID = "ProCreations/grug-9b"
PROCESSOR_ID = "deepreinforce-ai/Ornith-1.0-9B"  # grug-9b omits preprocessor configs; base model has them

print(f"Loading processor from {PROCESSOR_ID} …")
processor = AutoProcessor.from_pretrained(PROCESSOR_ID)

print(f"Loading model from {MODEL_ID} …")
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()
print("Model loaded.")


def _build_messages(history, image_path, user_text):
    """Build the messages list from chat history + new user input."""
    messages = []
    for msg in history:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        if isinstance(content, list):
            messages.append({"role": role, "content": content})
        else:
            messages.append({"role": role, "content": [{"type": "text", "text": str(content)}]})
    # Add the new user message
    user_content = []
    if image_path is not None:
        user_content.append({"type": "image", "image": image_path})
    user_content.append({"type": "text", "text": user_text})
    messages.append({"role": "user", "content": user_content})
    return messages


@spaces.GPU(duration=120)
def predict(image_path, user_text, max_new_tokens, temperature, top_p, enable_thinking):
    """Run a single-turn vision+text inference and return the response.

    Args:
        image_path: path to the uploaded image (or None for text-only).
        user_text: the user's text prompt.
        max_new_tokens: maximum number of tokens to generate.
        temperature: sampling temperature (1.0 = greedy-ish).
        top_p: nucleus sampling probability.
        enable_thinking: whether to emit <think> reasoning before the answer.

    Returns:
        The decoded text response.
    """
    if not user_text.strip() and image_path is None:
        return "Please provide some text or an image to analyze."

    messages = _build_messages([], image_path, user_text if user_text.strip() else "Describe this image.")

    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=enable_thinking,
    )

    # Build image inputs from the messages
    image_inputs = []
    for msg in messages:
        for item in msg.get("content", []):
            if isinstance(item, dict) and item.get("type") == "image":
                from PIL import Image
                image_inputs.append(Image.open(item["image"]).convert("RGB"))

    inputs = processor(
        text=[text],
        images=image_inputs if image_inputs else None,
        padding=True,
        return_tensors="pt",
    ).to("cuda")

    do_sample = temperature > 0.001
    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=do_sample,
            temperature=float(temperature) if do_sample else 1.0,
            top_p=float(top_p) if do_sample else 1.0,
        )

    # Strip the input tokens from the output
    generated = output_ids[0][inputs["input_ids"].shape[1]:]
    result = processor.decode(generated, skip_special_tokens=True, clean_up_tokenization_spaces=False)
    return result


@spaces.GPU(duration=120)
def chat_predict(message, history, image, max_new_tokens, temperature, top_p, enable_thinking):
    """Multi-turn chat with optional image. Gradio passes history as a list of
    [user, assistant] tuples.

    Args:
        message: the user's latest text message.
        history: list of (user_msg, assistant_msg) tuples.
        image: optional uploaded image path.
        max_new_tokens: maximum tokens to generate.
        temperature: sampling temperature.
        top_p: nucleus sampling probability.
        enable_thinking: whether to emit <think> reasoning.

    Returns:
        The assistant's text response.
    """
    if not message.strip() and image is None:
        return "Please provide a message or an image."

    # Convert Gradio history format to messages
    messages = []
    for user_msg, assistant_msg in history:
        # Rebuild each prior turn as a content list
        user_content = []
        # We can't perfectly reconstruct which prior messages had images,
        # so we store images as text references in history
        user_content.append({"type": "text", "text": user_msg})
        messages.append({"role": "user", "content": user_content})
        if assistant_msg:
            messages.append({"role": "assistant", "content": [{"type": "text", "text": assistant_msg}]})

    # Add the current message
    user_content = []
    if image is not None:
        user_content.append({"type": "image", "image": image})
    user_content.append({"type": "text", "text": message if message.strip() else "Describe this image."})
    messages.append({"role": "user", "content": user_content})

    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=enable_thinking,
    )

    # Collect images from messages
    image_inputs = []
    for msg in messages:
        for item in msg.get("content", []):
            if isinstance(item, dict) and item.get("type") == "image":
                from PIL import Image
                image_inputs.append(Image.open(item["image"]).convert("RGB"))

    inputs = processor(
        text=[text],
        images=image_inputs if image_inputs else None,
        padding=True,
        return_tensors="pt",
    ).to("cuda")

    do_sample = temperature > 0.001
    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=do_sample,
            temperature=float(temperature) if do_sample else 1.0,
            top_p=float(top_p) if do_sample else 1.0,
        )

    generated = output_ids[0][inputs["input_ids"].shape[1]:]
    result = processor.decode(generated, skip_special_tokens=True, clean_up_tokenization_spaces=False)
    return result


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown(
        """
        # 🪨 Grug-9B Vision-Language Demo

        **ProCreations/grug-9b** — a 9B-parameter reasoning VLM (fine-tuned from Ornith-1.0-9B / Qwen3.5)
        that "thinks small" — producing concise reasoning instead of verbose chain-of-thought.

        Upload an image and ask a question, or just type a prompt. The model will respond
        with a short reasoning trace followed by its answer.

        [Model card](https://huggingface.co/ProCreations/grug-9b) · [Base model](https://huggingface.co/deepreinforce-ai/Ornith-1.0-9B)
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            image_input = gr.Image(
                label="Upload Image (optional)",
                type="filepath",
                height=320,
            )
            text_input = gr.Textbox(
                label="Prompt",
                placeholder="Ask something about the image, or type a prompt…",
                lines=3,
            )
            submit_btn = gr.Button("Submit", variant="primary")

        with gr.Column(scale=4):
            output_text = gr.Textbox(
                label="Response",
                lines=16,
                max_lines=30,
                show_copy_button=True,
            )

    with gr.Accordion("Advanced settings", open=False):
        with gr.Row():
            max_tokens = gr.Slider(
                label="Max new tokens", minimum=64, maximum=2048, value=512, step=64,
            )
            temperature = gr.Slider(
                label="Temperature", minimum=0.0, maximum=2.0, value=0.7, step=0.05,
            )
            top_p = gr.Slider(
                label="Top-p", minimum=0.1, maximum=1.0, value=0.9, step=0.05,
            )
            thinking = gr.Checkbox(label="Enable thinking (<think> tag)", value=True)

    gr.Examples(
        examples=[
            ["examples/astronaut.jpg", "What is happening in this image?", 512, 0.7, 0.9, True],
            ["examples/cat_tabby.jpg", "Describe this cat in detail.", 512, 0.7, 0.9, True],
            ["examples/bird_bee_eater.jpg", "What species is this bird? What is it doing?", 512, 0.7, 0.9, True],
        ],
        inputs=[image_input, text_input, max_tokens, temperature, top_p, thinking],
        outputs=output_text,
        fn=predict,
        cache_examples=True,
        cache_mode="lazy",
    )

    submit_btn.click(
        fn=predict,
        inputs=[image_input, text_input, max_tokens, temperature, top_p, thinking],
        outputs=output_text,
        api_name="predict",
    )

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