"""LBS-27B — agentic GUI grounding with zoom-in crops.

Give the model a screenshot and an instruction. It looks at the full image,
then asks to crop up to three regions (`crop_image` text actions). The demo
runs each crop, sends the zoomed image back with its mapping to the original
image, and repeats until the model answers with a click point normalized to
[0,1000] on the original screenshot.

The prompts, action parsing, crop geometry and coordinate mapping come from
the authors' evaluation client: `evaluation/common/agentic_pipeline.py` in
https://github.com/inclusionAI/UI-Venus/tree/LBS (Apache-2.0). The only
change is the transport: generation runs locally with `transformers` on
ZeroGPU, not through a vLLM OpenAI endpoint.
"""

import os

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

import spaces  # noqa: E402  must precede torch

import ast  # noqa: E402
import json  # noqa: E402
import math  # noqa: E402
import re  # noqa: E402
import secrets  # noqa: E402
import threading  # noqa: E402
import time  # noqa: E402
from io import BytesIO  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402
from PIL import Image, ImageDraw, ImageFont  # noqa: E402
from transformers import (  # noqa: E402
    AutoModelForImageTextToText,
    AutoProcessor,
    TextIteratorStreamer,
)

MODEL_ID = "inclusionAI/LBS-27B"
HERE = Path(__file__).parent
EX = HERE / "examples"

print(f"Loading {MODEL_ID} ...", flush=True)
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID, dtype=torch.bfloat16, attn_implementation="sdpa"
).to("cuda")
model.eval()
# The checkpoint ships `use_cache: false` (a training leftover); turn the cache on
# so decoding is usable.
model.config.use_cache = True
model.config.text_config.use_cache = True
model.generation_config.use_cache = True
print("Model loaded.", flush=True)

# ---------------------------------------------------------------------------
# Protocol constants and prompts, copied from evaluation/common/agentic_pipeline.py
# ---------------------------------------------------------------------------
CROP_UPSCALE_FACTOR = 4.0
CROP_MAX_PIXELS = 4_000_000
CROP_MAX_DIMENSION = 4096
MAX_CROP_IMAGE_CALLS = 3
MAX_CROP_IMAGE_CALLS_PER_TURN = 2
MAX_ACTION_CORRECTIONS = 2
API_IMAGE_QUALITY = 90
TRUNCATED_ASSISTANT_CONTENT_CHARS = 2_000
DEGENERATE_RESCUE_MAX_TOKENS = 512

COORDINATE_RE = re.compile(
    r"\[\s*(-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)\s*,\s*"
    r"(-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)\s*\]"
)
ACTION_BLOCK_RE = re.compile(r"<action>\s*(.*?)\s*</action>", re.DOTALL | re.IGNORECASE)

CROP_IMAGE_DESCRIPTION = (
    "Crop an image that has already been supplied in this conversation. image_path "
    "must be the exact original path or a path returned by an earlier crop_image call. "
    "x1, y1, x2, y2 are absolute pixel coordinates in that referenced image's LOCAL "
    "coordinate system, with (x1,y1) at the top-left and (x2,y2) at the bottom-right. "
    "The result provides a new image path, dimensions, and the crop region mapped into "
    "the ORIGINAL image. Returned crops are resized for inspection, so nested crops must "
    "use the returned width and height. Crops may be nested. The final answer must always "
    "use relative coordinates of the ORIGINAL image normalized to the [0,1000] range, "
    "never crop-local coordinates. When writing a crop_image action, format each of x1, "
    "y1, x2, and y2 with exactly two digits after the decimal point."
)

