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

import spaces
import torch
import json
import re
import math
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoProcessor, AutoModelForImageTextToText

MODEL_ID = "PhoneBuddyAI/PhoneBuddy-4B-RealApp"

# Build tool-call format tags as variables to avoid issues with XML-like tokens
_TC_OPEN = chr(60) + "tool_call" + chr(62)
_TC_CLOSE = chr(60) + "/tool_call" + chr(62)
_THINK_OPEN = chr(60) + "think" + chr(62)
_THINK_CLOSE = chr(60) + "/think" + chr(62)


SYSTEM_PROMPT = (
    "You are a GUI Agent. Given an instruction, the current screenshot, "
    "and the history of operations, you need to predict how to fulfill "
    "the user's request and provide the accurate invocation command. "
    "Please note that coordinate values must be scaled to a range of 0 to 1000.\n\n"
    "# Tools\n\n"
    "You can call one or more of the following functions to complete "
    "the user's request.\n\n"
    "Below is the complete list of tools supported by the system:\n"
    "<tools>\n"
    '{"type": "function", "function": {"name": "click", "description": "Click on a specified coordinate position on the screen (coordinate range 0-1000)", "parameters": {"type": "object", "properties": {"points": {"description": "A list of click coordinates, formatted as [[x, y]]", "type": "array"}}, "required": ["points"]}}}' + "\n"
    '{"type": "function", "function": {"name": "double_click", "description": "Double-click on a specified coordinate position on the screen", "parameters": {"type": "object", "properties": {"points": {"description": "Click coordinates [[x, y]]", "type": "array"}, "interval": {"description": "Interval between two clicks (milliseconds)", "type": "integer"}}, "required": ["points"]}}}' + "\n"
    '{"type": "function", "function": {"name": "long_press", "description": "Long press on a specified coordinate position on the screen", "parameters": {"type": "object", "properties": {"points": {"description": "Long press coordinates [[x, y]]", "type": "array"}, "duration": {"description": "Long press duration (milliseconds)", "type": "integer"}}, "required": ["points"]}}}' + "\n"
    '{"type": "function", "function": {"name": "type", "description": "Type text in the currently focused input field", "parameters": {"type": "object", "properties": {"text": {"description": "The text content to type", "type": "string"}}, "required": ["text"]}}}' + "\n"
    '{"type": "function", "function": {"name": "scroll", "description": "Scroll from start coordinates to target coordinates (for scrolling pages)", "parameters": {"type": "object", "properties": {"points": {"description": "Start and end coordinates for scrolling [[x1, y1], [x2, y2]]", "type": "array"}, "duration": {"description": "Scroll duration (milliseconds)", "type": "integer"}}, "required": ["points"]}}}' + "\n"
    '{"type": "function", "function": {"name": "drag", "description": "Drag an element from start coordinates to target coordinates", "parameters": {"type": "object", "properties": {"points": {"description": "Start and end coordinates for dragging [[x1, y1], [x2, y2]]", "type": "array"}, "duration": {"description": "Drag duration (milliseconds)", "type": "integer"}}, "required": ["points"]}}}' + "\n"
    '{"type": "function", "function": {"name": "button_press", "description": "Press a phone physical/virtual button", "parameters": {"type": "object", "properties": {"type": {"description": "Button type: back/home/menu/enter", "type": "string", "enum": ["back", "home", "menu", "enter"]}}, "required": ["type"]}}}' + "\n"
    '{"type": "function", "function": {"name": "open_app", "description": "Open an app by package name", "parameters": {"type": "object", "properties": {"package": {"description": "App package name", "type": "string"}}, "required": ["package"]}}}' + "\n"
    '{"type": "function", "function": {"name": "close_app", "description": "Close an app by package name", "parameters": {"type": "object", "properties": {"package": {"description": "App package name", "type": "string"}}, "required": ["package"]}}}' + "\n"
    '{"type": "function", "function": {"name": "wait", "description": "Wait for a specified duration", "parameters": {"type": "object", "properties": {"time": {"description": "Wait duration (milliseconds)", "type": "integer"}}, "required": ["time"]}}}' + "\n"
    '{"type": "function", "function": {"name": "output", "description": "Output information to the user", "parameters": {"type": "object", "properties": {"text": {"description": "The text content to output", "type": "string"}}, "required": ["text"]}}}' + "\n"
    '{"type": "function", "function": {"name": "finish", "description": "Mark the task as complete and output the final result", "parameters": {"type": "object", "properties": {"text": {"description": "Description or result of the completed task", "type": "string"}}, "required": ["text"]}}}' + "\n"
    "</tools>\n\n"
    "When making a function call, first output your thought process in natural language, "
    "then make the function call.\n"
    "The format for each function call is as follows:\n"
    + _TC_OPEN + "\n"
    + '{"name": <function-name>, "arguments": <args-json-object>}\n'
    + _TC_CLOSE
)

