import spaces  # MUST come before torch / transformers (ZeroGPU)

import json
import os
import re
import time

import gradio as gr
import torch
from PIL import Image, ImageDraw
from transformers import AutoModelForImageTextToText, AutoProcessor

# GroundGUI-8B: a Qwen3-VL-8B derivative trained for GUI grounding with the
# ScaleCUA-recipe per-platform tool-call system prompts (frozen verbatim in
# prompts.json, extracted from ScaleCUAToolCallPrompt._GROUNDING_SYSTEM in
# github.com/PrentisAI/GroundGUI eval/prompts.py).
MODEL = "PrentisAI/GroundGUI-8B"

# Eval-side image budget (run_eval.sh: MAX_IMAGE_PIXELS=16777216). Larger
# screenshots are downscaled with smart_resize (factor = patch 16 * merge 2 =
# 32, the Qwen3-VL grid) before encoding; the predicted point is mapped back
# onto the ORIGINAL image afterwards, so the overlay is unaffected.
FACTOR = 32
MAX_PIXELS = 16_777_216

HERE = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(HERE, "prompts.json"), encoding="utf-8") as f:
    GROUNDING_SYSTEM = json.load(f)

TOOL_CALL_OPEN = chr(60) + "tool_call" + chr(62)
TOOL_CALL_CLOSE = chr(60) + "/" + "tool_call" + chr(62)


def smart_resize(h, w, factor=FACTOR, min_pixels=56 * 56, max_pixels=MAX_PIXELS):
    """Qwen3-VL smart_resize: snap sides to the factor grid under a pixel budget.

    Mirrors qwen_vl_utils.smart_resize as used by the GroundGUI eval
    (IMAGE_FACTOR=32, MAX_IMAGE_PIXELS=16777216).
    """
    h_bar = round(h / factor) * factor
    w_bar = round(w / factor) * factor
    if h_bar * w_bar > max_pixels:
        beta = (h * w / max_pixels) ** 0.5
        h_bar = round(h / beta / factor) * factor
        w_bar = round(w / beta / factor) * factor
    if h_bar * w_bar < min_pixels:
        beta2 = (min_pixels / (h * w)) ** 0.5
        h_bar = round(h * beta2 / factor) * factor
        w_bar = round(w * beta2 / factor) * factor
    return max(h_bar, factor), max(w_bar, factor)


def fit(image):
    """Downscale a screenshot to the eval's vision budget (factor-32 grid)."""
    w, h = image.size
    nh, nw = smart_resize(h, w)
    if (nh, nw) != (h, w):
        image = image.resize((nw, nh), Image.LANCZOS)
    return image


def _try_tool_blobs(blobs):
    """Decode tool-call JSON blobs (last first) -> (point, refusal, name, action).

    Returns (None, None, "", "") when no blob decodes into a usable payload.
    """
    for blob in reversed(blobs):
        try:
            obj = json.loads(blob)
        except json.JSONDecodeError:
            continue
        if not isinstance(obj, dict):
            continue
        name = obj.get("name") or ""
        args = obj.get("arguments") or {}
        if not isinstance(args, dict):
            continue
        if args.get("action") in ("terminate", "abstain"):
            return None, True, name, args.get("action")
        coord = args.get("coordinate")
        if isinstance(coord, (list, tuple)) and len(coord) >= 2:
            try:
                return [float(coord[0]), float(coord[1])], False, name, args.get("action")
            except (TypeError, ValueError):
                continue
    return None, None, "", ""


