"""
Execute tools node for the CodeAct agent.

This node handles execution of tool calls including:
- execute_python_code: Running Python in sandbox
- load_functions: Loading tools into sandbox
- search_functions: Searching for available tools
- plan_agent: Creating agent plans
- code_and_save_agent: Generating and saving agent code
- write_agent_creation_checklist: Managing agent creation checklists
"""

import json
import re
from typing import TYPE_CHECKING, Literal

from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langgraph.types import Command
from loguru import logger

from universal_mcp.agents.codeact0.state import CodeActState
from universal_mcp.agents.codeact0.tools import (
    build_or_patch_code,
    create_or_update_plan,
    write_agent_creation_checklist,
)

if TYPE_CHECKING:
    from universal_mcp.agents.codeact0.agent import CodeActPlaybookAgent


def _handle_termination(
    state: CodeActState,
    state_updates: dict,
    ask_user: bool,
    ai_msg: str | None,
    human_msg: str | None,
) -> Command | None:
    remaining_steps = state["remaining_steps"]
    tool_messages = state_updates["messages"]

    if not ask_user and (remaining_steps <= 5):
        logger.warning("execute_tools: Less than 5 steps remaining, routing to exit")
        tool_messages[-1].content += (
            "Note that you had less than 5 steps (including AI message containing "
            "tool call and tool responses) remaining. So the run has been ended."
        )
        ask_user = True
        ai_msg = (
            "The run has been terminated due to running out of steps. You can ask "
            "the agent to run again if you want to continue, or try something else."
        )

    if ask_user:
        if ai_msg:
            tool_messages.append(AIMessage(content=ai_msg))
        logger.debug("execute_tools: Interrupting for user input via route_exit")
        return Command(goto="route_exit", update=state_updates)

    if human_msg:
        tool_messages.append(HumanMessage(content=human_msg))

    if remaining_steps <= 10:
        logger.warning("execute_tools: Less than 10 steps remaining, routing to call_model")
        tool_messages[-1].content += (
            "Note that you have less than 10 steps (including AI message containing "
            "tool call and tool responses) remaining. Either finish the remaining "
            "work in this step or inform the user why it is taking so long, and "
            "possible alternatives"
        )

    return None


