import ast
import re
from collections import defaultdict
from io import BytesIO
from typing import Any, Dict, List, Optional, Tuple, Union

import pytesseract
from PIL import Image
from pytesseract import Output

from mm_agents.vlaa_gui.memory.procedural_memory import PROCEDURAL_MEMORY
from mm_agents.vlaa_gui.core.mllm import LMMAgent
from mm_agents.vlaa_gui.utils.common_utils import (
    call_llm_safe,
    call_llm_safe_with_thinking,
    parse_single_code_from_string,
    resize_screenshot,
)
from mm_agents.vlaa_gui.agents.coding_agent import CodeAgent
from mm_agents.vlaa_gui.agents.search_agent import SearcherAgent

import logging

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


class ACI:
    def __init__(self):
        self.notes: List[str] = []


# Agent action decorator
def agent_action(func):
    func.is_agent_action = (
        True  # used to identify agent action methods to pass to the procedural memory
    )
    return func


def agent_perception(func):
    func.is_agent_perception = True
    return func


UBUNTU_APP_SETUP = """import subprocess;
import difflib;
import pyautogui;
pyautogui.press('escape');
time.sleep(0.5);
output = subprocess.check_output(['wmctrl', '-lx']);
output = output.decode('utf-8').splitlines();
window_titles = [line.split(None, 4)[2] for line in output];
closest_matches = difflib.get_close_matches('APP_NAME', window_titles, n=1, cutoff=0.1);
if closest_matches:
    closest_match = closest_matches[0];
    for line in output:
        if closest_match in line:
            window_id = line.split()[0]
            break;
subprocess.run(['wmctrl', '-ia', window_id])
subprocess.run(['wmctrl', '-ir', window_id, '-b', 'add,maximized_vert,maximized_horz'])
"""


SET_CELL_VALUES_CMD = """import uno
import subprocess

def identify_document_type(component):
    if component.supportsService("com.sun.star.sheet.SpreadsheetDocument"):
        return "Calc"

    if component.supportsService("com.sun.star.text.TextDocument"):
        return "Writer"

    if component.supportsService("com.sun.star.sheet.PresentationDocument"):
        return "Impress"

    return None

def cell_ref_to_indices(cell_ref):
    column_letters = ''.join(filter(str.isalpha, cell_ref))
    row_number = ''.join(filter(str.isdigit, cell_ref))

    col = sum((ord(char.upper()) - ord('A') + 1) * (26**idx) for idx, char in enumerate(reversed(column_letters))) - 1
    row = int(row_number) - 1
    return col, row

def set_cell_values(new_cell_values: dict[str, str], app_name: str = "Untitled 1", sheet_name: str = "Sheet1"):
    new_cell_values_idx = {{}}
    for k, v in new_cell_values.items():
        try:
            col, row = cell_ref_to_indices(k)
        except:
            col = row = None

        if col is not None and row is not None:
            new_cell_values_idx[(col, row)] = v

    # Clean up previous TCP connections.
    subprocess.run(
        'echo \"osworld-public-evaluation\" | sudo -S ss --kill --tcp state TIME-WAIT sport = :2002',
        shell=True,
        check=True,
        text=True,
        capture_output=True
    )

    # Dynamically allow soffice to listen on port 2002.
    subprocess.run(
        [
            "soffice",
            "--accept=socket,host=localhost,port=2002;urp;StarOffice.Service"
        ]
    )

    local_context = uno.getComponentContext()
    resolver = local_context.ServiceManager.createInstanceWithContext(
        "com.sun.star.bridge.UnoUrlResolver", local_context
    )
    context = resolver.resolve(
        f"uno:socket,host=localhost,port=2002;urp;StarOffice.ComponentContext"
    )
    desktop = context.ServiceManager.createInstanceWithContext(
        "com.sun.star.frame.Desktop", context
    )

    # Collect all LibreOffice-related opened windows.
    documents = []
    for i, component in enumerate(desktop.Components):
        title = component.Title
        doc_type = identify_document_type(component)
        documents.append((i, component, title, doc_type))

    # Find the LibreOffice Calc app and the sheet of interest.
    spreadsheet = [doc for doc in documents if doc[3] == "Calc"]
    selected_spreadsheet = [doc for doc in spreadsheet if doc[2] == app_name]
    if spreadsheet:
        try:
            if selected_spreadsheet:
                spreadsheet = selected_spreadsheet[0][1]
            else:
                spreadsheet = spreadsheet[0][1]

            sheet = spreadsheet.Sheets.getByName(sheet_name)
        except:
            raise ValueError(f"Could not find sheet {{sheet_name}} in {{app_name}}.")

        for (col, row), value in new_cell_values_idx.items():
            cell = sheet.getCellByPosition(col, row)

            # Set the cell value.
            if isinstance(value, (int, float)):
                cell.Value = value
            elif isinstance(value, str):
                if value.startswith("="):
                    cell.Formula = value
                else:
                    cell.String = value
            elif isinstance(value, bool):
                cell.Value = 1 if value else 0
            elif value is None:
                cell.clearContents(0)
            else:
                raise ValueError(f"Unsupported cell value type: {{type(value)}}")

    else:
        raise ValueError(f"Could not find LibreOffice Calc app corresponding to {{app_name}}.")

set_cell_values(new_cell_values={cell_values}, app_name="{app_name}", sheet_name="{sheet_name}")        
"""


