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 transformers import AutoModelForMultimodalLM, AutoProcessor
from PIL import Image, ImageDraw
import json
import re

MODEL_ID = "JingyuanHuang/GUI-RD-9B"

processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForMultimodalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    trust_remote_code=True,
).to("cuda").eval()

SYSTEM_PROMPT = (
    'You may call one or more functions to assist with the user query.\n\n'
    'You are provided with function signatures within <tools></tools> XML tags:\n'
    '<tools>\n'
    '{"name":"computer_use","description":"Use a mouse to interact with a computer.",'
    '"notes":"Click with the cursor tip centered on targets;avoid edges unless asked.'
    'Do not use other tools(type,key,scroll,left_click_drag).Only left_click are allowed.",'
    '"parameters":{"type":"object","required":["action"],"properties":{'
    '"action":{"type":"string","enum":["left_click"],"description":"The action to perform."},'
    '"coordinate":{"type":"array","description":"(x,y):pixels from left/top.Required for action=left_click."}'
    '}}}\n'
    '</tools>\n\n'
    'For each function call, return a JSON object with function name and arguments within <toolcall></toolcall> XML tags:\n'
    '<toolcall>\n{"name":"<function-name>","arguments":<args-json-object>}\n</toolcall>'
)


def parse_coordinate(response_text: str):
    """Extract (x, y) pixel coordinates from the model's tool-call response."""
    # Try parsing as JSON tool call first
    # Look for JSON with "coordinate" key
    match = re.search(r'\{[^{}]*"coordinate"\s*:\s*\[(\d+)\s*,\s*(\d+)\][^{}]*\}', response_text)
    if match:
        x = int(match.group(1))
        y = int(match.group(2))
        return x, y
    # Fallback: look for any [x, y] pattern
    match = re.search(r'\[(\d+)\s*,\s*(\d+)\]', response_text)
    if match:
        x = int(match.group(1))
        y = int(match.group(2))
        return x, y
    return None


def draw_marker(image: Image.Image, x: int, y: int) -> Image.Image:
    """Draw a crosshair marker and circle at the predicted coordinate."""
    annotated = image.copy()
    draw = ImageDraw.Draw(annotated)
    radius = max(10, min(image.width, image.height) // 50)
    # Draw a circle
    draw.ellipse(
        [x - radius, y - radius, x + radius, y + radius],
        outline="red",
        width=3,
    )
    # Draw crosshair lines
    line_len = radius + 8
    draw.line([(x - line_len, y), (x + line_len, y)], fill="red", width=2)
    draw.line([(x, y - line_len), (x, y + line_len)], fill="red", width=2)
    return annotated


@spaces.GPU(duration=120)
def predict(image: Image.Image, instruction: str):
    """Predict the click coordinate for a target element on a GUI screenshot.

    Args:
        image: A GUI screenshot (PNG/JPG).
        instruction: Natural-language description of the target element to click.

    Returns:
        A tuple of (annotated image, coordinate string, raw model output).
    """
    if image is None:
        return None, "Please upload an image.", ""
    if not instruction.strip():
        return None, "Please enter an instruction.", ""

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

    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = processor(
        text=[text],
        images=[image],
        padding=True,
        return_tensors="pt",
    ).to("cuda")

    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=False,
        )

    generated_ids_trimmed = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0]

    coord = parse_coordinate(output_text)
    if coord is not None:
        x, y = coord
        # Clamp to image bounds
        x = max(0, min(x, image.width - 1))
        y = max(0, min(y, image.height - 1))
        annotated = draw_marker(image, x, y)
        coord_str = f"({x}, {y})"
        return annotated, coord_str, output_text
    else:
        return image, "Could not parse coordinates from model output.", output_text


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(
        "# 🎯 GUI-RD-9B: GUI Grounding\n"
        "Upload a GUI screenshot and describe the element you want to click. "
        "The model predicts the pixel coordinate and marks it on the image.\n\n"
        "Based on [Trust the Right Teacher: Quality-Aware Self-Distillation for GUI Grounding](https://arxiv.org/abs/2606.18101)"
    )

    with gr.Row(elem_id="col-container"):
        with gr.Column(scale=1):
            input_img = gr.Image(label="GUI Screenshot", type="pil")
            instruction = gr.Textbox(
                label="Instruction",
                placeholder="e.g. click the search box",
                lines=2,
            )
            run_btn = gr.Button("Predict Coordinate", variant="primary")

        with gr.Column(scale=1):
            output_img = gr.Image(label="Predicted Click Location")
            coord_text = gr.Textbox(label="Predicted Coordinate (x, y)", interactive=False)
            with gr.Accordion("Raw Model Output", open=False):
                raw_output = gr.Textbox(label="Model Response", lines=4, interactive=False)

    gr.Examples(
        examples=[
            ["assets/web_6f93090a-81f6-489e-bb35-1a2838b18c01.png", "select search textfield"],
            ["assets/web_6f93090a-81f6-489e-bb35-1a2838b18c01.png", "switch to discussions"],
        ],
        inputs=[input_img, instruction],
        outputs=[output_img, coord_text, raw_output],
        fn=predict,
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=predict,
        inputs=[input_img, instruction],
        outputs=[output_img, coord_text, raw_output],
    )

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