import logging
import textwrap
import platform
from functools import partial

from mm_agents.vlaa_gui.agents.grounding import ACI
from mm_agents.vlaa_gui.agents.gate_agent import GateAgent
from mm_agents.vlaa_gui.core.module import BaseModule
from mm_agents.vlaa_gui.core.knowledge import KnowledgeBase
from mm_agents.vlaa_gui.memory.procedural_memory import PROCEDURAL_MEMORY
from mm_agents.vlaa_gui.core.engine import OpenAIEmbeddingEngine
from mm_agents.vlaa_gui.utils.common_utils import (
    call_llm_safe_with_thinking,
    call_llm_formatted,
    create_pyautogui_code,
    parse_single_code_from_string,
    sanitize_code,
    extract_first_agent_function,
)

from mm_agents.vlaa_gui.utils.formatters import (
    SINGLE_ACTION_FORMATTER,
    CODE_VALID_FORMATTER,
)

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


class Worker(BaseModule):
    def __init__(
        self,
        engine_params: dict,
        grounding_agent: ACI,
        reflection_engine_params: dict | None = None,
        local_kb_path: str = "agent_memory",
        embedding_engine=OpenAIEmbeddingEngine(),
        platform: str = platform.system().lower(),
        enable_reflection: bool = True,
        observation_type: str = "screenshot",
        planning_type: str = "iterative",
        search_engine: str | None = None,
        memory_type: str = "null",
        use_task_experience: bool = False,
        coding_agent_flag: bool = False,
        max_image_in_trajectory: int = 8,
        debug: bool = False,
        lexical_weight: float = 0.0,
        enable_gate: bool = True,
        loop_detection: bool = True,
        feasibility_check: bool = True,
    ):
        """
        Worker receives a subtask list and active subtask and generates the next action for the to execute.
        Args:
            engine_params: dict
                Parameters for the multimodal engine
            grounding_agent: Agent
                The grounding agent to use
            local_kb_path: str
                Path to knowledge base
            platform: str
                OS platform the agent runs on (darwin, linux, windows)
            enable_reflection: bool
                Whether to enable reflection
            use_task_experience: bool
                Whether to use task experience
        """
        super().__init__(engine_params, platform)

        self.temperature = self.engine_params.get("temperature")
        if self.temperature is None:
            self.temperature = 0.0
        self.max_image_in_trajectory = max_image_in_trajectory
        self.reflection_engine_params = reflection_engine_params or self.engine_params

        self.grounding_agent = grounding_agent
        self.local_kb_path = local_kb_path
        self.embedding_engine = embedding_engine
        self.enable_reflection = enable_reflection
        self.use_task_experience = use_task_experience
        self.observation_type = observation_type
        self.planning_type = planning_type
        self.memory_type = memory_type
        self.search_engine = search_engine
        self.lexical_weight = lexical_weight
        self.coding_agent_flag = coding_agent_flag
        self.debug = debug
        self.enable_gate = enable_gate
        self.loop_detection = loop_detection
        self.feasibility_check = feasibility_check
        self.procedural_memory = PROCEDURAL_MEMORY
        self.gate_agent = GateAgent(engine_params, platform) if enable_gate else None
        self.reset()

    def reset(self):
        if self.platform != "linux":
            skipped_actions = ["set_cell_values"]
        else:
            skipped_actions = []

        # Hide code agent action when coding is disabled or env/controller is unavailable
        if (
            not self.coding_agent_flag
            or not getattr(self.grounding_agent, "env", None)
            or not getattr(
                getattr(self.grounding_agent, "env", None), "controller", None
            )
        ):
            skipped_actions.append("call_code_agent")

        # Only expose call_search_agent when search_engine == 'search_agent'
        if self.search_engine != "search_agent":
            skipped_actions.append("call_search_agent")

        # Only expose save_to_knowledge when memory_type != 'null'
        if self.memory_type == "null":
            skipped_actions.append("save_to_knowledge")

        pm = self.procedural_memory
        sys_prompt = pm.construct_worker_procedural_memory(
            type(self.grounding_agent),
            skipped_actions=skipped_actions,
            observation_type=self.observation_type,
            coding_agent=self.coding_agent_flag,
            search_agent=self.search_engine == "search_agent",
            planning_type=self.planning_type,
            enable_gate=self.enable_gate,
            loop_detection=self.loop_detection,
            feasibility_check=self.feasibility_check,
        ).replace("CURRENT_OS", self.platform)
        if self.debug:
            # logger.debug("Worker System Prompt: %s", sys_prompt)
            print("WORKER SYSTEM PROMPT:\n", sys_prompt)

        self.generator_agent = self._create_agent(sys_prompt)
        if hasattr(pm, "get_reflection_prompt"):
            reflection_prompt = pm.get_reflection_prompt(
                self.loop_detection, self.feasibility_check
            )
        else:
            reflection_prompt = pm.REFLECTION_ON_TRAJECTORY
        self.reflection_agent = self._create_agent(
            reflection_prompt, self.reflection_engine_params
        )
        self.knowledge_base = KnowledgeBase(
            embedding_engine=self.embedding_engine,
            local_kb_path=self.local_kb_path,
            platform=self.platform,
            engine_params=self.engine_params,
            lexical_weight=self.lexical_weight,
        )

        self.turn_count = 0
        self.worker_history = []
        self.reflections = []
        self.screenshot_inputs = []
        self.rag_flag = False
        self.pending_verifier_feedback: dict | None = None
        self.gate_criteria_str = ""  # populated by gate agent at step 1

    def set_replan_feedback(self, verifier_feedback: dict | None) -> None:
        """Attach one-shot verifier feedback for the next fresh plan."""
        self.pending_verifier_feedback = verifier_feedback

    def flush_messages(self):
        max_trajector_length = self.max_image_in_trajectory
        # generator msgs are alternating [user, assistant], so 2 per round
        if len(self.generator_agent.messages) > 2 * max_trajector_length + 1:
            self.generator_agent.remove_message_at(1)
            self.generator_agent.remove_message_at(1)
        # reflector msgs are all [(user text, user image)], so 1 per round
        if len(self.reflection_agent.messages) > max_trajector_length + 1:
            self.reflection_agent.remove_message_at(1)

    @staticmethod
    def _get_message_blocks(message):
        if isinstance(message, dict):
            return message.get("content")
        return getattr(message, "parts", None)

    @staticmethod
    def _is_image_block(block) -> bool:
        if isinstance(block, dict):
            return "image" in str(block.get("type", ""))

        inline_data = getattr(block, "inline_data", None)
        if inline_data is not None:
            mime_type = getattr(inline_data, "mime_type", "") or ""
            if mime_type.startswith("image/"):
                return True

        file_data = getattr(block, "file_data", None)
        if file_data is not None:
            mime_type = getattr(file_data, "mime_type", "") or ""
            if mime_type.startswith("image/"):
                return True

        return False

    def _flush_agent_images(self, agent, max_images: int) -> None:
        if agent is None:
            return

        img_count = 0
        for i in range(len(agent.messages) - 1, -1, -1):
            blocks = self._get_message_blocks(agent.messages[i])
            if not blocks:
                continue
            for j in range(len(blocks) - 1, -1, -1):
                if self._is_image_block(blocks[j]):
                    img_count += 1
                    if img_count > max_images:
                        del blocks[j]

    def _format_code_agent_result_for_reflection(self, result: dict) -> str:
        """Format a concise code-agent execution note for reflection."""
        if not isinstance(result, dict):
            return ""

        def _trim(text: str, limit: int = 600) -> str:
            if len(text) <= limit:
                return text
            return text[: limit - 1] + "…"

        completion_reason = str(result.get("completion_reason", "UNKNOWN"))
        summary = _trim(str(result.get("summary", "")).strip(), limit=700)
        steps_executed = result.get("steps_executed", 0)
        budget = result.get("budget", 0)
        error = str(result.get("error", "")).strip()

        lines = [
            "EXECUTION NOTE (code agent):",
            f"- completion_reason: {completion_reason}",
            f"- steps_executed: {steps_executed}/{budget}",
        ]
        if summary:
            lines.append(f"- summary: {summary}")
        if error:
            lines.append(f"- error: {_trim(error, limit=400)}")
        return "\n".join(lines)

    def _should_save_memory(self, reflection: str) -> bool:
        return False

    @staticmethod
    def _format_verifier_feedback(verifier_feedback: dict | None) -> str:
        """Format verifier rejection details for the next replanning turn."""
        if not isinstance(verifier_feedback, dict):
            return ""

        reason = str(verifier_feedback.get("reason", "")).strip()
        missing_steps = str(verifier_feedback.get("missing_steps", "")).strip()
        if not reason and not missing_steps:
            return ""

        lines = [
            "VERIFIER REPLAN FEEDBACK:",
            "The previous attempt called agent.done(), but completion was rejected.",
            "Continue the same original task and address this feedback before calling agent.done() again.",
        ]
        if reason:
            lines.append(f"- Reason: {reason}")
        if missing_steps:
            lines.append(f"- Missing steps/evidence: {missing_steps}")
        return "\n".join(lines)

    def generate_reflection(
        self, instruction: str, obs: dict, verifier_feedback_note: str = ""
    ) -> tuple[str, str]:
        """
        Generate a reflection based on the current observation and instruction.

        Args:
            instruction (str): The task instruction.
            obs (dict): The current observation containing the screenshot.

        Returns:
            tuple[str, str]: The generated reflection text and thoughts, if any (turn_count > 0).

        Side Effects:
            - Updates reflection agent's history
            - Generates reflection response with API call
        """
        reflection = None
        reflection_thoughts = None
        if self.enable_reflection:
            # Load the initial message
            if self.turn_count == 0:
                text_content = textwrap.dedent(
                    f"""
                    Task Description: {instruction}
                    Current Trajectory below:
                    """
                )
                if verifier_feedback_note:
                    text_content += (
                        "\nPrevious completion attempt feedback:\n"
                        f"{verifier_feedback_note}\n"
                    )
                updated_sys_prompt = (
                    self.reflection_agent.system_prompt + "\n" + text_content
                )
                self.reflection_agent.add_system_prompt(updated_sys_prompt)
                self.reflection_agent.add_message(
                    text_content="The initial screen is provided. No action has been taken yet.",
                    image_content=obs["screenshot"],
                    role="user",
                )
            # Load the latest action
            else:
                reflection_input = self.worker_history[-1]
                code_result_note = self._format_code_agent_result_for_reflection(
                    getattr(self.grounding_agent, "last_code_agent_result", None)
                )
                if code_result_note:
                    reflection_input += "\n\n" + code_result_note
                self.reflection_agent.add_message(
                    text_content=reflection_input,
                    image_content=obs["screenshot"],
                    role="user",
                )
                reflection, reflection_thoughts = call_llm_safe_with_thinking(
                    self.reflection_agent,
                    temperature=self.temperature,
                    agent_type="reflection",
                )
                self.reflections.append(reflection)
                logger.info("REFLECTION THOUGHTS: %s", reflection_thoughts)
                logger.info("REFLECTION: %s", reflection)
        return reflection, reflection_thoughts

    def generate_next_action(
        self, obs: dict, instruction: str
    ) -> tuple[dict, list[str]]:
        """Generate next action iteratively

        Args:
            obs: Current UI state observation
            instruction: Natural language instruction
        Returns:
            Tuple containing executor info dictionary and list of action codes
        """
        assert self.planning_type == "iterative", (
            "Planner called with incorrect planning type."
        )

        self.grounding_agent.assign_screenshot(obs)
        self.grounding_agent.set_task_instruction(instruction)
        self.rag_flag = self.turn_count == 0
        verifier_feedback_note = ""
        if self.turn_count == 0:
            verifier_feedback_note = self._format_verifier_feedback(
                self.pending_verifier_feedback
            )

        def _has_memory(experience: str | None) -> bool:
            return bool(experience) and experience != "None"

        def _retrieve_iterative_experience(
            query: str, retrieval_stage: str
        ) -> tuple[str, str]:
            memory_keys: list[str] = []
            memory_chunks: list[str] = []

            # narrative: task-level memory
            if self.memory_type == "mixed":
                narrative_key, narrative_exp = (
                    self.knowledge_base.retrieve_narrative_experience(query)
                )
                if _has_memory(narrative_exp):
                    memory_keys.append(f"narrative:{narrative_key}")
                    memory_chunks.append(
                        f"[Narrative task experience]\n{narrative_exp}"
                    )
                    logger.debug(
                        "%s NARRATIVE MEMORY RETRIEVAL: %s\n%s",
                        retrieval_stage,
                        narrative_key,
                        narrative_exp.strip(),
                    )

            # episodic: episode/subtask-level memory
            if self.memory_type in ("mixed", "episodic"):
                episodic_key, episodic_exp = (
                    self.knowledge_base.retrieve_episodic_experience(query)
                )
                if _has_memory(episodic_exp):
                    memory_keys.append(f"episodic:{episodic_key}")
                    memory_chunks.append(f"[Episodic experience]\n{episodic_exp}")
                    logger.debug(
                        "%s EPISODIC MEMORY RETRIEVAL: %s\n%s",
                        retrieval_stage,
                        episodic_key,
                        episodic_exp.strip(),
                    )

            combined_key = " | ".join(memory_keys)
            combined_experience = "\n\n".join(memory_chunks)
            return combined_key, combined_experience

        if self.rag_flag:
            most_similar_memory = None
            retrieved_experience = None
            integrated_knowledge = None
            if self.memory_type in ("mixed", "episodic"):
                most_similar_memory, retrieved_experience = (
                    _retrieve_iterative_experience(
                        instruction, retrieval_stage="INITIAL"
                    )
                )
            # Retrieve knowledge from the web if search_engine is provided
            if self.search_engine is not None and self.search_engine != "search_agent":
                self.search_query = self.knowledge_base.formulate_query(
                    instruction, obs
                )
                retrieved_knowledge = self.knowledge_base.retrieve_knowledge(
                    instruction=instruction,
                    search_query=self.search_query,
                    search_engine=self.search_engine,
                )
                logger.info("RETRIEVED KNOWLEDGE FROM WEB: %s", retrieved_knowledge)

                if retrieved_knowledge is not None:
                    # Fuse the retrieved knowledge and experience
                    if self.memory_type in ("mixed", "episodic") and _has_memory(
                        retrieved_experience
                    ):
                        integrated_knowledge = self.knowledge_base.knowledge_fusion(
                            observation=obs,
                            instruction=instruction,
                            web_knowledge=retrieved_knowledge,
                            similar_task=most_similar_memory,
                            experience=retrieved_experience,
                        )
                    else:
                        integrated_knowledge = retrieved_knowledge
                    logger.info("INTEGRATED KNOWLEDGE: %s", integrated_knowledge)

            if _has_memory(retrieved_experience):
                integrated_knowledge = integrated_knowledge or retrieved_experience

            # Add the integrated knowledge to the task instruction in the system prompt
            if _has_memory(integrated_knowledge):
                instruction += f"\nYou may refer to some retrieved knowledge if you think they are useful.{integrated_knowledge}"

            # Load the task into the system prompt
            self.generator_agent.add_system_prompt(
                self.generator_agent.system_prompt.replace(
                    "TASK_DESCRIPTION", instruction
                )
            )

            # Inject tutorials from search agent if available
            if (
                hasattr(self.grounding_agent, "tutorials")
                and self.grounding_agent.tutorials
            ):
                tutorials_str = "\n\n### Tutorials Found by Search Agent:\n"
                for idx, tutorial in enumerate(self.grounding_agent.tutorials, start=1):
                    tutorials_str += f"#### Tutorial {idx}:\n{tutorial}\n\n"
                self.generator_agent.add_system_prompt(
                    self.generator_agent.system_prompt + tutorials_str
                )

        # Predict the next action(s) based on the current observation.
        generator_message = (
            ""
            if self.turn_count > 0
            else "The initial screen is provided. No action has been taken yet."
        )

        # Inject recon context on the first turn so the agent has environment
        # awareness gathered before execution started.
        if self.turn_count == 0 and obs.get("recon_context"):
            generator_message += (
                "\n\nENVIRONMENT RECONNAISSANCE (read-only pre-task inspection):\n"
                + obs["recon_context"]
                + "\n"
            )

        # Propose completion gate criteria at step 1
        if self.turn_count == 0 and self.gate_agent is not None:
            criteria = self.gate_agent.propose_criteria(instruction, obs)
            self.gate_criteria_str = GateAgent.format_criteria_for_worker(criteria)
            logger.info("Gate criteria proposed:\n%s", self.gate_criteria_str)

        # Inject gate criteria into every step's message
        if self.gate_criteria_str:
            generator_message += "\n" + self.gate_criteria_str + "\n"

        # Retrieve memory before generating action in middle steps (when memory_type is not null)
        if self.turn_count > 0 and self.memory_type != "null":
            # Build a contextual query based on current state and recent history
            recent_history = self.worker_history[-3:] if self.worker_history else []
            context_query = (
                f"Task: {instruction}\nRecent actions: {' '.join(recent_history)}"
            )

            retrieved_key, retrieved_experience = _retrieve_iterative_experience(
                context_query, retrieval_stage="MID-STEP"
            )
            if _has_memory(retrieved_experience):
                logger.debug(
                    "MID-STEP MEMORY RETRIEVAL (combined): %s\n%s",
                    retrieved_key,
                    retrieved_experience.strip(),
                )
                generator_message += (
                    f"\nRelevant experience from memory: {retrieved_experience}\n"
                )

        # Get the per-step reflection
        reflection, reflection_thoughts = self.generate_reflection(
            instruction, obs, verifier_feedback_note=verifier_feedback_note
        )
        if reflection:
            generator_message += f"REFLECTION: You may use this reflection on the previous action and overall trajectory:\n{reflection}\n"

            # Save memory after reflection when memory_type is not null
            # Check if reflection indicates a significant learning or milestone
            if self.memory_type != "null" and self.turn_count > 0:
                # Build trajectory from recent history for episodic memory
                if self.worker_history:
                    recent_trajectory = "\n".join(self.worker_history[-5:])
                    trajectory_key = f"Task: {instruction}\nStep: {self.turn_count}\nReflection: {reflection}"

                    # Check if reflection suggests saving (e.g., completed a milestone or learned something useful)
                    should_save = self._should_save_memory(reflection)
                    if should_save:
                        memory_trajectory = f"Trajectory:\n{recent_trajectory}\nReflection: {reflection}"
                        if self.memory_type in ("mixed", "episodic"):
                            self.knowledge_base.save_episodic_memory(
                                trajectory_key, memory_trajectory
                            )
                            logger.info(
                                "MID-STEP MEMORY SAVE: Saved episodic experience after reflection"
                            )

                        if self.memory_type == "mixed":
                            self.knowledge_base.save_narrative_memory(
                                trajectory_key, memory_trajectory
                            )
                            logger.info(
                                "MID-STEP MEMORY SAVE: Saved narrative experience after reflection"
                            )

        # Add finalized message to conversation
        generator_message += (
            f"\nCurrent Text Buffer = [{','.join(self.grounding_agent.notes)}]\n"
        )
        if verifier_feedback_note:
            generator_message += "\n" + verifier_feedback_note + "\n"

        # Add code agent result from previous step if available (from full task or subtask execution)
        if (
            hasattr(self.grounding_agent, "last_code_agent_result")
            and self.grounding_agent.last_code_agent_result is not None
        ):
            if not self.coding_agent_flag:
                raise ValueError(
                    "Code agent result found but coding agent is disabled."
                )
            code_result = self.grounding_agent.last_code_agent_result
            generator_message += "\nCODE AGENT RESULT:\n"
            generator_message += (
                f"Task/Subtask Instruction: {code_result['task_instruction']}\n"
            )
            generator_message += f"Steps Completed: {code_result['steps_executed']}\n"
            generator_message += f"Max Steps: {code_result['budget']}\n"
            generator_message += (
                f"Completion Reason: {code_result['completion_reason']}\n"
            )
            generator_message += f"Summary: {code_result['summary']}\n"
            if code_result["execution_history"]:
                generator_message += "Execution History:\n"
                for i, step in enumerate(code_result["execution_history"]):
                    action = step["action"]
                    # Format code snippets with proper backticks
                    if "```python" in action:
                        # Extract Python code and format it
                        code_start = action.find("```python") + 9
                        code_end = action.find("```", code_start)
                        if code_end != -1:
                            python_code = action[code_start:code_end].strip()
                            generator_message += (
                                f"Step {i + 1}: \n```python\n{python_code}\n```\n"
                            )
                        else:
                            generator_message += f"Step {i + 1}: \n{action}\n"
                    elif "```bash" in action:
                        # Extract Bash code and format it
                        code_start = action.find("```bash") + 7
                        code_end = action.find("```", code_start)
                        if code_end != -1:
                            bash_code = action[code_start:code_end].strip()
                            generator_message += (
                                f"Step {i + 1}: \n```bash\n{bash_code}\n```\n"
                            )
                        else:
                            generator_message += f"Step {i + 1}: \n{action}\n"
                    else:
                        generator_message += f"Step {i + 1}: \n{action}\n"
            generator_message += "\n"

            # Log the code agent result section for debugging (truncated execution history)
            log_message = "\nCODE AGENT RESULT:\n"
            log_message += (
                f"Task/Subtask Instruction: {code_result['task_instruction']}\n"
            )
            log_message += f"Steps Completed: {code_result['steps_executed']}\n"
            log_message += f"Max Steps: {code_result['budget']}\n"
            log_message += f"Completion Reason: {code_result['completion_reason']}\n"
            log_message += f"Summary: {code_result['summary']}\n"
            if code_result["execution_history"]:
                log_message += "Execution History (truncated):\n"
                # Only log first 3 steps and last 2 steps to keep logs manageable
                total_steps = len(code_result["execution_history"])
                for i, step in enumerate(code_result["execution_history"]):
                    if i < 3 or i >= total_steps - 2:  # First 3 and last 2 steps
                        action = step["action"]
                        if "```python" in action:
                            code_start = action.find("```python") + 9
                            code_end = action.find("```", code_start)
                            if code_end != -1:
                                python_code = action[code_start:code_end].strip()
                                log_message += (
                                    f"Step {i + 1}: ```python\n{python_code}\n```\n"
                                )
                            else:
                                log_message += f"Step {i + 1}: {action}\n"
                        elif "```bash" in action:
                            code_start = action.find("```bash") + 7
                            code_end = action.find("```", code_start)
                            if code_end != -1:
                                bash_code = action[code_start:code_end].strip()
                                log_message += (
                                    f"Step {i + 1}: ```bash\n{bash_code}\n```\n"
                                )
                            else:
                                log_message += f"Step {i + 1}: {action}\n"
                        else:
                            log_message += f"Step {i + 1}: {action}\n"
                    elif i == 3 and total_steps > 5:
                        log_message += f"... (truncated {total_steps - 5} steps) ...\n"

            # logger.info(
            #     f"WORKER_CODE_AGENT_RESULT_SECTION - Step {self.turn_count + 1}: Code agent result added to generator message:\n{log_message}"
            # )

            # Reset the code agent result after adding it to context
            self.grounding_agent.last_code_agent_result = None

        # Add search agent result from previous step if available
        if (
            hasattr(self.grounding_agent, "last_search_agent_result")
            and self.grounding_agent.last_search_agent_result is not None
        ):
            search_result = self.grounding_agent.last_search_agent_result
            generator_message += "\nSEARCH AGENT RESULT:\n"
            generator_message += f"Search Query: {search_result.get('query', '')}\n"
            generator_message += f"Search Completion Reason: {search_result.get('completion_reason', '')}\n"
            generator_message += "Search Result: "
            if search_result.get("completion_reason") == "DONE":
                generator_message += "Search completed successfully. The tutorial has been added to the system prompt.\n"
            elif search_result.get("completion_reason") == "FAIL":
                generator_message += f"Search failed. Hint: {search_result.get('final_answer', 'No details available')}\n"
            else:
                generator_message += f"Search ended: {search_result.get('final_answer', 'No details available')}\n"
            generator_message += "\n"

            logger.info(
                f"WORKER_SEARCH_AGENT_RESULT_SECTION - Step {self.turn_count + 1}: Search agent result added"
            )

            # Reset the search agent result after adding it to context
            self.grounding_agent.last_search_agent_result = None

        # On the final step, force the agent to make a terminal decision
        if obs.get("is_last_step"):
            generator_message += (
                "\n⚠️ FINAL STEP — This is your LAST action before the step budget is exhausted. "
                "You MUST choose one of:\n"
                "  - agent.done() if the task has been completed\n"
                "  - agent.fail() if the task cannot be completed or is infeasible\n"
                "Do NOT choose any other action. Any non-terminal action will be wasted.\n"
            )

        self.generator_agent.add_message(
            generator_message, image_content=obs["screenshot"], role="user"
        )
        # Generate the plan and next action
        format_checkers = [
            SINGLE_ACTION_FORMATTER,
            partial(CODE_VALID_FORMATTER, self.grounding_agent, obs),
        ]

        plan = call_llm_formatted(
            self.generator_agent,
            format_checkers,
            agent_type="executor",
            temperature=self.temperature,
        )

        # plan, plan_thoughts = split_thinking_response(plan)
        # NOTE: currently dropping thinking tokens from context
        self.worker_history.append(plan)
        logger.info("FULL PLAN:\n %s", plan)
        self.generator_agent.add_message(plan, role="assistant")

        plan_code = parse_single_code_from_string(plan.split("Grounded Action")[-1])
        plan_code = sanitize_code(plan_code)
        plan_code = extract_first_agent_function(plan_code)
        try:
            assert plan_code, "Plan code should not be empty"
            exec_code = create_pyautogui_code(self.grounding_agent, plan_code, obs)
        except Exception as e:
            logger.error("Error in parsing plan code: %s", e)
            exec_code = self.grounding_agent.wait(
                1.333
            )  # Skip a turn if the code cannot be evaluated

        executor_info = {
            "executor_plan": plan,
            "plan_code": plan_code,
            "exec_code": exec_code,
            "reflection": reflection,
            "reflection_thoughts": reflection_thoughts,
            "code_agent_output": (
                self.grounding_agent.last_code_agent_result
                if hasattr(self.grounding_agent, "last_code_agent_result")
                and self.grounding_agent.last_code_agent_result is not None
                else None
            ),
        }
        if verifier_feedback_note:
            self.pending_verifier_feedback = None
        self.turn_count += 1

        self.screenshot_inputs.append(obs["screenshot"])
        self.flush_messages()

        return executor_info, [exec_code]
