"""VLM utilities for DeviceAgent multi-model integration.

Supports multiple VLM backends with model-specific prompt templates and action parsers:
- UI-TARS 1.5 (ByteDance): SOTA GUI grounding, 0-1000 normalized coords
- MobileAgent v3.5 (X-PLUG): tool_call XML format
- Generic OpenAI-compatible: fallback for GPT-4o, Claude, Qwen, etc.

Uses httpx for async VLM API calls (no openai SDK dependency).
"""

from __future__ import annotations

import base64
import json
import logging
import math
import os
import re
from datetime import datetime
from io import BytesIO
from typing import Any

import httpx
from PIL import Image, ImageDraw

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Image resize (Qwen-VL style, replaces qwen_vl_utils.smart_resize)
# ---------------------------------------------------------------------------

_IMAGE_MIN_TOKENS = 4
_IMAGE_MAX_TOKENS = 16384
_MAX_RATIO = 200


def smart_resize(
    height: int,
    width: int,
    factor: int = 16,
    min_pixels: int | None = None,
    max_pixels: int | None = None,
) -> tuple[int, int]:
    """Rescale dimensions preserving aspect ratio, divisible by *factor*."""
    _max = max_pixels if max_pixels is not None else (_IMAGE_MAX_TOKENS * factor ** 2)
    _min = min_pixels if min_pixels is not None else (_IMAGE_MIN_TOKENS * factor ** 2)

    if max(height, width) / max(min(height, width), 1) > _MAX_RATIO:
        raise ValueError(f"Aspect ratio exceeds {_MAX_RATIO}")

    def _round(n: float) -> int:
        return max(factor, round(n / factor) * factor)

    def _floor(n: float) -> int:
        return max(factor, math.floor(n / factor) * factor)

    def _ceil(n: float) -> int:
        return max(factor, math.ceil(n / factor) * factor)

    h_bar = _round(height)
    w_bar = _round(width)

    if h_bar * w_bar > _max:
        beta = math.sqrt((height * width) / _max)
        h_bar = _floor(height / beta)
        w_bar = _floor(width / beta)
    elif h_bar * w_bar < _min:
        beta = math.sqrt(_min / (height * width))
        h_bar = _ceil(height * beta)
        w_bar = _ceil(width * beta)

    return h_bar, w_bar


# ---------------------------------------------------------------------------
# Image encoding
# ---------------------------------------------------------------------------


def pil_to_base64(image: Image.Image) -> str:
    buf = BytesIO()
    image.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def image_to_base64_url(image_path: str) -> str:
    """Load image, resize for VLM input, return data-URI string."""
    path = image_path.removeprefix("file://")
    img = Image.open(path)
    h, w = smart_resize(
        img.height, img.width, factor=28, min_pixels=3136, max_pixels=10035200,
    )
    img = img.resize((w, h))
    return f"data:image/png;base64,{pil_to_base64(img)}"


