import os

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

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import gradio as gr
from PIL import Image
from threading import Thread
from transformers import (
    AutoModelForImageTextToText,
    AutoProcessor,
    TextIteratorStreamer,
)

MODEL_ID = "SpatialAxiom/SpatialAxiom-9B"

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval()


def _strip_think(text: str) -> str:
    """Remove thinking-block wrapper if present (model should not produce it with thinking off)."""
    if "</think>" in text:
        return text.split("</think>", 1)[1].lstrip()
    return text


@spaces.GPU(duration=90)
def answer(
    image: str | None,
    video: str | None,
    question: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    progress: gr.Progress = gr.Progress(),
):
    """Answer a spatial reasoning question about an image or video.

    Args:
        image: An image file (room interior, scene, etc.).
        video: A video file (walkthrough, embodied video, etc.).
        question: A spatial reasoning question about the visual input.
        max_new_tokens: Maximum number of tokens to generate.
        temperature: Sampling temperature (0 = greedy).
        top_p: Nucleus sampling probability.

    Returns:
        The model's text answer to the spatial reasoning question.
    """
    if not question or not question.strip():
        raise gr.Error("Please enter a question.")
    if image is None and video is None:
        raise gr.Error("Please upload an image or video.")

    # Build the message content
    content = []
    if image is not None:
        img = Image.open(image).convert("RGB")
        content.append({"type": "image", "image": img})
    if video is not None:
        content.append({"type": "video", "url": video})
    content.append({"type": "text", "text": question})

    messages = [{"role": "user", "content": content}]

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

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

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

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

    output = ""
    for token in streamer:
        output += token
        yield _strip_think(output)

    thread.join()
    yield _strip_think(output).strip()


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

EXAMPLES = [
    ["example_images/living_room_blue_couch.jpg", "If I stand at the door facing the bed, is the chair to my left or right?", 512, 0.0, 0.95],
    ["example_images/living_room_minimalist.jpg", "Describe the spatial layout of this room. What objects are on the left, center, and right?", 512, 0.0, 0.95],
    ["example_images/library_interior.jpg", "How many bookshelves can you see? Describe their spatial arrangement relative to the seating area.", 512, 0.0, 0.95],
    ["example_images/cafe_interior.jpg", "If I'm sitting at the table in the foreground, what is behind me and to my sides?", 512, 0.0, 0.95],
]


def _answer_example(image, question, max_new_tokens, temperature, top_p):
    """Wrapper for gr.Examples so the video input can be left at its default (None)."""
    yield from answer(image, None, question, max_new_tokens, temperature, top_p)

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🧠 SpatialAxiom-9B\n"
            "An open spatial intelligence model for 3D relational inference, perspective "
            "taking, multi-view correspondence, and embodied video understanding. "
            "Upload an image or video and ask a spatial reasoning question.\n\n"
            "📊 [Model card](https://huggingface.co/SpatialAxiom/SpatialAxiom-9B) · "
            "🌐 [Project page](https://d2i-ai.github.io/SpatialAxiom/) · "
            "💻 [GitHub](https://github.com/D2I-ai/SpatialAxiom)"
        )

        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Image", type="filepath", sources=["upload", "webcam"]
                )
                video_input = gr.Video(label="Video (optional)")
                question = gr.Textbox(
                    label="Question",
                    placeholder="e.g. If I stand at the door facing the bed, is the chair to my left or right?",
                    lines=2,
                )
                run_btn = gr.Button("Ask", variant="primary")
            with gr.Column(scale=1):
                output = gr.Textbox(
                    label="Answer",
                    lines=20,
                    placeholder="The model's answer will appear here…",
                )

        with gr.Accordion("Advanced settings", open=False):
            max_new_tokens = gr.Slider(
                64, 2048, value=512, step=64, label="Max new tokens"
            )
            temperature = gr.Slider(
                0.0, 2.0, value=0.0, step=0.1, label="Temperature (0 = greedy)"
            )
            top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")

        gr.Examples(
            examples=EXAMPLES,
            inputs=[image_input, question, max_new_tokens, temperature, top_p],
            outputs=output,
            fn=_answer_example,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=answer,
        inputs=[image_input, video_input, question, max_new_tokens, temperature, top_p],
        outputs=output,
        api_name="answer",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)