"""M3 response parser — extracts ``<tool_call>{"name":"computer", ...}</tool_call>``
blocks from the model output and translates the 17-action computer-use
tool enum into pyautogui code that OSWorld's ``desktop_env`` executor
accepts.

Public API:
  * ``parse_m3_response(response, original_w, original_h, processed_w,
    processed_h, coordinate_type="relative")`` →
    ``(low_level_instruction: str, pyautogui_code: list[str])``
  * ``tool_action_to_pyautogui(action, args, scaled_xy)`` — single-action
    translator (also re-exported for callers that already have a tool_call).

Coordinate space: M3 emits ``[0, 1000]`` normalized integers (default
``coordinate_type="relative"``). Set ``coordinate_type="absolute"`` if you
ever feed the model a 1280x720 image and want the parser to scale based on
processed_w / processed_h instead.

Special tokens recognised by ``desktop_env.env.step``: ``WAIT``, ``DONE``,
``FAIL``, ``CALL_USER``. They are kept as standalone strings in the
returned ``pyautogui_code`` list (not ``"\\n".join``-ed) so the env can
detect them.
"""

from __future__ import annotations

import json
import logging
import re
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger("desktopenv.agent.m3.parser")


_VALID_ACTIONS = frozenset({
    "click", "left click", "right click",
    "left_click", "right_click", "double_click", "middle_click",
    "left_press", "triple_click",
    "left_mouse_down", "left_mouse_up",
    "mouse_move", "left_click_drag",
    "hold_key", "key", "type",
    "scroll",
    "wait", "screenshot", "fail", "done", "call_user",
    "terminate", "cursor_position",
    "zoom", "zoom_in",
})


_NORMALIZE_KEY: Dict[str, str] = {
    "ctrl": "ctrl", "control": "ctrl",
    "shift": "shift",
    "alt": "alt", "option": "altleft",
    "cmd": "command", "win": "win", "super": "win", "meta": "win",
    "return": "enter", "kp_enter": "enter",
    "esc": "escape", "escape": "escape",
    "tab": "tab", "iso_left_tab": "tab",
    "space": "space", " ": "space",
    "backspace": "backspace",
    "up": "up", "down": "down", "left": "left", "right": "right",
    "home": "home", "end": "end",
    "pageup": "pageup", "pgup": "pageup", "prior": "pageup",
    "pagedown": "pagedown", "pgdn": "pagedown", "next": "pagedown",
    "delete": "delete", "del": "delete",
    "insert": "insert", "ins": "insert",
    "kp_0": "num0", "kp_1": "num1", "kp_2": "num2", "kp_3": "num3", "kp_4": "num4",
    "kp_5": "num5", "kp_6": "num6", "kp_7": "num7", "kp_8": "num8", "kp_9": "num9",
    "kp_add": "add", "kp_subtract": "subtract",
    "kp_multiply": "multiply", "kp_divide": "divide",
    "menu": "apps",
}