def extract_point(response):
    """Parse the model's tool-call reply into a [x, y] point in 0-1000 space.

    Follows the official eval parse (Qwen3VLPrompt.extract_coordinates):
    the JSON object inside the tool-call block, terminate/abstain -> refusal,
    falling back to the GUI-Owl "(x, y)" then "[x, y]" regexes (last match).
    Tag-wrapped, bare, and fenced JSON forms are all accepted, since the
    chat template may or may not re-wrap the model's reply with the tags.

    Returns (point, refusal, tool_name, action_name); point is None if
    nothing parsed, refusal True only on an explicit abstain/terminate.
    """
    tag_blobs = re.findall(
        re.escape(TOOL_CALL_OPEN) + r"\s*(\{.*?\})\s*" + re.escape(TOOL_CALL_CLOSE),
        response, re.DOTALL,
    )
    point, refusal, name, action = _try_tool_blobs(tag_blobs)
    if point is not None or refusal:
        return point, refusal, name, action
    # bare or fenced JSON objects carrying an "arguments" payload
    bare_blobs = re.findall(r'\{\s*"name"[^{}]*"arguments"\s*:\s*\{.*?\}\s*\}', response, re.DOTALL)
    point, refusal, name, action = _try_tool_blobs(bare_blobs)
    if point is not None or refusal:
        return point, refusal, name, action
    # explicit refusal wording anywhere in the reply, as in the official eval
    if ('"action": "terminate"' in response or '"status": "failure"' in response
            or '"action": "abstain"' in response):
        return None, True, "", ""
    matches = re.findall(r"\((\d+),\s*(\d+)\)", response)
    if not matches:
        matches = re.findall(r"\[(\d+),\s*(\d+)\]", response)
    if not matches:
        return None, False, "", ""
    x, y = matches[-1]
    return [int(x), int(y)], False, "", ""


def norm_to_pixels(point, width, height):
    """Map a 0-1000 normalized point onto image pixels (GUIOwlPrompt.calculate_metrics)."""
    return round(point[0] / 1000.0 * width), round(point[1] / 1000.0 * height)