# ---------------------------------------------------------------------------
# System prompt (from MobileAgent v3.5)
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = '''# Tools

You may call one or more functions to assist with the user query.

You are provided with function signatures within <tools></tools> XML tags:
<tools>
{"type": "function", "function": {"name_for_human": "mobile_use", "name": "mobile_use", "description": "Use a touchscreen to interact with a mobile device, and take screenshots.
* This is an interface to a mobile device with touchscreen. You can perform actions like clicking, typing, swiping, etc.
* Some applications may take time to start or process actions, so you may need to wait and take successive screenshots to see the results of your actions.
* The screen's resolution is 1000x1000.
* Make sure to click any buttons, links, icons, etc with the cursor tip in the center of the element. Don't click boxes on their edges unless asked.", "parameters": {"properties": {"action": {"description": "The action to perform. The available actions are:
* `key`: Perform a key event on the mobile device.
    - This supports adb's `keyevent` syntax.
    - Examples: \\"volume_up\\", \\"volume_down\\", \\"power\\", \\"camera\\", \\"clear\\".
* `click`: Click the point on the screen with coordinate (x, y).
* `long_press`: Press the point on the screen with coordinate (x, y) for specified seconds.
* `swipe`: Swipe from the starting point with coordinate (x, y) to the end point with coordinates2 (x2, y2).
* `type`: Input the specified text into the activated input box.
* `system_button`: Press the system button.
* `open`: Open an app on the device.
* `wait`: Wait specified seconds for the change to happen.
* `answer`: Terminate the current task and output the answer.
* `interact`: Resolve the blocking window by interacting with the user.
* `terminate`: Terminate the current task and report its completion status.", "enum": ["key", "click", "long_press", "swipe", "type", "system_button", "open", "wait", "answer", "interact", "terminate"], "type": "string"}, "coordinate": {"description": "(x, y): The x (pixels from the left edge) and y (pixels from the top edge) coordinates to move the mouse to. Required only by `action=click`, `action=long_press`, and `action=swipe`.", "type": "array"}, "coordinate2": {"description": "(x, y): The x (pixels from the left edge) and y (pixels from the top edge) coordinates to move the mouse to. Required only by `action=swipe`.", "type": "array"}, "text": {"description": "Required only by `action=key`, `action=type`, `action=open`, `action=answer`,and `action=interact`.", "type": "string"}, "time": {"description": "The seconds to wait. Required only by `action=long_press` and `action=wait`.", "type": "number"}, "button": {"description": "Back means returning to the previous interface, Home means returning to the desktop, Menu means opening the application background menu, and Enter means pressing the enter. Required only by `action=system_button`", "enum": ["Back", "Home", "Menu", "Enter"], "type": "string"}, "status": {"description": "The status of the task. Required only by `action=terminate`.", "type": "string", "enum": ["success", "failure"]}}, "required": ["action"], "type": "object"}, "args_format": "Format the arguments as a JSON object."}}
</tools>

For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
<tool_call>
{"name": <function-name>, "arguments": <args-json-object>}
</tool_call>

# Response format

Response format for every step:
1) Action: a short imperative describing what to do in the UI.
2) A single <tool_call>...</tool_call> block containing only the JSON: {"name": <function-name>, "arguments": <args-json-object>}.

Rules:
- Output exactly in the order: Action, <tool_call>.
- Be brief: one for Action.
- Do not output anything else outside those two parts.
- If finishing, use action=terminate in the tool call.'''


# ---------------------------------------------------------------------------
# UI-TARS 1.5 system prompt (ByteDance, SOTA on AndroidWorld/ScreenSpot)
# Coordinates: 0-1000 normalized. Output: Thought/Action format.
# ---------------------------------------------------------------------------

UI_TARS_SYSTEM_PROMPT = """You are a GUI agent. You are given a task and your action history, with screenshots. You need to perform the next action to complete the task.

## Output Format
```
Thought: ...
Action: ...
```

## Action Space

click(point='<point>x1 y1</point>')
long_press(point='<point>x1 y1</point>', time='')
type(content='') #If you want to submit your input, use "\\n" at the end of `content`.
scroll(point='<point>x1 y1</point>', direction='down or up or right or left')
open_app(app_name='')
drag(start_point='<point>x1 y1</point>', end_point='<point>x2 y2</point>')
press_home()
press_back()
wait()
finished(content='xxx') # Use escape characters \\', \\", and \\n in content part to ensure we can parse the content in normal python string format.

## Note
- Use Chinese in `Thought` part.
- Write a small plan and finally summarize your next action (with its target element) in one sentence in `Thought` part.
"""


# ---------------------------------------------------------------------------
# Model profile detection
# ---------------------------------------------------------------------------


def detect_model_profile(model_name: str) -> str:
    """Return 'uitars' or 'mobile_agent' based on model identifier."""
    lower = model_name.lower()
    if "ui-tars" in lower or "uitars" in lower:
        return "uitars"
    return "mobile_agent"


def get_system_prompt(profile: str, instruction: str = "") -> str:
    if profile == "uitars":
        prompt = UI_TARS_SYSTEM_PROMPT
        if instruction:
            prompt += f"\n## User Instruction\n{instruction}"
        return prompt
    return SYSTEM_PROMPT


# ---------------------------------------------------------------------------
# UI-TARS action parser
# ---------------------------------------------------------------------------