def tool_action_to_pyautogui(action: str, args: Dict, scaled_xy) -> List[str]:
    """Translate one computer-use tool action to pyautogui code.

    ``scaled_xy(x, y)`` is a callback that maps the model's emitted
    (x, y) to actual screen pixels.

    Returns a list of code strings; special tokens (``"WAIT"``, ``"DONE"``,
    ``"FAIL"``, ``"CALL_USER"``) are returned as their own one-element list
    so the caller can keep them as standalone strings.
    """
    code: List[str] = []

    action_conversion = {"left click": "click", "right click": "right_click"}
    action = action_conversion.get(action, action)
    if action == "click":
        action = "left_click"

    text = args.get("text")
    coordinate = args.get("coordinate")
    start_coordinate = args.get("start_coordinate")
    scroll_direction = args.get("scroll_direction")
    scroll_amount = args.get("scroll_amount")
    duration = args.get("duration")

    if coordinate is not None:
        coordinate = list(scaled_xy(coordinate[0], coordinate[1]))
    if start_coordinate is not None:
        start_coordinate = list(scaled_xy(start_coordinate[0], start_coordinate[1]))

    if action in ("left_mouse_down", "left_mouse_up"):
        if text:
            for k in text.split('+'):
                code.append(f"pyautogui.keyDown('{k.strip().lower()}')")
        if coordinate is not None:
            x, y = coordinate
            code.append(
                f"pyautogui.mouseDown({x}, {y})" if action == "left_mouse_down"
                else f"pyautogui.mouseUp({x}, {y})"
            )
        else:
            code.append(
                "pyautogui.mouseDown()" if action == "left_mouse_down"
                else "pyautogui.mouseUp()"
            )
        if text:
            for k in reversed(text.split('+')):
                code.append(f"pyautogui.keyUp('{k.strip().lower()}')")
    elif action == "hold_key":
        if not isinstance(text, str):
            raise ValueError(f"{text} must be a string")
        for key in text.split('+'):
            code.append(f"pyautogui.keyDown('{key.strip().lower()}')")
    elif action in ("mouse_move", "left_click_drag"):
        if coordinate is None:
            raise ValueError(f"coordinate is required for {action}")
        if text is not None:
            raise ValueError(f"text is not accepted for {action}")
        if not isinstance(coordinate, (list, tuple)) or len(coordinate) != 2:
            raise ValueError(f"{coordinate} must be a tuple of length 2")
        if not all(isinstance(i, int) for i in coordinate):
            raise ValueError(f"{coordinate} must be a tuple of ints")
        x, y = coordinate[0], coordinate[1]
        if action == "mouse_move":
            code.append(f"pyautogui.moveTo({x}, {y}, duration={duration or 0.5})")
        else:
            if start_coordinate:
                if not isinstance(start_coordinate, (list, tuple)) or len(start_coordinate) != 2:
                    raise ValueError(f"{start_coordinate} must be a tuple of length 2")
                if not all(isinstance(i, int) for i in start_coordinate):
                    raise ValueError(f"{start_coordinate} must be a tuple of ints")
                sx, sy = start_coordinate
                code.append(f"pyautogui.moveTo({sx}, {sy}, duration={duration or 0.5})")
            code.append(f"pyautogui.dragTo({x}, {y}, duration={duration or 0.5})")
    elif action in ("key", "type"):
        if text is None:
            raise ValueError(f"text is required for {action}")
        if coordinate is not None:
            raise ValueError(f"coordinate is not accepted for {action}")
        if not isinstance(text, str):
            raise ValueError(f"{text} must be a string")
        if action == "key":
            key_conversion = {
                "page_down": "pagedown",
                "page_up": "pageup",
                "super_l": "win",
                "super": "command",
                "escape": "esc",
            }
            keys = [key_conversion.get(k.strip().lower(), k.strip().lower()) for k in text.split('+')]
            for k in keys:
                code.append(f"pyautogui.keyDown('{k}')")
            for k in reversed(keys):
                code.append(f"pyautogui.keyUp('{k}')")
        else:
            for char in text:
                if char == '\n':
                    code.append("pyautogui.press('enter')")
                elif char == "'":
                    code.append('pyautogui.press("\'")')
                elif char == '\\':
                    code.append("pyautogui.press('\\\\')")
                elif char == '"':
                    code.append("pyautogui.press('\"')")
                else:
                    code.append(f"pyautogui.press('{char}')")
    elif action == "scroll":
        if text is not None:
            for k in text.split('+'):
                code.append(f"pyautogui.keyDown('{k.strip().lower()}')")
        if coordinate is None:
            if scroll_direction in ("up", "down"):
                code.append(f"pyautogui.scroll({scroll_amount if scroll_direction == 'up' else -scroll_amount})")
            elif scroll_direction in ("left", "right"):
                code.append(f"pyautogui.hscroll({scroll_amount if scroll_direction == 'right' else -scroll_amount})")
        else:
            x, y = coordinate[0], coordinate[1]
            if scroll_direction in ("up", "down"):
                code.append(f"pyautogui.scroll({scroll_amount if scroll_direction == 'up' else -scroll_amount}, {x}, {y})")
            elif scroll_direction in ("left", "right"):
                code.append(f"pyautogui.hscroll({scroll_amount if scroll_direction == 'right' else -scroll_amount}, {x}, {y})")
        if text is not None:
            for k in reversed(text.split('+')):
                code.append(f"pyautogui.keyUp('{k.strip().lower()}')")
    elif action in ("left_click", "right_click", "double_click", "middle_click", "left_press", "triple_click"):
        if text:
            for k in text.split('+'):
                code.append(f"pyautogui.keyDown('{k.strip().lower()}')")
        if coordinate is not None:
            x, y = coordinate
            mapping = {
                "left_click":   f"pyautogui.click({x}, {y})",
                "right_click":  f"pyautogui.rightClick({x}, {y})",
                "double_click": f"pyautogui.doubleClick({x}, {y})",
                "middle_click": f"pyautogui.middleClick({x}, {y})",
                "triple_click": f"pyautogui.tripleClick({x}, {y})",
            }
            if action == "left_press":
                code.append(f"pyautogui.mouseDown({x}, {y})")
                code.append("import time as _t; _t.sleep(1)")
                code.append(f"pyautogui.mouseUp({x}, {y})")
            else:
                code.append(mapping[action])
        else:
            mapping_no_xy = {
                "left_click":   "pyautogui.click()",
                "right_click":  "pyautogui.rightClick()",
                "double_click": "pyautogui.doubleClick()",
                "middle_click": "pyautogui.middleClick()",
                "triple_click": "pyautogui.tripleClick()",
            }
            if action == "left_press":
                code.append("pyautogui.mouseDown()")
                code.append("import time as _t; _t.sleep(1)")
                code.append("pyautogui.mouseUp()")
            else:
                code.append(mapping_no_xy[action])
        if text:
            for k in reversed(text.split('+')):
                code.append(f"pyautogui.keyUp('{k.strip().lower()}')")
    elif action == "wait":
        code.append("pyautogui.sleep(0.5)")
    elif action == "fail":
        code.append("FAIL")
    elif action == "done":
        code.append("DONE")
    elif action == "call_user":
        code.append("CALL_USER")
    elif action == "screenshot":
        code.append("pyautogui.sleep(0.1)")
    elif action == "terminate":
        status = str(args.get("status", "success")).lower()
        code.append("FAIL" if status in ("failure", "fail") else "DONE")
    elif action == "cursor_position":
        code.append("pyautogui.sleep(0.1)")
    elif action in ("zoom", "zoom_in"):
        logger.info(f"action {action!r} → no-op sleep (env doesn't expose native zoom)")
        code.append("pyautogui.sleep(0.1)")
    else:
        raise ValueError(f"Invalid action: {action}")

    return code