# ACI primitives are parameterized by description, and coordinate generation uses a pretrained grounding model
class OSWorldACI(ACI):
    def __init__(
        self,
        platform: str,
        engine_params_for_generation: dict,
        engine_params_for_grounding: dict,
        engine_params_for_searcher: dict,
        env=None,
        width: int = 1920,
        height: int = 1080,
        resize_width: Optional[int] = None,
        observation_type: str = "screenshot",
        grounding_model_type: str = "MoG",  # "MoG" or "unified"
        code_agent_engine_params: dict = None,
        code_agent_budget: int = 20,
        debug: bool = False,
    ):
        self.platform = (
            platform  # Dictates how the switch_applications agent action works.
        )
        self.env = env
        self.engine_params_for_generation = engine_params_for_generation
        self.engine_params_for_grounding = engine_params_for_grounding
        self.engine_params_for_searcher = engine_params_for_searcher
        self.code_agent_engine_params = code_agent_engine_params
        # Search agent state (created lazily)
        self.searcher_agent = None
        self.last_search_agent_result = None
        self.last_code_agent_result = None
        self.tutorials: list[str] = []
        # Configure scaling
        self.width = width
        self.height = height
        self.resize_width = resize_width
        # Maintain state for save_to_knowledge
        self.notes = []
        # Screenshot used during ACI execution
        self.obs = None

        # Coordinates used during ACI execution
        self.coords1 = None
        self.coords2 = None

        self.debug = debug

        # Configure the visual grounding model responsible for coordinate generation
        self.grounding_model = LMMAgent(engine_params_for_grounding)
        self.grounding_model_type = grounding_model_type
        self.last_grounding_thoughts: Optional[str] = None

        # Configure text grounding agent
        self.text_span_agent = LMMAgent(
            engine_params=engine_params_for_generation,
            system_prompt=PROCEDURAL_MEMORY.PHRASE_TO_WORD_COORDS_PROMPT,
        )

        # Configure code agent
        self.code_agent = CodeAgent(
            code_agent_engine_params, code_agent_budget, echo=debug
        )

        self.current_task_instruction = None
        self._code_agent_constraint_fail_count = 0
        self.enable_zoom_grounding = False

        self.observation_type = observation_type

    def _get_grounding_coordinate_space(self) -> Tuple[int, int]:
        if {"grounding_width", "grounding_height"}.issubset(
            self.engine_params_for_grounding
        ):
            return (
                int(self.engine_params_for_grounding["grounding_width"]),
                int(self.engine_params_for_grounding["grounding_height"]),
            )
        return (1000, 1000)

    def _parse_xy_from_grounding_response(
        self, response: str
    ) -> Optional[Tuple[int, int]]:
        if not response:
            return None
        numericals = re.findall(r"-?\d+", response)
        if len(numericals) < 2:
            return None
        coord_w, coord_h = self._get_grounding_coordinate_space()
        x = max(0, min(int(numericals[0]), coord_w - 1))
        y = max(0, min(int(numericals[1]), coord_h - 1))
        return (x, y)

    def _ground_single_pass(
        self,
        ref_expr: str,
        screenshot_bytes: bytes,
        prompt_prefix: str = "",
    ) -> Tuple[int, int]:
        self.grounding_model.reset()
        coord_w, coord_h = self._get_grounding_coordinate_space()
        prompt = (
            f"{prompt_prefix}"
            # f"The image uses a coordinate system from (0, 0) at the top-left to ({coord_w}, {coord_h}) at the bottom-right.\n"
            f"Query: {ref_expr}\n"
            "Output only the coordinate of one point in your response.\n"
        )
        self.grounding_model.add_message(
            text_content=prompt,
            image_content=screenshot_bytes,
            put_text_last=True,
        )
        if getattr(self.grounding_model.engine, "thinking", False):
            response, thoughts = call_llm_safe_with_thinking(
                self.grounding_model,
                agent_type="grounding",
            )
        else:
            response = call_llm_safe(
                self.grounding_model,
                agent_type="grounding",
            )
            thoughts = None

        self.last_grounding_thoughts = thoughts

        logger.info(f"Grounding model response: {response}")

        response = response if isinstance(response, str) else str(response)

        # parsed = self._parse_xy_from_grounding_response(response)
        # if parsed is None:
        #     raise RuntimeError(f"Unable to parse grounding coordinates: {response}")
        # return parsed
        point_match = re.search(r"<point>\s*(\d+)\s+(\d+)\s*</point>", response)
        if point_match:
            return [int(point_match.group(1)), int(point_match.group(2))]
        numericals = re.findall(r"\d+", response)
        assert len(numericals) >= 2
        return [int(numericals[0]), int(numericals[1])]

    def _crop_centered_and_resize(
        self,
        screenshot_bytes: bytes,
        center_px: Tuple[int, int],
        crop_ratio: float,
    ) -> Tuple[bytes, Tuple[int, int], Tuple[int, int], Tuple[int, int]]:
        """
        Crop a window centered on `center_px` with (w,h)=crop_ratio*(W,H) and resize back to (W,H).

        Returns:
          zoomed_bytes, (left_px, top_px), (crop_w_px, crop_h_px), (img_w, img_h)
        """
        img = Image.open(BytesIO(screenshot_bytes))
        img_w, img_h = img.size

        crop_w_px = max(1, int(round(img_w * crop_ratio)))
        crop_h_px = max(1, int(round(img_h * crop_ratio)))
        crop_w_px = min(crop_w_px, img_w)
        crop_h_px = min(crop_h_px, img_h)

        cx, cy = center_px
        cx = max(0, min(cx, img_w - 1))
        cy = max(0, min(cy, img_h - 1))

        left_px = int(round(cx - crop_w_px / 2))
        top_px = int(round(cy - crop_h_px / 2))
        left_px = max(0, min(left_px, img_w - crop_w_px))
        top_px = max(0, min(top_px, img_h - crop_h_px))

        cropped = img.crop((left_px, top_px, left_px + crop_w_px, top_px + crop_h_px))
        try:
            resample = Image.Resampling.BICUBIC  # Pillow >= 9
        except AttributeError:  # pragma: no cover
            resample = Image.BICUBIC
        zoomed = cropped.resize((img_w, img_h), resample=resample)

        output = BytesIO()
        zoomed.save(output, format=img.format or "PNG")
        return (
            output.getvalue(),
            (left_px, top_px),
            (crop_w_px, crop_h_px),
            (img_w, img_h),
        )

    # Given the state and worker's referring expression, use the grounding model to generate (x,y)
    def generate_coords(self, ref_expr: str, obs: Dict) -> List[int]:
        if self.resize_width is not None:
            resize_height = int(self.resize_width / self.width * self.height)
            new_size = (self.resize_width, resize_height)
            screenshot_bytes = resize_screenshot(obs["screenshot"], new_size)
        else:
            screenshot_bytes = obs["screenshot"]

        coord_w, coord_h = self._get_grounding_coordinate_space()

        # Pass 1: coarse prediction on full image
        print(f"Generating coordinates for: {ref_expr}")
        coarse_x, coarse_y = self._ground_single_pass(ref_expr, screenshot_bytes)

        if not self.enable_zoom_grounding:
            return [coarse_x, coarse_y]

        # Pass 2: crop a zoomed-in window around the coarse point and refine
        try:
            img = Image.open(BytesIO(screenshot_bytes))
            img_w, img_h = img.size
        except Exception:
            return [coarse_x, coarse_y]

        if img_w <= 1 or img_h <= 1:
            return [coarse_x, coarse_y]

        center_px = (
            int(round(coarse_x * img_w / max(coord_w, 1))),
            int(round(coarse_y * img_h / max(coord_h, 1))),
        )

        try:
            zoomed_bytes, (left_px, top_px), (crop_w_px, crop_h_px), _ = (
                self._crop_centered_and_resize(
                    screenshot_bytes=screenshot_bytes,
                    center_px=center_px,
                    crop_ratio=self.zoom_grounding_crop_ratio,
                )
            )
        except Exception:
            return [coarse_x, coarse_y]

        try:
            zoom_x, zoom_y = self._ground_single_pass(
                ref_expr,
                zoomed_bytes,
                prompt_prefix="This is a zoomed-in crop of the original screenshot centered near the target.\n",
            )
        except Exception:
            return [coarse_x, coarse_y]

        scale_x = crop_w_px / img_w
        scale_y = crop_h_px / img_h
        left_coord = left_px / img_w * coord_w
        top_coord = top_px / img_h * coord_h

        refined_x = int(round(left_coord + zoom_x * scale_x))
        refined_y = int(round(top_coord + zoom_y * scale_y))

        refined_x = max(0, min(refined_x, coord_w - 1))
        refined_y = max(0, min(refined_y, coord_h - 1))
        return [refined_x, refined_y]

    # Calls pytesseract to generate word level bounding boxes for text grounding
    def get_ocr_elements(self, b64_image_data: str) -> Tuple[str, List]:
        image = Image.open(BytesIO(b64_image_data))
        image_data = pytesseract.image_to_data(image, output_type=Output.DICT)

        # Clean text by removing leading and trailing spaces and non-alphabetical characters, but keeping punctuation
        for i, word in enumerate(image_data["text"]):
            image_data["text"][i] = re.sub(
                r"^[^a-zA-Z\s.,!?;:\-\+]+|[^a-zA-Z\s.,!?;:\-\+]+$", "", word
            )

        ocr_elements = []
        ocr_table = "Text Table:\nWord id\tText\n"
        # Obtain the <id, text, group number, word number> for each valid element
        grouping_map = defaultdict(list)
        ocr_id = 0
        for i in range(len(image_data["text"])):
            block_num = image_data["block_num"][i]
            if image_data["text"][i]:
                grouping_map[block_num].append(image_data["text"][i])
                ocr_table += f"{ocr_id}\t{image_data['text'][i]}\n"
                ocr_elements.append(
                    {
                        "id": ocr_id,
                        "text": image_data["text"][i],
                        "group_num": block_num,
                        "word_num": len(grouping_map[block_num]),
                        "left": image_data["left"][i],
                        "top": image_data["top"][i],
                        "width": image_data["width"][i],
                        "height": image_data["height"][i],
                    }
                )
                ocr_id += 1

        return ocr_table, ocr_elements

    # Given the state and worker's text phrase, generate the coords of the first/last word in the phrase
    def generate_text_coords(
        self, phrase: str, obs: Dict, alignment: str = ""
    ) -> List[int]:
        ocr_table, ocr_elements = self.get_ocr_elements(obs["screenshot"])

        alignment_prompt = ""
        if alignment == "start":
            alignment_prompt = "**Important**: Output the word id of the FIRST word in the provided phrase.\n"
        elif alignment == "end":
            alignment_prompt = "**Important**: Output the word id of the LAST word in the provided phrase.\n"

        # Load LLM prompt
        self.text_span_agent.reset()
        self.text_span_agent.add_message(
            alignment_prompt + "Phrase: " + phrase + "\n" + ocr_table, role="user"
        )
        self.text_span_agent.add_message(
            "Screenshot:\n", image_content=obs["screenshot"], role="user"
        )

        # Obtain the target element
        response = call_llm_safe(self.text_span_agent)
        print("TEXT SPAN AGENT RESPONSE:", response)
        numericals = re.findall(r"\d+", response)
        if len(numericals) > 0:
            text_id = int(numericals[-1])
        else:
            text_id = 0
        elem = ocr_elements[text_id]

        # Compute the element coordinates
        if alignment == "start":
            coords = [elem["left"], elem["top"] + (elem["height"] // 2)]
        elif alignment == "end":
            coords = [elem["left"] + elem["width"], elem["top"] + (elem["height"] // 2)]
        else:
            coords = [
                elem["left"] + (elem["width"] // 2),
                elem["top"] + (elem["height"] // 2),
            ]
        return coords

    # Given the state and worker's text phrase, generate the coords using only LLM (no OCR)
    # Returns coordinates in screen pixel space (already resized), matching generate_text_coords behavior.
    def generate_text_coords_mllm(
        self, phrase: str, obs: Dict, alignment: str = ""
    ) -> List[int]:
        # Reset the grounding model state
        self.grounding_model.reset()

        coord_w, coord_h = self._get_grounding_coordinate_space()

        # Build the prompt based on alignment
        coord_space_prefix = f"The image uses a coordinate system from (0, 0) at the top-left to ({coord_w}, {coord_h}) at the bottom-right.\n"
        if alignment == "start":
            prompt = f"{coord_space_prefix}Find the coordinates of the START position of the text phrase '{phrase}' in the screenshot. Output only the x,y coordinates."
        elif alignment == "end":
            prompt = f"{coord_space_prefix}Find the coordinates of the END position of the text phrase '{phrase}' in the screenshot. Output only the x,y coordinates."
        else:
            prompt = f"{coord_space_prefix}Find the coordinates of the CENTER position of the text phrase '{phrase}' in the screenshot. Output only the x,y coordinates."

        self.grounding_model.add_message(
            text_content=prompt, image_content=obs["screenshot"], put_text_last=True
        )
        print(
            f"Generating text coordinates for phrase: {phrase} (alignment: {alignment})"
        )

        # Generate and parse coordinates
        response = call_llm_safe(self.grounding_model)
        print("RAW MLLM TEXT COORDS RESPONSE:", response)
        numericals = re.findall(r"\d+", response)
        if len(numericals) >= 2:
            grounding_coords = [int(numericals[0]), int(numericals[1])]
            # Convert from grounding space to screen pixel space so that
            # highlight_text_span (which uses coords directly) works correctly.
            return self.resize_coordinates(grounding_coords)
        else:
            # Fallback: raise an error if we can't parse coordinates
            raise ValueError(
                f"Could not parse coordinates from MLLM response: {response}"
            )

    # Takes a description based action and assigns the coordinates for any coordinate based action
    # Raises an error if function can't be parsed
    @agent_perception
    def assign_coordinates(self, plan: str, obs: Dict):
        assert self.observation_type == "screenshot" or self.observation_type == "mixed"
        # Reset coords from previous action generation
        self.coords1, self.coords2 = None, None

        try:
            # Extract the function name and args
            action = parse_single_code_from_string(plan.split("Grounded Action")[-1])
            function_name = re.match(r"(\w+\.\w+)\(", action).group(1)
            args = self.parse_function_args(action)
        except Exception as e:
            raise RuntimeError(f"Error in parsing grounded action: {e}") from e

        print(f"Generating coordinates for {function_name}")
        # arg0 is a description
        # For actions requiring one coordinate
        if (
            function_name in ["agent.click", "agent.type", "agent.scroll"]
            and len(args) >= 1
            and args[0] is not None
        ):
            self.coords1 = self.generate_coords(args[0], obs)
        # arg0 and arg1 are descriptions
        # For actions requiring two coordinates
        elif function_name == "agent.drag_and_drop" and len(args) >= 2:
            self.coords1 = self.generate_coords(args[0], obs)
            self.coords2 = self.generate_coords(args[1], obs)
        # arg0 and arg1 are text phrases
        # For text-based highlighting
        elif function_name == "agent.highlight_text_span" and len(args) >= 2:
            if self.grounding_model_type == "MoG":
                self.coords1 = self.generate_text_coords(
                    args[0], obs, alignment="start"
                )
                self.coords2 = self.generate_text_coords(args[1], obs, alignment="end")
            elif self.grounding_model_type == "unified":
                self.coords1 = self.generate_text_coords_mllm(
                    args[0], obs, alignment="start"
                )
                self.coords2 = self.generate_text_coords_mllm(
                    args[1], obs, alignment="end"
                )

    def assign_screenshot(self, obs: Dict):
        self.obs = obs

    def set_task_instruction(self, task_instruction: str):
        """Set the current task instruction for the code agent."""
        self.current_task_instruction = task_instruction

    # Resize from grounding model dim into OSWorld dim (1920 * 1080)
    def resize_coordinates(self, coordinates: List[int]) -> List[int]:
        # User explicitly passes the grounding model dimensions
        if {"grounding_width", "grounding_height"}.issubset(
            self.engine_params_for_grounding
        ):
            grounding_width = self.engine_params_for_grounding["grounding_width"]
            grounding_height = self.engine_params_for_grounding["grounding_height"]
        # Default to (1000, 1000), which is UI-TARS resizing
        else:
            print("Using UI-TARS grounding model dimensions")
            grounding_width = 1000
            grounding_height = 1000

        return [
            round(coordinates[0] * self.width / grounding_width),
            round(coordinates[1] * self.height / grounding_height),
        ]

    # Given a generated ACI function, returns a list of argument values, where descriptions are at the front of the list
    def parse_function_args(self, function: str) -> List[str]:
        tree = ast.parse(function)
        call_node = tree.body[0].value

        def safe_eval(node):
            if isinstance(
                node, ast.Constant
            ):  # Handles literals like numbers, strings, etc.
                return node.value
            else:
                return ast.unparse(node)  # Return as a string if not a literal

        positional_args = [safe_eval(arg) for arg in call_node.args]
        keyword_args = {kw.arg: safe_eval(kw.value) for kw in call_node.keywords}

        res = []

        for key, val in keyword_args.items():
            if "description" in key:
                res.append(val)

        for arg in positional_args:
            res.append(arg)

        return res

    @agent_action
    def click(
        self,
        element_description: str,
        num_clicks: int = 1,
        button_type: str = "left",
        hold_keys: List = [],
    ):
        """Click on the element
        Args:
            element_description:str, a detailed descriptions of which element to click on. This description should be at least a full sentence.
            num_clicks:int, number of times to click the element
            button_type:str, which mouse button to press can be "left", "middle", or "right"
            hold_keys:List, list of keys to hold while clicking
        """
        coords1 = self.coords1
        if coords1 is None:
            coords1 = self.generate_coords(element_description, self.obs)

        x, y = self.resize_coordinates(coords1)

        command = "import pyautogui; "

        for k in hold_keys:
            command += f"pyautogui.keyDown({repr(k)}); "
        command += f"""import pyautogui; pyautogui.click({x}, {y}, clicks={num_clicks}, button={repr(button_type)}); """
        for k in hold_keys:
            command += f"pyautogui.keyUp({repr(k)}); "
        # Return pyautoguicode to click on the element
        return command

    @agent_action
    def double_click(
        self,
        element_description: str,
        hold_keys: List = [],
    ):
        """Double-click on the element
        Args:
            element_description:str, a detailed descriptions of which element to double-click on. This description should be at least a full sentence.
            hold_keys:List, list of keys to hold while double-clicking
        """
        return self.click(
            element_description, num_clicks=2, button_type="left", hold_keys=hold_keys
        )

    @agent_action
    def switch_applications(self, app_code):
        """Switch to a different application that is already open
        Args:
            app_code:str the code name of the application to switch to from the provided list of open applications
        """
        if self.platform == "darwin":
            return f"import pyautogui; import time; pyautogui.hotkey('command', 'space', interval=0.5); pyautogui.typewrite({repr(app_code)}); pyautogui.press('enter'); time.sleep(1.0)"
        elif self.platform == "linux":
            return UBUNTU_APP_SETUP.replace("APP_NAME", app_code)
        elif self.platform == "windows":
            return f"import pyautogui; import time; pyautogui.hotkey('win', 'd', interval=0.5); pyautogui.typewrite({repr(app_code)}); pyautogui.press('enter'); time.sleep(1.0)"
        else:
            assert False, (
                f"Unsupported platform: {self.platform}. Supported platforms are: darwin, linux, windows."
            )

    @agent_action
    def open(self, app_or_filename: str):
        """Open any application or file with name app_or_filename. Use this action to open applications or files on the desktop, do not open manually.
        Args:
            app_or_filename:str, the name of the application or filename to open
        """
        if self.platform == "linux":
            return f"import pyautogui; pyautogui.hotkey('win'); time.sleep(0.5); pyautogui.write({repr(app_or_filename)}); time.sleep(1.0); pyautogui.hotkey('enter'); time.sleep(0.5)"
        elif self.platform == "darwin":
            return f"import pyautogui; import time; pyautogui.hotkey('command', 'space', interval=0.5); pyautogui.typewrite({repr(app_or_filename)}); pyautogui.press('enter'); time.sleep(1.0)"
        elif self.platform == "windows":
            return (
                "import pyautogui; import time; "
                "pyautogui.hotkey('win'); time.sleep(0.5); "
                f"pyautogui.write({repr(app_or_filename)}); time.sleep(1.0); "
                "pyautogui.press('enter'); time.sleep(0.5)"
            )
        else:
            assert False, (
                f"Unsupported platform: {self.platform}. Supported platforms are: darwin, linux, windows."
            )

    @agent_action
    def type(
        self,
        element_description: Optional[str] = None,
        text: str = "",
        overwrite: bool = False,
        enter: bool = False,
    ):
        """Type text/unicode into a specific element
        Args:
            element_description:str, a detailed description of which element to enter text in. This description should be at least a full sentence.
            text:str, the text to type
            overwrite:bool, Assign it to True if the text should overwrite the existing text, otherwise assign it to False. Using this argument clears all text in an element.
            enter:bool, Assign it to True if the enter key should be pressed after typing the text, otherwise assign it to False.
        """
        command = "import pyautogui; "
        command += (
            "\ntry:\n"
            "    import pyperclip\n"
            "except ImportError:\n"
            "    import subprocess\n"
            "    subprocess.run('echo \"osworld-public-evaluation\" | sudo -S apt-get install -y xclip xsel', shell=True, check=True)\n"
            "    subprocess.check_call([subprocess.sys.executable, '-m', 'pip', 'install', 'pyperclip'])\n"
            "    import pyperclip\n\n"
        )

        if element_description is not None:
            coords1 = self.coords1
            if coords1 is None:
                coords1 = self.generate_coords(element_description, self.obs)
            x, y = self.resize_coordinates(coords1)
            command += f"pyautogui.click({x}, {y}); "

        if overwrite:
            command += (
                f"pyautogui.hotkey({repr('command' if self.platform == 'darwin' else 'ctrl')}, 'a'); "
                "pyautogui.press('backspace'); "
            )

        # Check if text contains Unicode characters that pyautogui.write() can't handle
        has_unicode = any(ord(char) > 127 for char in text)

        if has_unicode:
            # Use clipboard method for Unicode characters
            command += f"pyperclip.copy({repr(text)}); "
            command += f"pyautogui.hotkey({repr('command' if self.platform == 'darwin' else 'ctrl')}, 'v'); "
        else:
            # Use regular pyautogui.write() for ASCII text
            command += f"pyautogui.write({repr(text)}); "

        if enter:
            command += "pyautogui.press('enter'); "
        return command

    @agent_action
    def save_to_knowledge(self, text: List[str]):
        """Save facts, elements, texts, etc. to a long-term knowledge bank for reuse during this task. Can be used for copy-pasting text, saving elements, etc.
        Args:
            text:List[str] the text to save to the knowledge
        """
        self.notes.extend(text)
        return """WAIT"""

    @agent_action
    def drag_and_drop(
        self, starting_description: str, ending_description: str, hold_keys: List = []
    ):
        """Drag from the starting description to the ending description
        Args:
            starting_description:str, a very detailed description of where to start the drag action. This description should be at least a full sentence.
            ending_description:str, a very detailed description of where to end the drag action. This description should be at least a full sentence.
            hold_keys:List list of keys to hold while dragging
        """
        coords1 = (
            self.coords1
            if self.coords1 is not None
            else self.generate_coords(starting_description, self.obs)
        )
        coords2 = (
            self.coords2
            if self.coords2 is not None
            else self.generate_coords(ending_description, self.obs)
        )
        x1, y1 = self.resize_coordinates(coords1)
        x2, y2 = self.resize_coordinates(coords2)

        command = "import pyautogui; "

        command += f"pyautogui.moveTo({x1}, {y1}); "
        # TODO: specified duration?
        for k in hold_keys:
            command += f"pyautogui.keyDown({repr(k)}); "
        command += f"pyautogui.dragTo({x2}, {y2}, duration=1., button='left'); pyautogui.mouseUp(); "
        for k in hold_keys:
            command += f"pyautogui.keyUp({repr(k)}); "

        # Return pyautoguicode to drag and drop the elements

        return command

    @agent_action
    def highlight_text_span(
        self, starting_phrase: str, ending_phrase: str, button: str = "left"
    ):
        """Highlight a text span between a provided starting phrase and ending phrase. Use this to highlight words, lines, and paragraphs.
        Args:
            starting_phrase:str, the phrase that denotes the start of the text span you want to highlight. If you only want to highlight one word, just pass in that single word.
            ending_phrase:str, the phrase that denotes the end of the text span you want to highlight. If you only want to highlight one word, just pass in that single word.
            button:str, the button to use to highlight the text span. Defaults to "left". Can be "left", "right", or "middle".
        """
        coords1 = (
            self.coords1
            if self.coords1 is not None
            else self.generate_text_coords(starting_phrase, self.obs, alignment="start")
        )
        coords2 = (
            self.coords2
            if self.coords2 is not None
            else self.generate_text_coords(ending_phrase, self.obs, alignment="end")
        )
        x1, y1 = coords1
        x2, y2 = coords2

        command = "import pyautogui; "
        command += f"pyautogui.moveTo({x1}, {y1}); "
        command += f"pyautogui.dragTo({x2}, {y2}, duration=1., button='{button}'); pyautogui.mouseUp(); "

        # Return pyautoguicode to drag and drop the elements
        return command

    @agent_action
    def set_cell_values(
        self, cell_values: Dict[str, Any], app_name: str, sheet_name: str
    ):
        """Use this to set individual cell values in a spreadsheet. For example, setting A2 to "hello" would be done by passing {"A2": "hello"} as cell_values. The sheet must be opened before this command can be used.
        Args:
            cell_values: Dict[str, Any], A dictionary of cell values to set in the spreadsheet. The keys are the cell coordinates in the format "A1", "B2", etc.
                Supported value types include: float, int, string, bool, formulas.
            app_name: str, The name of the spreadsheet application. For example, "Some_sheet.xlsx".
            sheet_name: str, The name of the sheet in the spreadsheet. For example, "Sheet1".
        """
        return SET_CELL_VALUES_CMD.format(
            cell_values=cell_values, app_name=app_name, sheet_name=sheet_name
        )

    @agent_action
    def scroll(self, element_description: str, clicks: int, shift: bool = False):
        """Scroll the element in the specified direction
        Args:
            element_description:str, a very detailed description of which element to enter scroll in. This description should be at least a full sentence.
            clicks:int, the number of clicks to scroll can be positive (up) or negative (down).
            shift:bool, whether to use shift+scroll for horizontal scrolling
        """
        coords1 = self.generate_coords(element_description, self.obs)
        x, y = self.resize_coordinates(coords1)

        if shift:
            return f"import pyautogui; import time; pyautogui.moveTo({x}, {y}); time.sleep(0.5); pyautogui.hscroll({clicks})"
        else:
            return f"import pyautogui; import time; pyautogui.moveTo({x}, {y}); time.sleep(0.5); pyautogui.vscroll({clicks})"

    @agent_action
    def call_code_agent(self, task: str = None):
        """Call the code agent to execute code for tasks or subtasks that can be completed solely with coding.

        Args:
            task: str, the task or subtask to execute. If None, uses the current full task instruction.

        **🚨 CRITICAL GUIDELINES:**
        - **ONLY pass a task parameter for SPECIFIC subtasks** (e.g., "Calculate sum of column B", "Filter data by date")
        - **NEVER pass a task parameter for full tasks** - let it default to the original task instruction
        - **NEVER rephrase or modify the original task** - this prevents hallucination corruption
        - **If unsure, omit the task parameter entirely** to use the original task instruction

        Use this for tasks that can be fully accomplished through code execution, particularly for:
        - Spreadsheet applications (LibreOffice Calc, Excel): data processing, filtering, sorting, calculations, formulas, data analysis
        - Document editors (LibreOffice Writer, Word): text processing, content editing, formatting, document manipulation
        - Code editors (VS Code, text editors): code editing, file processing, text manipulation, configuration
        - Data analysis tools: statistical analysis, data transformation, reporting
        - File management: bulk operations, file processing, content extraction
        - System utilities: configuration, setup, automation
        """
        logger.info("=" * 50)
        logger.info("GROUNDING AGENT: Calling Code Agent")
        logger.info("=" * 50)

        # **CRITICAL**: Only use provided task for specific subtasks, otherwise use original task instruction
        if task is not None:
            # This is a subtask — use the provided task but carry over
            # constraints from the original instruction so the code agent
            # does not bypass application-level restrictions.
            task_to_execute = task
            if self.current_task_instruction:
                from mm_agents.vlaa_gui.agents.coding_agent import extract_constraints

                parent_constraints = extract_constraints(self.current_task_instruction)
                if parent_constraints:
                    constraint_note = (
                        "\n\nIMPORTANT — The original task has the following constraints "
                        "(these apply to this subtask too):\n"
                        + "\n".join(f"  - {c}" for c in parent_constraints)
                    )
                    task_to_execute = task + constraint_note
            logger.info(f"Executing SUBTASK: {task_to_execute}")
        else:
            # This is a full task - use the original task instruction to prevent hallucination
            task_to_execute = self.current_task_instruction
            logger.info(f"Executing FULL TASK: {task_to_execute}")

        if task_to_execute:
            # If the code agent previously FAILed due to constraints on this
            # task, don't let the worker retry — return the FAIL immediately.
            from mm_agents.vlaa_gui.agents.coding_agent import extract_constraints

            has_constraints = bool(
                extract_constraints(self.current_task_instruction or "")
            )
            if has_constraints and self._code_agent_constraint_fail_count >= 1:
                logger.warning(
                    "Code agent already FAILed %d time(s) due to task constraints. "
                    "Returning cached FAIL to prevent constraint override.",
                    self._code_agent_constraint_fail_count,
                )
                self.last_code_agent_result = {
                    "task_instruction": task_to_execute,
                    "completion_reason": "FAIL",
                    "summary": (
                        "The code agent previously determined this task is infeasible "
                        "due to constraints in the original instruction. The task "
                        "cannot be accomplished without violating those constraints."
                    ),
                    "execution_history": [],
                    "steps_executed": 0,
                    "budget": getattr(self.code_agent, "budget", 0),
                }
                return "import time; time.sleep(2.222)"

            logger.info(
                "Observation keys: %s", list(self.obs.keys()) if self.obs else []
            )
            screenshot = self.obs.get("screenshot", "") if self.obs else ""
            logger.info(f"Screenshot available: {'Yes' if screenshot else 'No'}")

            logger.info("Executing code agent...")
            try:
                result = self.code_agent.execute(
                    task_to_execute, screenshot, self.env.controller
                )
            except Exception as exc:
                logger.exception("Code agent execution failed")
                result = {
                    "task_instruction": task_to_execute,
                    "completion_reason": "ERROR",
                    "summary": f"Code agent execution failed: {exc}",
                    "execution_history": [],
                    "steps_executed": 0,
                    "budget": getattr(self.code_agent, "budget", 0),
                    "error": str(exc),
                }

            if not isinstance(result, dict):
                result = {
                    "task_instruction": task_to_execute,
                    "completion_reason": "ERROR",
                    "summary": f"Code agent returned non-dict result: {type(result).__name__}",
                    "execution_history": [],
                    "steps_executed": 0,
                    "budget": getattr(self.code_agent, "budget", 0),
                }

            # Normalize required fields so downstream reasoning always receives
            # an explicit success/failure signal.
            result.setdefault("task_instruction", task_to_execute)
            result.setdefault("completion_reason", "UNKNOWN")
            result.setdefault("summary", "")
            result.setdefault("execution_history", [])
            result.setdefault("steps_executed", 0)
            result.setdefault("budget", getattr(self.code_agent, "budget", 0))

            # Track constraint-related FAILs to prevent override on retry
            if (
                has_constraints
                and str(result.get("completion_reason", "")).upper() == "FAIL"
            ):
                self._code_agent_constraint_fail_count += 1
                logger.info(
                    "Code agent FAILed with active constraints (count=%d).",
                    self._code_agent_constraint_fail_count,
                )

            # Store the result for the worker to access
            self.last_code_agent_result = result

            logger.info("Code agent execution completed")
            logger.info(
                f"Result - Completion reason: {result.get('completion_reason')}"
            )
            logger.info(f"Steps executed: {result.get('steps_executed')}")
            logger.info(f"Summary: {result.get('summary')}")

            logger.info("=" * 50)
            logger.info("GROUNDING AGENT: Code Agent Call Finished")
            logger.info("=" * 50)

            # Return code to be executed in the environment
            return "import time; time.sleep(2.222)"
        else:
            logger.warning("No task instruction available for code agent call")
            self.last_code_agent_result = {
                "task_instruction": None,
                "completion_reason": "ERROR",
                "summary": "Code agent was invoked without a task instruction.",
                "execution_history": [],
                "steps_executed": 0,
                "budget": getattr(self.code_agent, "budget", 0),
            }
            return "import time; time.sleep(1.111)"

    @agent_action
    def call_search_agent(self, query: str):
        """
        Calls a specialized 'Searcher Agent' to find a detailed, step-by-step tutorial on the internet for a specific GUI action.
        Args:
            query:str, the search phrase or question for the tutorial. The formulation of this query is critical for success and must follow the guidelines below.

        **Query Formulation Guidelines:**

        Your query must be a well-defined question targeting a **single, specific action** within a **specific application**. To get the best results, adhere to these rules:

        1.  **Start with "How to":** Your query must begin with the phrase "How to" to frame it as a request for instructions.
        2.  **Include the Application Name:** Always specify the name of the software you are working in (e.g., "GIMP", "Google Chrome", "Libreoffice Writer").
        3.  **Focus on a Single Intent:** The query should represent one clear goal. Do not combine multiple steps or tasks into one query.
        4.  **Be Specific, Not Abstract:** Ask a concrete question. Avoid repeating the user's high-level or abstract instructions.
        5.  **Decompose Complex Tasks:** If the user's overall instruction involves multiple actions (e.g., "download a file and then email it"), and you are stuck on one part, search *only for that specific part*.

        **Examples:**

        *   **User's Overall Instruction:** "Please help me download my latest bank statement and then send it to my accountant."
            *   **Correct Query (if stuck on downloading):** "How to download a bank statement from the Bank of America website?"
            *   **Correct Query (if stuck on attaching a file):** "How to attach a file to an email in Gmail?"
            *   **Incorrect Query:** "Download my bank statement and email it to my accountant" *(This query is too broad, contains multiple sub-tasks, and does not start with "How to".)*
        """
        logger.info("=" * 50)
        logger.info("ACI: Calling Search Agent(query=%s)", query)
        logger.info("=" * 50)

        # Lazily create the searcher agent on first use
        if self.searcher_agent is None:
            self.searcher_agent = SearcherAgent.create(
                engine_params=self.engine_params_for_searcher,
                search_env=self.env,
                grounder_agent=self,
                platform=self.platform,
                client_password="osworld-public-evaluation",
            )

        if hasattr(self, "result_dir"):
            self.searcher_agent.result_dir = self.result_dir

        try:
            result = self.searcher_agent.search(query=query, obs=self.obs)
        except Exception as exc:
            logger.exception("Search agent failed")
            result = {
                "query": query,
                "completion_reason": "FAIL",
                "tutorial_notes": [],
                "execution_history": [],
                "steps_executed": 0,
                "budget": getattr(self.searcher_agent, "budget", 0),
                "final_answer": f"Search agent error: {exc}",
            }

        self.last_search_agent_result = result
        if result.get("completion_reason") == "DONE" and result.get("final_answer"):
            self.tutorials.append(result["final_answer"])
            logger.info(f"Search agent retrieved info: {result['final_answer']}")

        # Return sleep command since actual work happened out-of-band
        return "import time; time.sleep(2.222)"

    @agent_action
    def hotkey(self, keys: List):
        """Press a hotkey combination
        Args:
            keys:List the keys to press in combination in a list format (e.g. ['ctrl', 'c'])
        """
        # add quotes around the keys
        keys = [f"'{key}'" for key in keys]
        return f"import pyautogui; pyautogui.hotkey({', '.join(keys)})"

    @agent_action
    def hold_and_press(self, hold_keys: List, press_keys: List):
        """Hold a list of keys and press a list of keys
        Args:
            hold_keys:List, list of keys to hold
            press_keys:List, list of keys to press in a sequence
        """

        press_keys_str = "[" + ", ".join([f"'{key}'" for key in press_keys]) + "]"
        command = "import pyautogui; "
        for k in hold_keys:
            command += f"pyautogui.keyDown({repr(k)}); "
        command += f"pyautogui.press({press_keys_str}); "
        for k in hold_keys:
            command += f"pyautogui.keyUp({repr(k)}); "

        return command

    @agent_action
    def wait(self, time: float):
        """Wait for a specified amount of time
        Args:
            time:float the amount of time to wait in seconds
        """
        return f"""import time; time.sleep({time})"""

    @agent_action
    def done(
        self,
        return_value: Optional[Union[Dict, str, List, Tuple, int, float, bool]] = None,
    ):
        """End the current task with a success and the required return value"""
        self.returned_info = return_value
        return """DONE"""

    @agent_action
    def fail(self):
        """Report that the current task is infeasible and cannot be completed. Use this when the task is impossible due to missing features, absent hardware, missing resources, or violated constraints."""
        return """FAIL"""