def parse_uitars_action(output_text: str) -> dict[str, Any]:
    """Parse UI-TARS Thought/Action output into a normalized action dict.

    Examples:
        click(point='<point>500 300</point>')
        type(content='hello world')
        scroll(point='<point>500 500</point>', direction='up')
        open_app(app_name='WeChat')
        drag(start_point='<point>100 200</point>', end_point='<point>300 400</point>')
        press_home()
        press_back()
        finished(content='task done')
    """
    action_line = ""
    for line in output_text.strip().splitlines():
        stripped = line.strip()
        if stripped.startswith("Action:"):
            action_line = stripped[7:].strip()
            break
    if not action_line:
        lines = output_text.strip().splitlines()
        for line in reversed(lines):
            stripped = line.strip()
            if "(" in stripped and not stripped.startswith("Thought"):
                action_line = stripped
                break

    if not action_line:
        raise ValueError(f"No Action line found in UI-TARS output")

    func_match = re.match(r"(\w+)\((.*)\)\s*$", action_line, re.DOTALL)
    if not func_match:
        raise ValueError(f"Cannot parse UI-TARS action: {action_line[:100]}")

    func_name = func_match.group(1)
    args_str = func_match.group(2).strip()

    def _extract_point(s: str) -> list[int]:
        m = re.search(r"<point>\s*(\d+)\s+(\d+)\s*</point>", s)
        if m:
            return [int(m.group(1)), int(m.group(2))]
        m = re.search(r"(\d+)\s*,\s*(\d+)", s)
        if m:
            return [int(m.group(1)), int(m.group(2))]
        return [0, 0]

    def _extract_kwarg(s: str, key: str) -> str:
        m = re.search(rf"{key}\s*=\s*['\"](.+?)['\"]", s)
        return m.group(1) if m else ""

    result: dict[str, Any] = {}

    if func_name == "click":
        result = {"action": "click", "coordinate": _extract_point(args_str)}
    elif func_name == "long_press":
        time_val = _extract_kwarg(args_str, "time")
        result = {
            "action": "long_press",
            "coordinate": _extract_point(args_str),
            "time": float(time_val) if time_val else 0.8,
        }
    elif func_name == "type":
        content = _extract_kwarg(args_str, "content")
        content = content.replace("\\n", "\n").replace("\\'", "'").replace('\\"', '"')
        result = {"action": "type", "text": content}
    elif func_name == "scroll":
        direction = _extract_kwarg(args_str, "direction") or "down"
        point = _extract_point(args_str) if "<point>" in args_str else [500, 500]
        _SCROLL_DELTA = 300
        dx, dy = {"up": (0, -_SCROLL_DELTA), "down": (0, _SCROLL_DELTA),
                   "left": (-_SCROLL_DELTA, 0), "right": (_SCROLL_DELTA, 0)}.get(direction, (0, _SCROLL_DELTA))
        result = {
            "action": "swipe",
            "coordinate": point,
            "coordinate2": [point[0] + dx, point[1] + dy],
        }
    elif func_name == "drag":
        start = [0, 0]
        end = [0, 0]
        m_start = re.search(r"start_point\s*=\s*['\"](.+?)['\"]", args_str)
        m_end = re.search(r"end_point\s*=\s*['\"](.+?)['\"]", args_str)
        if m_start:
            start = _extract_point(m_start.group(1))
        if m_end:
            end = _extract_point(m_end.group(1))
        result = {"action": "swipe", "coordinate": start, "coordinate2": end}
    elif func_name == "open_app":
        app = _extract_kwarg(args_str, "app_name")
        result = {"action": "open", "text": app}
    elif func_name == "press_home":
        result = {"action": "system_button", "button": "Home"}
    elif func_name == "press_back":
        result = {"action": "system_button", "button": "Back"}
    elif func_name == "wait":
        result = {"action": "wait", "time": 2}
    elif func_name == "finished":
        content = _extract_kwarg(args_str, "content")
        result = {"action": "terminate", "status": "success", "text": content}
    else:
        raise ValueError(f"Unknown UI-TARS action: {func_name}")

    return result


# ---------------------------------------------------------------------------
# Message construction (multi-turn with image history)
# ---------------------------------------------------------------------------


def build_messages(
    image_path: str,
    instruction: str,
    history: list[dict[str, str]],
    model_name: str,
    history_n: int = 4,
) -> list[dict[str, Any]]:
    """Build multi-turn VLM messages in DashScope format.

    These are later converted to OpenAI format by ``convert_to_openai_messages``.
    Auto-selects prompt template based on model profile.
    """
    profile = detect_model_profile(model_name)

    weekdays = ["Monday", "Tuesday", "Wednesday", "Thursday",
                "Friday", "Saturday", "Sunday"]
    today = datetime.today()
    date_str = today.strftime("%Y-%m-%d") + " " + weekdays[today.weekday()]

    if profile == "uitars":
        return _build_uitars_messages(
            image_path, instruction, history, date_str, history_n,
        )
    return _build_mobile_agent_messages(
        image_path, instruction, history, model_name, date_str, history_n,
    )


