import spaces  # MUST come before torch / transformers
import torch
import gradio as gr
from transformers import AutoProcessor, AutoModelForImageTextToText

MODEL_ID = "webbrain-one/webbrain-vl-2-450M"

# Production six-section system prompt from WebBrain's test/vision/prompt.mjs
SYSTEM_PROMPT = (
    "You are the vision subsystem of a web-automation agent. "
    "A screenshot of the current browser viewport is attached. "
    "Describe what is on screen so the planning agent can decide its next action.\n\n"
    "Format — keep it terse, structured, no flowery prose:\n\n"
    "1) Page purpose: one line (e.g. \"GitHub repo issue list\", "
    "\"Gmail compose\", \"Stripe checkout form\").\n"
    "2) Visible text: list the EXACT strings on buttons, links, headings, tabs, "
    "and menu items. Quote them verbatim. Do not paraphrase.\n"
    "3) Inputs: list each visible form field with its label, placeholder, "
    "current value, and whether it is focused/disabled.\n"
    "4) State signals: loading spinners, toasts, modals, error banners, "
    "success messages, CAPTCHAs, cookie/consent banners, overlays.\n"
    "5) Blockers: anything that would prevent the next likely action "
    "(overlay, disabled submit, missing data, auth prompt).\n"
    "6) Unknowns: if you cannot read something clearly, say so. "
    "Do not guess numbers, names, or identifiers.\n\n"
    "Rules: no prose intro, no conclusion, no \"this screenshot shows...\", "
    "no layout description unless it matters (e.g. \"left nav is collapsed\"). "
    "If the page is blank or still loading, say that in one line and stop."
)

USER_TEXT = (
    "Describe this screenshot of the current browser viewport "
    "for a web-automation agent. Follow the format in the system prompt."
)

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


@spaces.GPU(duration=30)
def analyze_screenshot(image, max_new_tokens=800, temperature=0.0):
    """Analyze a browser screenshot and return a structured six-section observation.

    Args:
        image: A browser viewport screenshot (PIL Image).
        max_new_tokens: Maximum number of tokens to generate.
        temperature: Sampling temperature (0 = deterministic).
    """
    if image is None:
        return "Please upload a browser screenshot image."

    conversation = [
        {
            "role": "system",
            "content": [{"type": "text", "text": SYSTEM_PROMPT}],
        },
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": USER_TEXT},
            ],
        },
    ]

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

    gen_kwargs = {
        "max_new_tokens": int(max_new_tokens),
    }
    if temperature > 0:
        gen_kwargs["temperature"] = float(temperature)
    else:
        gen_kwargs["do_sample"] = False

    with torch.no_grad():
        output_ids = model.generate(**inputs, **gen_kwargs)

    # Decode only the new tokens
    generated_ids = output_ids[0, inputs["input_ids"].shape[1]:]
    result = processor.decode(generated_ids, skip_special_tokens=True)
    return result.strip()


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

with gr.Blocks() as demo:
    gr.Markdown(
        "# 🖥️ WebBrain VL 2 450M — Browser Screenshot Understanding\n"
        "Upload a browser screenshot to get a structured six-section observation "
        "for web-automation planning. Fine-tuned from "
        "[LFM2.5-VL-450M](https://huggingface.co/LiquidAI/LFM2.5-VL-450M)."
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Browser Screenshot",
                    type="pil",
                    height=400,
                )
                run_btn = gr.Button("Analyze Screenshot", variant="primary")

                with gr.Accordion("Advanced settings", open=False):
                    max_tokens = gr.Slider(
                        label="Max new tokens",
                        minimum=64,
                        maximum=2048,
                        value=800,
                        step=64,
                    )
                    temp = gr.Slider(
                        label="Temperature",
                        minimum=0.0,
                        maximum=1.0,
                        value=0.0,
                        step=0.05,
                    )

            with gr.Column(scale=1):
                output = gr.Textbox(
                    label="Structured Observation",
                    lines=24,
                    max_lines=50,
                )

        run_btn.click(
            fn=analyze_screenshot,
            inputs=[image_input, max_tokens, temp],
            outputs=output,
            api_name="analyze",
        )

        gr.Examples(
            examples=[
                ["example_signin.png", 800, 0.0],
                ["example_search.png", 800, 0.0],
                ["example_checkout.png", 800, 0.0],
            ],
            inputs=[image_input, max_tokens, temp],
            outputs=output,
            fn=analyze_screenshot,
            cache_examples=True,
            cache_mode="lazy",
        )


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