def draw_marker(image, point):
    """Draw a crosshair + circle at the predicted click point.

    A white underlay keeps the marker visible on both light and dark UIs.
    """
    out = image.copy()
    d = ImageDraw.Draw(out)
    x, y = point
    r = max(10, min(out.width, out.height) // 60)
    # white underlay, then orange marker on top
    d.ellipse([x - r, y - r, x + r, y + r], outline=(255, 255, 255), width=6)
    d.ellipse([x - r, y - r, x + r, y + r], outline=(255, 64, 0), width=3)
    for (x0, y0, x1, y1) in [
        (x - r * 1.8, y, x + r * 1.8, y),  # horizontal
        (x, y - r * 1.8, x, y + r * 1.8),  # vertical
    ]:
        d.line([x0, y0, x1, y1], fill=(255, 255, 255), width=9)
        d.line([x0, y0, x1, y1], fill=(255, 64, 0), width=4)
    return out


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


@spaces.GPU(duration=30)
def predict(instruction: str, screenshot, platform: str = "desktop", max_new_tokens: int = 256):
    """Locate a UI element in a screenshot.

    Args:
        instruction: referring expression naming the element to find,
            e.g. "the redo button".
        screenshot: a GUI screenshot (desktop, web, or mobile).
        platform: one of "desktop", "web", "mobile" -- selects the
            per-platform grounding tool schema the model was trained with
            (computer_use / browser_use / mobile_use).
        max_new_tokens: generation length cap for the emitted tool call.

    Returns:
        Tuple of (annotated screenshot with the predicted click point,
        raw model output, predicted action, coordinates on the original
        image, inference seconds).
    """
    if screenshot is None:
        raise gr.Error("Please provide a GUI screenshot first.")
    if not instruction or not instruction.strip():
        raise gr.Error("Please enter an instruction, e.g. 'the redo button'.")
    platform = platform if platform in GROUNDING_SYSTEM else "desktop"

    if screenshot.mode != "RGB":
        screenshot = screenshot.convert("RGB")
    orig_w, orig_h = screenshot.size

    # Eval-order prompt: system = frozen per-platform grounding prompt,
    # user = image first, then the bare instruction (scalecua_toolcall).
    messages = [
        {"role": "system", "content": GROUNDING_SYSTEM[platform]},
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": instruction.strip()},
            ],
        },
    ]
    prompt = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = processor(
        text=[prompt], images=[fit(screenshot)], return_tensors="pt"
    ).to(model.device)

    t0 = time.perf_counter()
    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
            use_cache=True,
        )
    elapsed = time.perf_counter() - t0

    response = processor.decode(
        output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True
    ).strip()

    # Parse the full reply: the tags may or may not be re-emitted by the
    # template; extract_point handles tagged, bare, and fenced JSON forms.
    point, refusal, tool_name, action_name = extract_point(response)
    # display: strip the (non-special) tool-call tags if present, else raw
    display = response
    if TOOL_CALL_OPEN in display:
        display = display.replace(TOOL_CALL_OPEN, "").replace(TOOL_CALL_CLOSE, "").strip()
    if refusal or point is None:
        note = "Model abstained: element not found on screen." if refusal else "No coordinates parsed from the model output."
        return screenshot.copy(), display, "abstain" if refusal else "", note, f"{elapsed:.1f}s"

    px = norm_to_pixels(point, orig_w, orig_h)
    clipped = (max(0, min(orig_w - 1, px[0])), max(0, min(orig_h - 1, px[1])))
    annotated = draw_marker(screenshot, clipped)
    coords_text = f"click point ({clipped[0]}, {clipped[1]}) px on {orig_w}x{orig_h} image; model space [{point[0]:g}, {point[1]:g}] / 1000"
    action_text = (tool_name or "?") + "." + (action_name or "?") + f" -> ({clipped[0]}, {clipped[1]})"
    return annotated, display, action_text, coords_text, f"{elapsed:.1f}s"


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

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # :mag: GroundGUI · GUI Grounding

            **GroundGUI-8B** ([PrentisAI/GroundGUI-8B](https://huggingface.co/PrentisAI/GroundGUI-8B))
            is an 8B GUI-grounding vision-language model fine-tuned from Qwen3-VL-8B-Instruct with
            high-quality grounding data and distillation-to-reinforcement learning.
            Give it a GUI screenshot and a short instruction, and it emits the click
            action with coordinates in a **normalized 0–1000 space**, which are mapped onto
            your original image and marked with a crosshair.

            Paper: [GroundGUI: GUI Grounding by High-Quality Data and Distillation-to-Reinforcement Learning](https://huggingface.co/papers/2610.19100) ·
            Code: [github.com/PrentisAI/GroundGUI](https://github.com/PrentisAI/GroundGUI)

            Benchmarks: **95.2** ScreenSpot-v2 · **63.9** ScreenSpot-Pro · **86.8** MMBench-GUI L2 ·
            **38.3** UI-Vision · **65.4** OSWorld-G.
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                screenshot_in = gr.Image(
                    type="pil", label="GUI screenshot", height=420,
                    sources=["upload", "clipboard"],
                )
                instruction = gr.Textbox(
                    label="Instruction",
                    placeholder="e.g. the redo button",
                    lines=2,
                )
                platform = gr.Radio(
                    ["desktop", "web", "mobile"], value="desktop",
                    label="Platform (tool schema used in the system prompt)",
                )
                run_btn = gr.Button("Ground it", variant="primary")
                with gr.Accordion("Advanced", open=False):
                    max_new_tokens = gr.Slider(
                        64, 512, value=256, step=32, label="Max new tokens",
                    )
            with gr.Column(scale=1):
                screenshot_out = gr.Image(
                    type="pil", label="Annotated screenshot (predicted click point)",
                    height=420,
                )
                action_out = gr.Textbox(label="Predicted action", lines=1)
                coords_out = gr.Textbox(label="Coordinates", lines=2)
                raw_out = gr.Textbox(label="Raw model output", lines=6)
                time_out = gr.Textbox(label="Inference time", lines=1)

        gr.Examples(
            examples=[
                ["redo", "examples/word_redo.png", "desktop"],
                ["Split editor.", "examples/vscode_split_editor.png", "desktop"],
                ["mute", "examples/windows_mute.png", "desktop"],
                ["take a photo", "examples/linux_take_photo.png", "desktop"],
            ],
            inputs=[instruction, screenshot_in, platform],
            fn=predict,
            outputs=[screenshot_out, raw_out, action_out, coords_out, time_out],
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        predict,
        inputs=[instruction, screenshot_in, platform, max_new_tokens],
        outputs=[screenshot_out, raw_out, action_out, coords_out, time_out],
        api_name="predict",
    )

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