def _build_uitars_messages(
    image_path: str,
    instruction: str,
    history: list[dict[str, str]],
    date_str: str,
    history_n: int,
) -> list[dict[str, Any]]:
    """Build messages for UI-TARS models (Thought/Action format)."""
    full_instruction = f"Today is {date_str}. {instruction}"
    sys_prompt = get_system_prompt("uitars", full_instruction)

    messages: list[dict[str, Any]] = [
        {"role": "system", "content": [{"text": sys_prompt}]},
    ]

    history_window = history[-history_n:]
    for item in history_window:
        messages.append({
            "role": "user",
            "content": [{"image": "file://" + item["image"]}],
        })
        messages.append({
            "role": "assistant",
            "content": [{"text": item["output"]}],
        })

    messages.append({
        "role": "user",
        "content": [{"image": "file://" + image_path}],
    })
    return messages


def _build_mobile_agent_messages(
    image_path: str,
    instruction: str,
    history: list[dict[str, str]],
    model_name: str,
    date_str: str,
    history_n: int,
) -> list[dict[str, Any]]:
    """Build messages for MobileAgent / generic VLM (tool_call format)."""
    current_step = len(history)
    history_start = max(0, current_step - history_n)

    previous_actions: list[str] = []
    for i in range(history_start):
        if i < len(history):
            text = history[i]["output"]
            if model_name.endswith(".mem"):
                if "<tool_call>" in text:
                    text = text.split("<tool_call>")[0].strip()
            elif "Action:" in text and "<tool_call>" in text:
                text = text.split("Action:")[1].split("<tool_call>")[0].strip()
            previous_actions.append(f"Step {i + 1}: {text}")

    prev_str = "\n".join(previous_actions) if previous_actions else "None"

    instruction_prompt = (
        f"Please generate the next move according to the UI screenshot, "
        f"instruction and previous actions.\n\n"
        f"Instruction: Today's date is: {date_str}. {instruction}\n\n"
        f"Previous actions:\n{prev_str}"
    )

    messages: list[dict[str, Any]] = [
        {"role": "system", "content": [{"text": SYSTEM_PROMPT}]},
    ]

    history_window = history[-history_n:]
    if history_window:
        for idx, item in enumerate(history_window):
            if idx == 0:
                messages.append({
                    "role": "user",
                    "content": [
                        {"text": instruction_prompt},
                        {"image": "file://" + item["image"]},
                    ],
                })
            else:
                messages.append({
                    "role": "user",
                    "content": [{"image": "file://" + item["image"]}],
                })
            messages.append({
                "role": "assistant",
                "content": [{"text": item["output"]}],
            })
        messages.append({
            "role": "user",
            "content": [{"image": "file://" + image_path}],
        })
    else:
        messages.append({
            "role": "user",
            "content": [
                {"text": instruction_prompt},
                {"image": "file://" + image_path},
            ],
        })

    return messages