def create_execute_tools_node(agent_instance: "CodeActPlaybookAgent"):
    """
    Factory function that creates the execute_tools node with access to agent instance.

    Args:
        agent_instance: The CodeActPlaybookAgent instance

    Returns:
        The execute_tools node function
    """

    async def execute_tools(state: CodeActState) -> Command[Literal["call_model", "route_entry"]]:
        """Execute tool calls"""
        if not agent_instance.is_initialized:
            return Command(goto="route_entry")
        last_message = state["messages"][-1]
        tool_calls = last_message.tool_calls if isinstance(last_message, AIMessage) else []
        logger.debug(f"execute_tools: Processing {len(tool_calls)} tool calls")

        tool_messages = []
        state_updates = {
            "selected_tool_ids": [],
        }

        ask_user = False
        ai_msg = None
        human_msg = None

        for tool_call in tool_calls:
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]
            logger.debug(f"execute_tools: Executing tool '{tool_name}'")
            try:
                if tool_name == "execute_python_code":
                    code = tool_args["snippet"]
                    display_output_to_user = tool_args.get("display_output_to_user", False)
                    logger.debug(f"execute_tools: Executing Python code in sandbox (length={len(code)} chars)")
                    # Context is already in sandbox.namespace (loaded in route_entry)
                    output, new_context = await agent_instance.sandbox.handle_execute_python_code(
                        code, display_output_to_user
                    )
                    state_updates["context"] = new_context
                    if display_output_to_user and not re.search(r"error", output, re.IGNORECASE):
                        ask_user = True
                        tool_result = "Code output added to response"
                        ai_msg = output
                    else:
                        tool_result = output
                    logger.info("execute_tools: Python code execution completed successfully")
                elif tool_name == "load_functions":
                    # The tool now does all the work of validation and formatting.
                    logger.debug(f"execute_tools: Loading functions: {tool_args}")
                    (
                        tool_result,
                        new_context_for_sandbox,
                        valid_tools,
                        unconnected_links,
                    ) = await agent_instance.meta_tools[  # noqa: E501
                        "load_functions"
                    ].ainvoke(tool_args)
                    # We still need to update the sandbox context for `execute_python_code`
                    state_updates["selected_tool_ids"].extend(valid_tools)
                    if state_updates["selected_tool_ids"]:
                        await agent_instance.sandbox.update_context(new_context_for_sandbox)
                        logger.info(f"execute_tools: Loaded {len(valid_tools)} functions into sandbox")
                    if unconnected_links:
                        ask_user = True
                        ai_msg = f"Please login to the following app(s) using the following links and let me know in order to proceed:\n {unconnected_links} "
                        logger.info(f"execute_tools: User authentication required for {len(unconnected_links)} app(s)")

                elif tool_name == "search_functions":
                    logger.debug(f"execute_tools: Searching functions with query: {tool_args}")
                    tool_result = await agent_instance.meta_tools["search_functions"].ainvoke(tool_args)
                    logger.info("execute_tools: Function search completed")

                elif tool_name == "plan_agent":
                    logger.debug("execute_tools: Creating or updating agent plan")
                    state_updates["plan"], tool_result = await create_or_update_plan(
                        self=agent_instance, state=state, plan=state_updates.get("plan", state.get("plan", None))
                    )
                    logger.info("execute_tools: Agent plan created/updated")
                    # TODO: This should be an interrupt
                    if agent_instance.eval_mode:
                        human_msg = "Yes, the plan looks great"
                        logger.debug("execute_tools: eval_mode enabled, auto-approving plan")
                    else:
                        ask_user = True
                        logger.debug("execute_tools: Requesting user approval for plan")

                elif tool_name == "schedule_code_and_save_agent":
                    logger.debug("execute_tools: Building or patching or scheduling agent")
                    cron_schedule = tool_args.get("cron")
                    update_code = tool_args.get("update_code", True)
                    update_general_info = tool_args.get("update_general_info", True)
                    generate_html = tool_args.get("generate_html", False)
                    (
                        tool_result,
                        state_updates["context"],
                        state_updates["agent_name"],
                        state_updates["agent_description"],
                    ) = await build_or_patch_code(
                        self=agent_instance,
                        state=state,
                        cron=cron_schedule,
                        update_code=update_code,
                        update_general_info=update_general_info,
                        generate_html=generate_html,
                    )
                    state_updates["agent_id"] = tool_result.get("id")
                    logger.info(f"execute_tools: Agent code saved/scheduled (name={state_updates.get('agent_name')})")
                    if agent_instance.eval_mode:
                        human_msg = "Now test the agent. Do not ask me any questions, assume parameters appropriately."
                        logger.debug("execute_tools: eval_mode enabled, auto-proceeding to testing")
                elif tool_name == "write_agent_creation_checklist":
                    logger.debug("execute_tools: Writing agent creation checklist")
                    tool_result = await write_agent_creation_checklist.ainvoke(tool_args)
                    logger.info("execute_tools: Agent creation checklist written")
                else:
                    raise Exception(
                        f"Unexpected tool call: {tool_call['name']}. "
                        "tool calls must be one of 'execute_python_code', 'load_functions', 'search_functions', 'plan_agent', or 'schedule_code_and_save_agent', 'write_agent_creation_checklist'. For using functions, call them in code using 'execute_python_code'."
                    )
            except Exception as e:
                logger.warning(f"execute_tools: Tool '{tool_name}' execution failed: {str(e)}")
                tool_result = str(e)

            tool_message = ToolMessage(
                content=json.dumps(tool_result),
                name=tool_call["name"],
                tool_call_id=tool_call["id"],
            )
            tool_messages.append(tool_message)
        state_updates["messages"] = tool_messages
        early_return_command = _handle_termination(
            state=state,
            state_updates=state_updates,
            ask_user=ask_user,
            ai_msg=ai_msg,
            human_msg=human_msg,
        )
        if early_return_command:
            return early_return_command

        logger.debug(f"execute_tools: Completed {len(tool_calls)} tool executions, routing back to call_model")
        return Command(
            goto="call_model",
            update=state_updates,
        )

    return execute_tools