SYSTEM_PROMPT_TEMPLATE = """You are an agentic GUI grounding model.

Task instruction:
{instruction}

Original image metadata:
- image_path: {image_path}
- resolution: {width}x{height} pixels (width x height)
- origin: [0,0] at the top-left

Goal:
Output the center point of the position corresponding to the task instruction. The final
coordinate MUST be relative to the ORIGINAL image and normalized to [0.00,1000.00], where
[0.00,0.00] is the original image's top-left and [1000.00,1000.00] is its bottom-right.
Never output absolute pixels or coordinates local to a crop. Preserve sub-unit precision:
the two final coordinate values MUST each have exactly two digits after the decimal point
(for example, [512.34,76.50]). If the task is infeasible (for example, the requested UI
element or action is absent from the image), use [-1.00,-1.00].

Available local action:
- crop_image: {crop_description}

Text-action protocol:
Every turn must first emit visible reasoning inside <think>...</think>, then end with
exactly one of these choices:
- To inspect one region:
  <think>brief reasoning for the crop</think>
  <action>crop_image(image_path="exact path", x1=100.00, y1=200.00, x2=500.00, y2=600.00)</action>
- To inspect two independent regions in the same turn:
  <think>brief reasoning for both crops</think>
  <action>crop_image(image_path="exact path", x1=100.00, y1=200.00, x2=500.00, y2=600.00)</action>
  <action>crop_image(image_path="exact path", x1=600.00, y1=200.00, x2=900.00, y2=600.00)</action>
- To return the final point:
  <think>brief reasoning for the final point</think>
  [x.xx,y.yy]

crop_image is a textual action interpreted locally after your answer. It is NOT an API
function call. Do not emit JSON, <tool_call> XML, or native tool_calls. Keep x1, y1, x2,
and y2 as four separately named numeric arguments.

Tool rules:
1. Before returning the final coordinate, you SHOULD make at least one successful
   crop_image call to inspect the target more precisely. A direct first-turn coordinate is
   discouraged. Skip cropping only when the task is clearly infeasible or cropping cannot
   reasonably improve localization.
2. You may make at most {max_tool_calls} crop_image calls in total. Never exceed this
   budget; use each crop purposefully and return the final coordinate as soon as the target
   is localized. A single assistant turn may contain one or two crop_image actions, and
   every action counts separately toward the total budget. Never emit more than two
   actions in one turn or mix crop actions with a final coordinate.
3. Tool arguments are local absolute pixels in the image_path passed to that call. Write
   x1, y1, x2, and y2 with exactly two digits after the decimal point, including trailing
   zeros (for example, x1=100.00). Two actions emitted in the same turn must refer only
   to image paths already available before that turn; the second action cannot depend on
   the result of the first action from the same turn.
4. Each result states the crop's region both in original-image pixels and original-image
   normalized [0,1000] coordinates. Use that mapping to convert any point found in a crop
   back to the ORIGINAL image's normalized [0,1000] coordinate system.
5. After crop_image actions, the following single user message contains all result
   metadata and all successful crop images in action order.

Final-answer rule:
When ready, output exactly <think>...</think> followed by a newline and [x.xx,y.yy], with
no prose, Markdown, or action wrapper after the coordinate. Both values must contain
exactly two digits after the decimal point, including trailing zeros. The point must be in
the ORIGINAL image's normalized [0,1000] coordinate system. If the task is infeasible, use
[-1.00,-1.00].
"""