def convert_to_openai_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Convert DashScope-style messages to OpenAI chat completion format."""
    converted = []
    for msg in messages:
        role = msg["role"]
        content_items = msg.get("content", [])
        new_content: list[dict[str, Any]] = []

        for item in content_items:
            if "text" in item:
                new_content.append({"type": "text", "text": item["text"]})
            elif "image" in item:
                new_content.append({
                    "type": "image_url",
                    "image_url": {"url": image_to_base64_url(item["image"])},
                })

        if role == "system" and len(new_content) == 1 and new_content[0]["type"] == "text":
            converted.append({"role": "system", "content": new_content[0]["text"]})
        elif role == "assistant" and len(new_content) == 1 and new_content[0]["type"] == "text":
            converted.append({"role": "assistant", "content": new_content[0]["text"]})
        else:
            converted.append({"role": role, "content": new_content})

    return converted


# ---------------------------------------------------------------------------
# Async VLM API call (httpx, OpenAI-compatible)
# ---------------------------------------------------------------------------

_MAX_VLM_RETRIES = 3


async def call_vlm(
    api_key: str,
    base_url: str,
    model: str,
    openai_messages: list[dict[str, Any]],
    temperature: float = 0.0,
    max_tokens: int = 4096,
) -> str:
    """Call an OpenAI-compatible VLM and return the text response."""
    url = f"{base_url.rstrip('/')}/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    body = {
        "model": model,
        "messages": openai_messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }

    last_err: Exception | None = None
    for attempt in range(_MAX_VLM_RETRIES):
        try:
            async with httpx.AsyncClient(timeout=90) as client:
                resp = await client.post(url, json=body, headers=headers)
                resp.raise_for_status()
                data = resp.json()
            return data["choices"][0]["message"]["content"]
        except Exception as exc:
            last_err = exc
            logger.warning("VLM call attempt %d failed: %s", attempt + 1, exc)
            if attempt < _MAX_VLM_RETRIES - 1:
                import asyncio
                await asyncio.sleep(3 * (attempt + 1))

    raise RuntimeError(f"VLM call failed after {_MAX_VLM_RETRIES} retries: {last_err}")


# ---------------------------------------------------------------------------
# App name resolution via LLM
# ---------------------------------------------------------------------------


async def resolve_app_name(
    instruction: str,
    installed_app_names: list[str],
    api_key: str,
    base_url: str,
    model: str = "",
) -> str:
    """Use LLM to resolve which installed app should be opened."""
    if not installed_app_names:
        return ""

    app_list_str = ", ".join(installed_app_names)
    prompt = f'''You are an app resolver. Given a user instruction and installed app names, determine which app to open.

Input:
User instruction: "{instruction}"
Installed apps: "{app_list_str}"

Rules:
- Only select from the given app name list.
- If the instruction names a specific app and it's in the list, return it.
- If the instruction is generic, pick any matching candidate.
- If no match, return empty string.

Output only JSON: {{"reason": "...", "app": "..."}}'''

    url = f"{base_url.rstrip('/')}/chat/completions"
    try:
        async with httpx.AsyncClient(timeout=30) as client:
            resp = await client.post(
                url,
                json={
                    "model": model,
                    "messages": [
                        {"role": "system", "content": "You are a helpful assistant."},
                        {"role": "user", "content": prompt},
                    ],
                    "temperature": 0.0,
                },
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
            )
            resp.raise_for_status()
            text = resp.json()["choices"][0]["message"]["content"]
    except Exception as exc:
        logger.warning("App resolver LLM call failed: %s", exc)
        return ""

    try:
        cleaned = text
        if "```json" in cleaned:
            cleaned = cleaned.split("```json")[1].split("```")[0]
        parsed = json.loads(cleaned)
        return parsed.get("app", "")
    except (json.JSONDecodeError, IndexError):
        logger.warning("App resolver JSON parse failed: %s", text[:200])
        return ""


# ---------------------------------------------------------------------------
# Screenshot annotation (debug visualization)
# ---------------------------------------------------------------------------


def annotate_screenshot(
    image_path: str,
    action_params: dict[str, Any],
    save_path: str,
) -> str | None:
    """Draw action annotations on a screenshot for debugging."""
    image = Image.open(image_path)
    draw = ImageDraw.Draw(image)
    action_type = action_params.get("action", "")

    if action_type == "click":
        r = 15
        cx, cy = action_params.get("coordinate", [0, 0])
        draw.ellipse((cx - r, cy - r, cx + r, cy + r), fill="red", outline="red")
    elif action_type in ("scroll", "swipe"):
        x1, y1 = action_params.get("coordinate", [0, 0])
        x2, y2 = action_params.get("coordinate2", [x1, y1])
        draw.line((x1, y1, x2, y2), fill="red", width=3)
        angle = math.atan2(y2 - y1, x2 - x1)
        s = 12
        ax1 = x2 - s * math.cos(angle - math.pi / 6)
        ay1 = y2 - s * math.sin(angle - math.pi / 6)
        ax2 = x2 - s * math.cos(angle + math.pi / 6)
        ay2 = y2 - s * math.sin(angle + math.pi / 6)
        draw.polygon([(x2, y2), (ax1, ay1), (ax2, ay2)], fill="red")
    elif action_type == "long_press":
        r = 20
        cx, cy = action_params.get("coordinate", [0, 0])
        draw.ellipse((cx - r, cy - r, cx + r, cy + r), fill=None, outline="orange", width=3)
    else:
        return None

    image.save(save_path)
    return save_path