# Load model and processor at module scope
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()


def _lenient_json_loads(s):
    s = s.strip()
    try:
        return json.loads(s)
    except Exception:
        pass
    start = s.find("{")
    end = s.rfind("}")
    if start != -1 and end != -1 and end > start:
        s = s[start : end + 1]
    s = s.replace("\u201c", '"').replace("\u201d", '"').replace("\u2018", "'").replace("\u2019", "'")
    try:
        return json.loads(s)
    except Exception:
        pass
    s2 = re.sub(r",\s*([}\]])", r"\1", s)
    try:
        return json.loads(s2)
    except Exception:
        pass
    if '"' not in s2:
        try:
            return json.loads(s2.replace("'", '"'))
        except Exception:
            pass
    raise ValueError(f"Could not parse JSON from: {s[:200]!r}")


def parse_model_response(response):
    """Parse the model response to extract thought and tool call."""
    if not response:
        return None
    try:
        # Extract thought using think tags
        think = ""
        think_pattern = _THINK_OPEN + "(.*?)" + _THINK_CLOSE
        think_match = re.search(think_pattern, response, re.DOTALL)
        if think_match:
            think = think_match.group(1).strip()

        # Remove thinking from response
        remaining = re.sub(think_pattern, "", response, flags=re.DOTALL).strip()

        # Extract tool call
        tc_pattern = re.escape(_TC_OPEN) + r"\s*(.*?)\s*" + re.escape(_TC_CLOSE)
        tc = re.search(tc_pattern, remaining, re.DOTALL)
        if not tc:
            tc = re.search(tc_pattern, response, re.DOTALL)
        if not tc:
            # Try to grab JSON after tool_call open tag
            tc_pattern2 = re.escape(_TC_OPEN) + r"\s*(\{.*\})"
            tc = re.search(tc_pattern2, response, re.DOTALL)
        if not tc:
            return None

        thought_text = remaining.split(_TC_OPEN)[0].strip() if _TC_OPEN in remaining else ""
        full_thought = chr(10).join(x for x in (think, thought_text) if x).strip()

        obj = _lenient_json_loads(tc.group(1).strip())
        if not isinstance(obj, dict):
            return None
        name = (obj.get("name") or "").strip()
        args = obj.get("arguments", {}) or {}
        if not isinstance(args, dict):
            args = {}
        if name in ("open_app", "close_app") and "package" in args:
            args["app"] = args.pop("package")
        return {"action": name.lower(), "cot": full_thought, "args": args}
    except Exception:
        return None


def _extract_point(args, index=0):
    """Extract [x, y] from points/coordinate fields."""
    coord = None
    for key in ("points", "coordinate", "point", "coordinates"):
        if key in args and args[key] is not None:
            coord = args[key]
            break
    if coord is None:
        return None
    if isinstance(coord, list) and coord:
        if isinstance(coord[0], list):
            if index < len(coord) and len(coord[index]) >= 2:
                return [int(coord[index][0]), int(coord[index][1])]
            return None
        if len(coord) >= 2:
            return [int(coord[0]), int(coord[1])]
    return None


def _scale_point(pt, w, h):
    """Scale normalized 0-1000 coordinates to pixel coordinates."""
    x = int(pt[0] * w / 1000)
    y = int(pt[1] * h / 1000)
    x = max(0, min(x, w - 1))
    y = max(0, min(y, h - 1))
    return x, y