# ---------------------------------------------------------------------------
# Parsing (point answer mode of the reference parser)
# ---------------------------------------------------------------------------
def _detect_degenerate(text, *, max_unit=64, min_reps=10):
    if not text:
        return None
    n = len(text)
    for unit in range(1, min(max_unit, n // min_reps) + 1):
        if n % unit == 0 and text[:unit] * (n // unit) == text:
            return {"unit_chars": unit, "reps": n // unit}
    return None


def _coordinate_after_think(text):
    value = text or ""
    closes = list(re.finditer(r"</think\s*>", value, re.IGNORECASE))
    if closes:
        matches = list(COORDINATE_RE.finditer(value[closes[-1].end():]))
        if matches:
            return matches[-1]
    return None


def _final_coordinate_match(text):
    after = _coordinate_after_think(text)
    if after:
        return after
    matches = list(COORDINATE_RE.finditer(text or ""))
    return matches[-1] if matches else None


def _extract_call_source(text):
    match = re.search(r"\b(?:crop_image|crop|click|finished)\s*\(", text, re.IGNORECASE)
    if not match:
        return None
    start = match.start()
    open_index = text.find("(", start)
    depth, quote, escaped = 0, None, False
    for index in range(open_index, len(text)):
        char = text[index]
        if quote is not None:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == quote:
                quote = None
            continue
        if char in {"'", '"'}:
            quote = char
        elif char == "(":
            depth += 1
        elif char == ")":
            depth -= 1
            if depth == 0:
                return text[start: index + 1]
    return None


def _literal_value(node):
    if isinstance(node, ast.Constant) and node.value is Ellipsis:
        return "..."
    try:
        return ast.literal_eval(node)
    except (ValueError, TypeError, SyntaxError) as exc:
        raise ValueError("Action arguments must be Python literals") from exc


def _parse_function_action(source):
    try:
        expression = ast.parse(source.strip(), mode="eval").body
    except SyntaxError as exc:
        raise ValueError(f"Invalid action syntax: {source}") from exc
    if not isinstance(expression, ast.Call) or not isinstance(expression.func, ast.Name):
        raise ValueError("Expected a function-style action")
    if any(k.arg is None for k in expression.keywords):
        raise ValueError("Action does not support **kwargs")
    arguments = {str(k.arg): _literal_value(k.value) for k in expression.keywords}
    positional = [_literal_value(item) for item in expression.args]
    name = expression.func.id
    if positional:
        keys = ("image_path", "x1", "y1", "x2", "y2") if name.lower() in {"crop", "crop_image"} else ("box",)
        if len(positional) > len(keys):
            raise ValueError("Too many positional action arguments")
        for key, value in zip(keys, positional):
            arguments.setdefault(key, value)
    return name, arguments


def _repair_malformed_crop_call(source):
    if not re.match(r"\s*(?:crop_image|crop)\s*\(", source, re.IGNORECASE):
        return None
    path_match = re.search(r"image_path\s*=\s*(['\"])(.*?)\1", source, re.DOTALL | re.IGNORECASE)
    number = r"(-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)"
    pair_match = re.search(
        rf"x1\s*=\s*{number}\s*,\s*{number}\s*,\s*x2\s*=\s*{number}\s*,\s*{number}",
        source, re.IGNORECASE,
    )
    if not pair_match:
        return None
    arguments = dict(zip(("x1", "y1", "x2", "y2"), [float(v) for v in pair_match.groups()]))
    if path_match:
        arguments["image_path"] = path_match.group(2)
    return arguments


def _normalize_crop_arguments(arguments):
    normalized = dict(arguments)
    box = normalized.pop("box", None)
    if box is not None and not all(k in normalized for k in ("x1", "y1", "x2", "y2")):
        if not isinstance(box, (list, tuple)) or len(box) != 4:
            raise ValueError("crop_image box must contain x1, y1, x2, y2")
        normalized.update(dict(zip(("x1", "y1", "x2", "y2"), box)))
    first, second = normalized.get("x1"), normalized.get("x2")
    if "y1" not in normalized and isinstance(first, (list, tuple)) and len(first) == 2:
        normalized["x1"], normalized["y1"] = first
    if "y2" not in normalized and isinstance(second, (list, tuple)) and len(second) == 2:
        normalized["x2"], normalized["y2"] = second
    return normalized


def _parse_single_text_action(text):
    value = text or ""
    after = _coordinate_after_think(value)
    if after:
        return {"type": "final", "point": [float(after.group(1)), float(after.group(2))]}
    trailing = _final_coordinate_match(value)
    if trailing and not value[trailing.end():].strip():
        return {"type": "final", "point": [float(trailing.group(1)), float(trailing.group(2))]}

    blocks = ACTION_BLOCK_RE.findall(value)
    candidate = blocks[-1].strip() if blocks else value.strip()
    call_source = _extract_call_source(candidate)
    if call_source:
        try:
            name, arguments = _parse_function_action(call_source)
        except ValueError:
            repaired = _repair_malformed_crop_call(call_source)
            name = "crop_image" if repaired else ""
            arguments = repaired or {}
        lowered = name.lower()
        if lowered in {"crop", "crop_image"}:
            try:
                return {"type": "crop", "arguments": _normalize_crop_arguments(arguments),
                        "source": call_source}
            except ValueError:
                return None
        if lowered in {"click", "finished"}:
            point = arguments.get("box", arguments.get("point"))
            if point is None and "x" in arguments and "y" in arguments:
                point = (arguments["x"], arguments["y"])
            if isinstance(point, (list, tuple)) and len(point) == 2:
                try:
                    return {"type": "final", "point": [float(point[0]), float(point[1])]}
                except (TypeError, ValueError):
                    return None

    try:
        json_action = json.loads(candidate)
    except (json.JSONDecodeError, TypeError):
        json_action = None
    if isinstance(json_action, dict):
        name = str(json_action.get("name") or json_action.get("action") or "")
        arguments = json_action.get("arguments") or {}
        if name.lower() in {"crop", "crop_image"} and isinstance(arguments, dict):
            try:
                return {"type": "crop", "arguments": _normalize_crop_arguments(arguments),
                        "source": candidate}
            except ValueError:
                return None

    coordinate = _final_coordinate_match(value)
    if coordinate:
        return {"type": "final", "point": [float(coordinate.group(1)), float(coordinate.group(2))]}
    return None


def parse_text_actions(text):
    """Return one final answer or up to two crop actions, or None if invalid."""
    value = text or ""
    matches = list(ACTION_BLOCK_RE.finditer(value))
    if matches:
        if len(matches) > MAX_CROP_IMAGE_CALLS_PER_TURN:
            return None
        if _final_coordinate_match(value):
            return None
        actions = [_parse_single_text_action(m.group(1).strip()) for m in matches]
        if any(a is None for a in actions):
            return None
        if len(actions) > 1 and any(a["type"] != "crop" for a in actions):
            return None
        return actions
    action = _parse_single_text_action(value)
    return [action] if action is not None else None


# ---------------------------------------------------------------------------
# Images and crops
# ---------------------------------------------------------------------------
def jpeg_roundtrip(image, quality=API_IMAGE_QUALITY):
    """The reference client sends every image to the model as JPEG q90."""
    if image.mode in {"RGBA", "LA"} or "transparency" in image.info:
        rgba = image.convert("RGBA")
        rgb = Image.new("RGB", rgba.size, "white")
        rgb.paste(rgba, mask=rgba.getchannel("A"))
    else:
        rgb = image.convert("RGB")
    buf = BytesIO()
    rgb.save(buf, format="JPEG", quality=quality, optimize=True, subsampling=2)
    buf.seek(0)
    out = Image.open(buf)
    out.load()
    return out.convert("RGB")


def _json_two_decimals(value):
    if isinstance(value, float) and math.isfinite(value):
        return f"{value:.2f}"
    if isinstance(value, dict):
        return "{" + ", ".join(
            f"{json.dumps(k, ensure_ascii=False)}: {_json_two_decimals(v)}" for k, v in value.items()
        ) + "}"
    if isinstance(value, (list, tuple)):
        return "[" + ", ".join(_json_two_decimals(v) for v in value) + "]"
    return json.dumps(value, ensure_ascii=False)


def execute_crop(arguments, registry, crop_dir, crop_index):
    """Port of AgenticGroundingPipeline._execute_crop, with crops kept in memory."""
    try:
        source_path = str(arguments["image_path"]).strip()
        if source_path not in registry:
            raise ValueError("image_path is not the original image or a path returned by crop_image")
        source = registry[source_path]
        values = [float(arguments[k]) for k in ("x1", "y1", "x2", "y2")]
        if not all(math.isfinite(v) for v in values):
            raise ValueError("Crop coordinates must be finite")
        x1, y1, x2, y2 = values
        if x2 <= x1 or y2 <= y1:
            raise ValueError("Expected x1 < x2 and y1 < y2")
        left = max(0, min(source["width"], math.floor(x1)))
        top = max(0, min(source["height"], math.floor(y1)))
        right = max(0, min(source["width"], math.ceil(x2)))
        bottom = max(0, min(source["height"], math.ceil(y2)))
        if right - left < 2 or bottom - top < 2:
            raise ValueError("Clamped crop must be at least 2x2 pixels")

        cropped = source["image"].crop((left, top, right, bottom))
        cw, ch = cropped.size
        scale = min(
            CROP_UPSCALE_FACTOR,
            math.sqrt(CROP_MAX_PIXELS / (cw * ch)),
            CROP_MAX_DIMENSION / cw,
            CROP_MAX_DIMENSION / ch,
        )
        ow = max(1, min(CROP_MAX_DIMENSION, math.floor(cw * scale)))
        oh = max(1, min(CROP_MAX_DIMENSION, math.floor(ch * scale)))
        if ow * oh > CROP_MAX_PIXELS:
            if ow >= oh:
                ow = max(1, CROP_MAX_PIXELS // oh)
            else:
                oh = max(1, CROP_MAX_PIXELS // ow)
        if (ow, oh) != cropped.size:
            cropped = cropped.resize((ow, oh), resample=Image.Resampling.LANCZOS)

        original = source["original_region"]
        psx = (original[2] - original[0]) / source["width"]
        psy = (original[3] - original[1]) / source["height"]
        region = [
            original[0] + left * psx,
            original[1] + top * psy,
            original[0] + right * psx,
            original[1] + bottom * psy,
        ]
        osx = (region[2] - region[0]) / ow
        osy = (region[3] - region[1]) / oh
        W, H = source["original_width"], source["original_height"]
        out_path = f"{crop_dir}/crop_{crop_index:02d}.png"
        registry[out_path] = {
            "image": cropped,
            "width": ow,
            "height": oh,
            "original_region": region,
            "original_width": W,
            "original_height": H,
        }
        normalized = [region[0] * 1000 / W, region[1] * 1000 / H, region[2] * 1000 / W, region[3] * 1000 / H]
        return {
            "ok": True,
            "image_path": out_path,
            "width": ow,
            "height": oh,
            "crop_in_parent": [left, top, right, bottom],
            "crop_size_before_resize": [cw, ch],
            "resize_scale": round(scale, 2),
            "original_image_region_pixels": [round(v, 2) for v in region],
            "original_image_region_0_1000": [round(v, 2) for v in normalized],
            "mapping_to_original_0_1000": {
                "x": f"({region[0]:.2f} + crop_x * {osx:.2f}) * 1000 / {W}",
                "y": f"({region[1]:.2f} + crop_y * {osy:.2f}) * 1000 / {H}",
            },
            "reminder": "Final coordinates must be normalized [0,1000] in the original image.",
        }
    except Exception as exc:
        return {"ok": False, "error": f"{type(exc).__name__}: {exc}"}


def crop_results_message(events, registry):
    """One user message with every crop result and successful crop image."""
    content, images = [], []
    for position, event in enumerate(events, start=1):
        result = event["result"]
        if len(events) == 1:
            heading = "crop_image action result:"
        else:
            heading = (f"crop_image action result {position}/{len(events)} "
                       f"(tool_call_index={event['tool_call_index']}):")
        lines = [heading, _json_two_decimals(result)]
        path = result.get("image_path") if result.get("ok") else None
        if path:
            lines.append(f"image_path: {path}")
        content.append({"type": "text", "text": "\n".join(lines)})
        if path:
            content.append({"type": "image"})
            images.append(jpeg_roundtrip(registry[path]["image"]))
    content.append({
        "type": "text",
        "text": (
            "Inspect all successful crops and account for any failed actions. "
            "Reply with <think>...</think>, then either one or two independent "
            "crop_image <action> blocks whose coordinates have exactly two "
            "decimal places, or the final [x.xx,y.yy]. Do not mix crop actions "
            "with a final answer."
        ),
    })
    return {"role": "user", "content": content}, images


def visible_content(generated):
    """The template pre-fills `<think>\\n`; put it back like the reference client does."""
    content = generated or ""
    has_open = re.search(r"<think\b[^>]*>", content, re.IGNORECASE)
    has_close = re.search(r"</think\s*>", content, re.IGNORECASE)
    if has_close and (not has_open or has_close.start() < has_open.start()):
        return f"<think>\n{content.lstrip()}"
    return content


# ---------------------------------------------------------------------------
# Drawing
# ---------------------------------------------------------------------------
CROP_COLORS = [(33, 150, 243), (255, 152, 0), (156, 39, 176)]
POINT_COLOR = (229, 57, 53)


def _font(size):
    for path in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
                 "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):
        try:
            return ImageFont.truetype(path, size)
        except OSError:
            continue
    try:
        return ImageFont.load_default(size=size)
    except TypeError:
        return ImageFont.load_default()


def _label(draw, xy, text, color, font, bounds):
    l, t, r, b = draw.textbbox((0, 0), text, font=font)
    w, h = r - l, b - t
    x = min(max(xy[0], 4), bounds[0] - w - 10)
    y = min(max(xy[1], 4), bounds[1] - h - 10)
    draw.rectangle([x, y, x + w + 8, y + h + 6], fill=color)
    draw.text((x + 4 - l, y + 3 - t), text, fill=(255, 255, 255), font=font)


def annotate(image, crop_regions, point_px):
    canvas = image.convert("RGB").copy()
    draw = ImageDraw.Draw(canvas)
    W, H = canvas.size
    unit = max(2, round(max(W, H) / 640))
    font = _font(max(14, unit * 8))
    for i, region in enumerate(crop_regions):
        color = CROP_COLORS[i % len(CROP_COLORS)]
        draw.rectangle(region, outline=color, width=unit * 2)
        _label(draw, (region[0], region[1] - unit * 12), f"crop {i + 1}", color, font, (W, H))
    if point_px is not None:
        x, y = point_px
        r = unit * 10
        draw.ellipse([x - r, y - r, x + r, y + r], outline=POINT_COLOR, width=unit * 2)
        draw.line([x - 2 * r, y, x + 2 * r, y], fill=POINT_COLOR, width=unit)
        draw.line([x, y - 2 * r, x, y + 2 * r], fill=POINT_COLOR, width=unit)
        _label(draw, (x + r + unit * 3, y - r), f"click ({round(x)}, {round(y)})", POINT_COLOR, font, (W, H))
    return canvas


def render_transcript(turns):
    parts = []
    for turn in turns:
        if turn["role"] == "model":
            parts.append(f"**Turn {turn['turn']} · LBS-27B**\n\n```text\n{turn['text'].strip()}\n```")
        else:
            parts.append(f"**Client →** {turn['text']}")
    return "\n\n".join(parts) if parts else "_Waiting for the model…_"


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
def _stream_generate(messages, images, *, thinking, effort, max_new_tokens, temperature):
    template_kwargs = {"thinking": thinking, "enable_thinking": thinking}
    if thinking:
        template_kwargs["reasoning_effort"] = effort
    prompt = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True, **template_kwargs
    )
    inputs = processor(text=[prompt], images=images or None, return_tensors="pt").to(model.device)
    streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
    sample = float(temperature) > 0.0
    holder = {"error": None, "n_tokens": 0}

    def worker():
        try:
            with torch.inference_mode():
                out = model.generate(
                    **inputs,
                    streamer=streamer,
                    max_new_tokens=int(max_new_tokens),
                    do_sample=sample,
                    temperature=float(temperature) if sample else None,
                    top_p=0.95 if sample else None,
                    top_k=20 if sample else None,
                    use_cache=True,
                )
            holder["n_tokens"] = int(out.shape[1] - inputs["input_ids"].shape[1])
        except BaseException as exc:  # surfaced after the stream drains
            holder["error"] = exc
            streamer.end()

    thread = threading.Thread(target=worker, daemon=True)
    thread.start()
    text = ""
    for chunk in streamer:
        text += chunk
        yield text, None
    thread.join()
    if holder["error"] is not None:
        raise gr.Error(f"Generation failed: {holder['error']}")
    finish = "length" if holder["n_tokens"] >= int(max_new_tokens) else "stop"
    yield text, finish


def _duration(image=None, instruction="", reasoning_effort="xhigh", max_new_tokens=4096,
              temperature=0.0, *args, **kwargs):
    # Measured on the live Space (xlarge): 35-66 s per run for 1-3 crops, cold start
    # included. Kept tight because xlarge bills ZeroGPU quota at 2x.
    return 80


@spaces.GPU(duration=_duration, size="xlarge")
def ground(image, instruction, reasoning_effort="xhigh", max_new_tokens=4096, temperature=0.0):
    """Locate the UI element described by an instruction on a screenshot, zooming in with crops.

    Args:
        image: the screenshot (PIL image).
        instruction: what to click, e.g. "open the sort dialog".
        reasoning_effort: chat-template reasoning effort, "xhigh" (default), "medium" or "low".
        max_new_tokens: generation budget per model turn.
        temperature: 0 for greedy decoding (the authors' evaluation setting).

    Returns:
        (annotated screenshot, list of crops, markdown summary, JSON result, transcript markdown)
    """
    if image is None:
        raise gr.Error("Upload a screenshot first.")
    if not instruction or not instruction.strip():
        raise gr.Error("Describe the element to locate.")
    if isinstance(image, str):
        image = Image.open(image)
    original = image.convert("RGB")
    W, H = original.size
    started = time.perf_counter()

    conv_id = secrets.token_hex(8)
    crop_dir = f"/tmp/lbs/conversation_{conv_id}"
    orig_path = f"/tmp/lbs/conversation_{conv_id}/original.png"
    registry = {orig_path: {
        "image": original, "width": W, "height": H,
        "original_region": [0.0, 0.0, float(W), float(H)],
        "original_width": W, "original_height": H,
    }}
    system_prompt = SYSTEM_PROMPT_TEMPLATE.format(
        instruction=instruction.strip().rstrip("."),
        image_path=orig_path, width=W, height=H,
        crop_description=CROP_IMAGE_DESCRIPTION,
        max_tool_calls=MAX_CROP_IMAGE_CALLS,
    )
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": [{"type": "text", "text": f"image_path: {orig_path}"},
                                     {"type": "image"}]},
    ]
    images = [jpeg_roundtrip(original)]

    turns, crops_gallery, crop_regions = [], [], []
    tool_count = tool_attempts = corrections = 0
    current_path = orig_path
    final_point = None
    error = None
    thinking_next, max_tokens_next = True, None

    def snapshot(status, point_px=None):
        return (
            annotate(original, crop_regions, point_px),
            list(crops_gallery),
            status,
            {},
            render_transcript(turns),
        )

    yield snapshot("Looking at the full screenshot…")

    max_turns = MAX_CROP_IMAGE_CALLS + MAX_ACTION_CORRECTIONS + 1
    for turn in range(1, max_turns + 1):
        thinking = thinking_next
        budget = max_tokens_next or int(max_new_tokens)
        thinking_next, max_tokens_next = True, None
        turns.append({"role": "model", "turn": turn, "text": ""})
        last_emit = 0.0
        generated, finish = "", None
        for generated, finish in _stream_generate(
            messages, images, thinking=thinking, effort=reasoning_effort,
            max_new_tokens=budget, temperature=temperature,
        ):
            now = time.perf_counter()
            if finish is None and now - last_emit < 0.5:
                continue
            last_emit = now
            turns[-1]["text"] = (visible_content(generated) if finish else
                                 (("<think>\n" if thinking else "") + generated))
            if finish is None:
                yield snapshot(f"Turn {turn}: model is thinking…")
        content = visible_content(generated) if thinking else generated
        turns[-1]["text"] = content
        degenerate = _detect_degenerate(generated) or _detect_degenerate(content)
        actions = parse_text_actions(content)

        if actions is None:
            if corrections >= MAX_ACTION_CORRECTIONS:
                error = (f"The model did not return a parseable crop action or final coordinate "
                         f"after {corrections + 1} attempts.")
                break
            corrections += 1
            truncated = finish == "length"
            if degenerate:
                messages.append({"role": "assistant", "content": "(repeated degenerate output, discarded)"})
                thinking_next, max_tokens_next = False, DEGENERATE_RESCUE_MAX_TOKENS
                note = ("Your previous reply degenerated into a repeated token loop and was discarded. "
                        "Do not reason, plan, or request crops now. Reply immediately with exactly "
                        "[x.xx,y.yy] coordinate and nothing else.")
            elif truncated:
                short = content[:TRUNCATED_ASSISTANT_CONTENT_CHARS] + " ...[truncated]" \
                    if len(content) > TRUNCATED_ASSISTANT_CONTENT_CHARS else content
                messages.append({"role": "assistant", "content": short})
                note = ("Your previous reply hit the token limit and was cut off before producing any "
                        "crop action or final answer. Do not reason, plan, or request crops now. Reply "
                        "immediately with exactly [x.xx,y.yy] coordinate and nothing else.")
            else:
                messages.append({"role": "assistant", "content": content})
                note = ("Your previous answer was not valid. Reply with either one or two independent "
                        "crop actions, each in its own <action>crop_image(...)</action> block, or one "
                        "final [x.xx,y.yy] coordinate. Do not emit more than two actions or mix actions "
                        "with a final answer. Every crop coordinate value must have exactly two digits "
                        "after the decimal point.")
            messages.append({"role": "user", "content": note})
            turns.append({"role": "client", "text": f"_correction:_ {note}"})
            yield snapshot(f"Turn {turn}: invalid reply, asking the model to correct it…")
            continue

        if actions[0]["type"] == "final":
            final_point = actions[0]["point"]
            break

        messages.append({"role": "assistant", "content": content})
        batch_source = current_path
        events = []
        for action in actions:
            tool_attempts += 1
            arguments = dict(action.get("arguments") or {})
            arguments.setdefault("image_path", batch_source)
            if tool_count >= MAX_CROP_IMAGE_CALLS:
                result = {"ok": False, "error": (
                    f"Crop budget exhausted: maximum is {MAX_CROP_IMAGE_CALLS}. "
                    "Return the final [x.xx,y.yy] now, using exactly two decimal places.")}
            else:
                tool_count += 1
                result = execute_crop(arguments, registry, crop_dir, tool_count)
            events.append({"tool_call_index": tool_attempts, "result": result})
            if result.get("ok"):
                current_path = result["image_path"]
                region = result["original_image_region_pixels"]
                crop_regions.append(region)
                crops_gallery.append((
                    registry[current_path]["image"],
                    f"crop {len(crop_regions)} · original px "
                    f"[{region[0]:.0f}, {region[1]:.0f}, {region[2]:.0f}, {region[3]:.0f}] · "
                    f"{result['width']}×{result['height']}",
                ))
                turns.append({"role": "client", "text": (
                    f"executed crop {len(crop_regions)}: original region "
                    f"`[{region[0]:.0f}, {region[1]:.0f}, {region[2]:.0f}, {region[3]:.0f}]`, "
                    f"zoomed ×{result['resize_scale']:.2f} → {result['width']}×{result['height']}")})
            else:
                turns.append({"role": "client", "text": f"crop failed: `{result['error']}`"})
        msg, new_images = crop_results_message(events, registry)
        messages.append(msg)
        images.extend(new_images)
        yield snapshot(f"Turn {turn}: executed {len(events)} crop(s), sending the zoomed view back…")
    else:
        error = "The conversation used up the model-turn limit without a final coordinate."

    elapsed = time.perf_counter() - started
    print(f"[timing] {len(turns)} steps, {tool_count} crops, {elapsed:.1f}s", flush=True)

    result = {
        "instruction": instruction.strip(),
        "image_size": [W, H],
        "point_0_1000": final_point,
        "point_pixels": None,
        "infeasible": False,
        "crops": [[round(v, 2) for v in r] for r in crop_regions],
        "model_turns": sum(1 for t in turns if t["role"] == "model"),
        "seconds": round(elapsed, 1),
    }
    point_px = None
    if final_point is not None and final_point[0] < 0 and final_point[1] < 0:
        result["infeasible"] = True
        status = "### The model judged the target **absent / infeasible** (`[-1.00, -1.00]`)."
    elif final_point is not None:
        point_px = (final_point[0] / 1000 * W, final_point[1] / 1000 * H)
        result["point_pixels"] = [round(point_px[0], 1), round(point_px[1], 1)]
        status = (f"### Click at **({point_px[0]:.0f}, {point_px[1]:.0f})** px\n"
                  f"normalized `[{final_point[0]:.2f}, {final_point[1]:.2f}]` on a {W}×{H} screenshot · "
                  f"{tool_count} crop(s) · {result['model_turns']} model turn(s) · {elapsed:.0f}s")
    else:
        status = f"### No final coordinate\n{error}"
        result["error"] = error
    yield (annotate(original, crop_regions, point_px), list(crops_gallery), status, result,
           render_transcript(turns))


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

EXAMPLES = [
    [str(EX / "quartus.png"), "Programmer"],
    [str(EX / "excel.png"), "open the sort dialog"],
    [str(EX / "calc.png"), "Click the cell in column Q, row 35 in the spreadsheet of LibreOffice Calc."],
    [str(EX / "note_style.png"),
     'Click on the style that the word "NOTE" is written in a light yellow font color in the pop-up window'],
    [str(EX / "blender.png"), "Edit the Distance Min value for adjusting the minimum distance value."],
]

with gr.Blocks(title="LBS-27B GUI grounding") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🔍 LBS-27B: GUI grounding with zoom-in crops\n"
            "[LBS-27B](https://huggingface.co/inclusionAI/LBS-27B) (*Learning to Look Beyond Screens*) is a "
            "27B vision-language model trained to **look closer before it answers**. Give it a screenshot and an "
            "instruction. It can request up to three `crop_image` zoom-ins; the demo runs each crop, sends the "
            "magnified region back, and the model returns the click point on the original screenshot.  \n"
            "[Model](https://huggingface.co/inclusionAI/LBS-27B) · "
            "[Code](https://github.com/inclusionAI/UI-Venus/tree/LBS)"
        )
        with gr.Row():
            with gr.Column(scale=1):
                image_in = gr.Image(type="pil", label="Screenshot", height=420)
                instruction = gr.Textbox(
                    label="Instruction", lines=2,
                    placeholder='e.g. "open the sort dialog" or "close the plotting window"',
                )
                run = gr.Button("Locate", variant="primary", size="lg")
                with gr.Accordion("Advanced settings", open=False):
                    effort = gr.Radio(["xhigh", "medium", "low"], value="xhigh", label="Reasoning effort",
                                      info="Chat-template setting; the authors' evaluation uses xhigh.")
                    max_tokens = gr.Slider(256, 4096, value=4096, step=128, label="Max new tokens per turn")
                    temperature = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Temperature",
                                            info="0 = greedy, as in the authors' evaluation.")
            with gr.Column(scale=1):
                annotated = gr.Image(label="Prediction (red) and crops the model asked for", height=420)
                status = gr.Markdown()
                crops = gr.Gallery(label="Zoomed crops the model inspected", columns=3, height=220)
        with gr.Accordion("Trajectory (model turns and client actions)", open=True):
            transcript = gr.Markdown()
        with gr.Accordion("JSON result", open=False):
            result_json = gr.JSON()

        outputs = [annotated, crops, status, result_json, transcript]
        run.click(ground, [image_in, instruction, effort, max_tokens, temperature], outputs,
                  api_name="ground")
        instruction.submit(ground, [image_in, instruction, effort, max_tokens, temperature], outputs,
                           api_name=False)

        gr.Examples(EXAMPLES, [image_in, instruction], outputs, fn=ground,
                    cache_examples=True, cache_mode="lazy",
                    label="Screenshots from ScreenSpot-Pro and OSWorld-G (from the authors' trajectories)")
        gr.Markdown(
            "Coordinates come back normalized to [0,1000] on the original image and are drawn in pixels here. "
            "The model predicts a location only; nothing is clicked.  \n"
            "Example screenshots: [ScreenSpot-Pro](https://huggingface.co/datasets/likaixin/ScreenSpot-Pro) (MIT) "
            "and [OSWorld-G](https://github.com/xlang-ai/OSWorld-G) (Apache-2.0)."
        )

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