def parse_m3_response(
    response: str,
    original_width: int = None,
    original_height: int = None,
    processed_width: int = None,
    processed_height: int = None,
    coordinate_type: str = "relative",
) -> Tuple[str, List[str]]:
    """Parse a single M3 response into ``(low_level_instruction, pyautogui_code)``.

    Steps:
    1. ``[INFEASIBLE]`` token → return ``("[INFEASIBLE]", ["FAIL"])``.
    2. Walk lines, picking up ``Action: ...`` and ``<tool_call>...</tool_call>``
       blocks. Bare JSON object lines that look like a tool call are also
       accepted (some upstream configs strip the xml wrapper).
    3. For each tool call, run ``tool_action_to_pyautogui`` with a coord
       scaler chosen by ``coordinate_type``.
    4. Stop-sequence recovery: if the upstream ate ``</tool_call>`` (set as a
       stop string) we still process the buffered body at end-of-input.

    Returns ``(low_level_instruction, pyautogui_code)`` where each tool call
    contributes one string element (multi-line code joined by ``"\\n"``) so
    ``desktop_env.env.step`` runs the whole action atomically.
    """
    if not response or not response.strip():
        return "", []
    if "[INFEASIBLE]" in response:
        return "[INFEASIBLE]", ["FAIL"]

    def scaled_xy(x, y):
        if not (original_width and original_height):
            return int(x), int(y)
        if coordinate_type == "absolute":
            if processed_width and processed_height:
                sx = original_width / processed_width
                sy = original_height / processed_height
                return int(x * sx), int(y * sy)
            return int(x), int(y)
        return int(x * original_width / 1000), int(y * original_height / 1000)

    low_level_instruction = ""
    pyautogui_code: List[str] = []

    def emit(json_str: str) -> None:
        try:
            tool_call = json.loads(json_str)
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse tool_call JSON: {e} — body: {json_str[:200]}")
            return
        if tool_call.get("name") != "computer":
            logger.warning(f"Expected name='computer', got {tool_call.get('name')!r} — skipping")
            return
        args = tool_call.get("arguments") or tool_call.get("input") or {}
        action = args.get("action")
        if not action:
            logger.warning(f"tool_call missing 'action': {tool_call}")
            return

        # BPE tokenizer artifact recovery (M3 quirks observed in production).
        if isinstance(action, str) and "_ " in action:
            fixed = action.replace("_ ", "_")
            logger.info(f"[parser] BPE-fix action: {action!r} → {fixed!r}")
            action = fixed
            args = dict(args, action=fixed)
        if isinstance(args, dict):
            renamed = {}
            changed = False
            for k, v in args.items():
                if isinstance(k, str) and "_ " in k:
                    nk = k.replace("_ ", "_")
                    logger.info(f"[parser] BPE-fix arg key: {k!r} → {nk!r}")
                    renamed[nk] = v
                    changed = True
                else:
                    renamed[k] = v
            if changed:
                args = renamed
        if isinstance(action, str) and "_lick" in action and action not in _VALID_ACTIONS:
            candidate = action.replace("_lick", "_click")
            if candidate in _VALID_ACTIONS:
                logger.info(f"[parser] BPE click-recover: {action!r} → {candidate!r}")
                action = candidate
                args = dict(args, action=candidate)
        if (isinstance(action, str)
                and action not in _VALID_ACTIONS
                and (action.endswith("_") or len(action) < 3)):
            logger.warning(f"[parser] truncated/short action {action!r} → no-op")
            pyautogui_code.append("pyautogui.sleep(0.1)")
            return

        try:
            code = tool_action_to_pyautogui(action, args, scaled_xy)
        except ValueError as e:
            logger.warning(f"[parser] invalid action {action!r}: {e} → no-op")
            pyautogui_code.append("pyautogui.sleep(0.1)")
            return
        if not code:
            return
        if len(code) == 1 and code[0] in ("FAIL", "DONE", "CALL_USER", "WAIT"):
            pyautogui_code.append(code[0])
        else:
            pyautogui_code.append("\n".join(code))

    inside = False
    buffer: List[str] = []
    for line in response.split("\n"):
        stripped = line.strip()
        if not stripped:
            continue
        lower = stripped.lower()
        if lower.startswith("action:"):
            if not low_level_instruction:
                low_level_instruction = stripped.split(":", 1)[-1].strip()
            continue
        if stripped.startswith("<tool_call>"):
            inside = True
            continue
        if stripped.startswith("</tool_call>"):
            if buffer:
                emit("\n".join(buffer))
                buffer = []
            inside = False
            continue
        if inside:
            buffer.append(stripped)
            continue
        if stripped.startswith("{") and stripped.endswith("}"):
            try:
                obj = json.loads(stripped)
                if isinstance(obj, dict) and "name" in obj and "arguments" in obj:
                    emit(stripped)
            except json.JSONDecodeError:
                pass

    if inside and buffer:
        emit("\n".join(buffer))

    return low_level_instruction, pyautogui_code


def wrap_for_history(raw_text: str) -> str:
    """Re-wrap bare JSON tool calls with ``<tool_call>...</tool_call>`` so the
    chat template renders past assistant turns in the same format the model
    was trained to expect.

    Also recovers from upstream stop-sequence configs that consume
    ``</tool_call>`` before it reaches ``message.content`` — appends the
    missing closing tag so history stays well-formed.
    """
    out = raw_text or ""
    if "<tool_call>" in out:
        if "</tool_call>" not in out:
            out = out.rstrip() + "\n</tool_call>"
        return out
    new_lines = []
    for line in out.split("\n"):
        stripped = line.strip()
        if (stripped.startswith("{")
                and stripped.endswith("}")
                and '"name"' in stripped
                and '"arguments"' in stripped):
            indent = line[: len(line) - len(line.lstrip())]
            new_lines.append(f"{indent}<tool_call>\n{stripped}\n</tool_call>")
        else:
            new_lines.append(line)
    return "\n".join(new_lines)