def visualize_action(image, action_name, args):
    """Draw the predicted action on the screenshot."""
    img = image.copy()
    draw = ImageDraw.Draw(img)
    w, h = img.size

    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", max(20, w // 40))
    except Exception:
        font = ImageFont.load_default()

    if action_name in ("click", "double_click", "long_press"):
        pt = _extract_point(args)
        if pt:
            x, y = _scale_point(pt, w, h)
            r = max(15, w // 50)
            draw.ellipse([x-r, y-r, x+r, y+r], outline=(255, 0, 0), width=max(5, w // 200))
            if action_name == "double_click":
                draw.ellipse([x-r//2, y-r//2, x+r//2, y+r//2], outline=(255, 100, 0), width=max(3, w // 300))
            elif action_name == "long_press":
                draw.ellipse([x-r-5, y-r-5, x+r+5, y+r+5], outline=(0, 0, 255), width=max(3, w // 300))
            label = f"{action_name} ({x},{y})"
            draw.text((10, 10), label, fill=(255, 0, 0), font=font)
            return img

    elif action_name in ("scroll", "drag", "swipe"):
        p1 = _extract_point(args, 0)
        p2 = _extract_point(args, 1)
        if p1 and p2:
            x1, y1 = _scale_point(p1, w, h)
            x2, y2 = _scale_point(p2, w, h)
            r = max(10, w // 60)
            draw.ellipse([x1-r, y1-r, x1+r, y1+r], outline=(0, 255, 0), width=max(5, w // 200))
            draw.ellipse([x2-r, y2-r, x2+r, y2+r], outline=(255, 0, 0), width=max(5, w // 200))
            draw.line([(x1, y1), (x2, y2)], fill=(255, 200, 0), width=max(5, w // 200))
            angle = math.atan2(y2 - y1, x2 - x1)
            arrow_len = max(20, w // 30)
            for sign in [1, -1]:
                ax = x2 - arrow_len * math.cos(angle - sign * 0.4)
                ay = y2 - arrow_len * math.sin(angle - sign * 0.4)
                draw.line([(x2, y2), (ax, ay)], fill=(255, 200, 0), width=max(3, w // 250))
            label = f"{action_name} ({x1},{y1}) -> ({x2},{y2})"
            draw.text((10, 10), label, fill=(255, 0, 0), font=font)
            return img

    elif action_name == "type":
        text = args.get("text", "")
        label = f"type: {text[:50]}"
        draw.text((10, 10), label, fill=(0, 100, 255), font=font)
        return img

    elif action_name == "button_press":
        btn = args.get("type", "")
        label = f"button_press: {btn}"
        draw.text((10, 10), label, fill=(255, 100, 0), font=font)
        return img

    elif action_name in ("open_app", "close_app"):
        app = args.get("app", args.get("package", ""))
        label = f"{action_name}: {app}"
        draw.text((10, 10), label, fill=(0, 200, 100), font=font)
        return img

    elif action_name in ("finish", "output", "answer"):
        text = args.get("text", "")
        label = f"{action_name}: {text[:80]}"
        draw.text((10, 10), label, fill=(128, 0, 128), font=font)
        return img

    draw.text((10, 10), action_name, fill=(255, 0, 0), font=font)
    return img


def format_action_text(action_name, args):
    """Format the action as readable text."""
    if action_name in ("click", "double_click", "long_press"):
        pt = _extract_point(args)
        if pt:
            return f"{action_name} at normalized coordinates [{pt[0]}, {pt[1]}] (0-1000 scale)"
    elif action_name in ("scroll", "drag", "swipe"):
        p1 = _extract_point(args, 0)
        p2 = _extract_point(args, 1)
        if p1 and p2:
            return f"{action_name} from [{p1[0]}, {p1[1]}] to [{p2[0]}, {p2[1]}] (0-1000 scale)"
    elif action_name == "type":
        return f"type text: {args.get('text', '')}"
    elif action_name == "button_press":
        return f"press {args.get('type', '')} button"
    elif action_name in ("open_app", "close_app"):
        return f"{action_name}: {args.get('app', args.get('package', ''))}"
    elif action_name in ("finish", "output", "answer"):
        return f"{action_name}: {args.get('text', '')}"
    elif action_name == "wait":
        return f"wait {args.get('time', 1000)}ms"
    return f"{action_name}({json.dumps(args, ensure_ascii=False)})"


@spaces.GPU(duration=120)
def predict_action(screenshot, instruction):
    """Predict the next phone action given a screenshot and instruction.

    Args:
        screenshot: A phone screenshot image.
        instruction: The task instruction (e.g., "Open the Contacts app").

    Returns:
        A tuple of (visualized_action_image, action_text, raw_response).
    """
    if screenshot is None:
        return None, "Please upload a phone screenshot.", ""
    if not instruction.strip():
        return None, "Please provide an instruction.", ""

    if isinstance(screenshot, str):
        screenshot = Image.open(screenshot)
    img = screenshot.convert("RGB")

    # Build messages for the chat template
    messages = [
        {"role": "system", "content": ""},
        {"role": "user", "content": [
            {"type": "text", "text": SYSTEM_PROMPT},
            {"type": "image", "image": img},
            {"type": "text", "text": f"# Instruction\n{instruction}"},
        ]},
    ]

    # Apply chat template
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    # Process inputs
    inputs = processor(
        text=[text], images=[img], padding=True, return_tensors="pt"
    ).to("cuda")

    # Generate
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=2048,
            do_sample=False,
            temperature=1.0,
            top_p=1.0,
        )

    # Decode only the generated part
    input_len = inputs["input_ids"].shape[1]
    generated_ids = output_ids[0][input_len:]
    response = processor.decode(generated_ids, skip_special_tokens=False)

    # Parse the response
    parsed = parse_model_response(response)
    if parsed is None:
        return img, "Could not parse model output.", response

    action_name = parsed["action"]
    args = parsed["args"]
    cot = parsed["cot"]

    # Visualize the action on the screenshot
    vis_img = visualize_action(img, action_name, args)

    # Format the output text
    action_text = format_action_text(action_name, args)
    if cot:
        action_text = f"Thought: {cot}\n\nAction: {action_text}"

    return vis_img, action_text, response


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

with gr.Blocks() as demo:
    gr.Markdown(
        "# PhoneBuddy: Agentic Phone Use\n"
        "Upload a phone screenshot and an instruction. "
        "The model predicts the next action (click, swipe, type, etc.) "
        "and visualizes it on the screenshot.\n\n"
        "Model: [PhoneBuddy-4B-RealApp](https://huggingface.co/PhoneBuddyAI/PhoneBuddy-4B-RealApp) | "
        "Paper: [arXiv:2606.23049](https://arxiv.org/abs/2606.23049) | "
        "Code: [GitHub](https://github.com/PhoneBuddyAI/phonebuddy)"
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                screenshot_input = gr.Image(
                    label="Phone Screenshot",
                    type="pil",
                    height=500,
                )
                instruction_input = gr.Textbox(
                    label="Instruction",
                    placeholder="e.g., Open the Contacts app and add a new contact",
                    lines=2,
                )
                run_btn = gr.Button("Predict Action", variant="primary")

            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="Visualized Action",
                    type="pil",
                    height=500,
                )
                output_text = gr.Textbox(
                    label="Predicted Action",
                    lines=6,
                )

        with gr.Accordion("Raw Model Output", open=False):
            raw_output = gr.Textbox(
                label="Raw Response",
                lines=10,
                interactive=False,
            )

    run_btn.click(
        fn=predict_action,
        inputs=[screenshot_input, instruction_input],
        outputs=[output_image, output_text, raw_output],
        api_name="predict_action",
    )

    gr.Examples(
        examples=[
            ["example_home_screen.png", "Open the Phone app to make a call"],
            ["example_home_screen.png", "Search for weather on Google"],
            ["example_settings_screen.png", "Turn on Wi-Fi"],
            ["example_settings_screen.png", "Check the battery percentage"],
        ],
        inputs=[screenshot_input, instruction_input],
        outputs=[output_image, output_text, raw_output],
        fn=predict_action,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)