"""
Route entry node for the CodeAct agent.

This node is the entry point for the agent graph, handling:
- Agent loading from registry
- Sandbox initialization (tools and context)
- Agent script injection (for playbook agents)
- State recovery (for corrupted tool calls)
- Routing to appropriate next node

Note: System prompt creation is delegated to the prompts module
(see universal_mcp.agents.codeact0.prompts.build_complete_system_instructions)
"""

from typing import TYPE_CHECKING, Literal

from langchain_core.messages import AIMessage, RemoveMessage
from langchain_core.runnables import RunnableConfig
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langgraph.types import Command
from loguru import logger

from universal_mcp.agents.codeact0.prompts import build_complete_system_instructions
from universal_mcp.agents.codeact0.state import CodeActState
from universal_mcp.agents.codeact0.utils import detect_pending_tool

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


async def _load_agent_from_registry(
    agent_instance: "CodeActPlaybookAgent",
    state: CodeActState,
) -> None:
    """
    Load agent from agent_builder_registry if available and not already loaded.

    Updates agent_instance.agent and agent_instance.tools_config in place.
    """
    if not agent_instance.agent_builder_registry or agent_instance.agent:
        return

    # Set agent_id/template_id if provided in state
    if state.get("agent_id"):
        agent_instance.agent_builder_registry.default_agent_id = state["agent_id"]
    elif state.get("template_id"):
        agent_instance.agent_builder_registry.default_template_id = state["template_id"]

    # Load agent from registry
    agent_instance.agent = await agent_instance.agent_builder_registry.get_agent()
    agent_instance.tools_config = agent_instance.agent.tools if agent_instance.agent else {}

    logger.debug(
        f"Loaded agent from registry with {len(agent_instance.tools_config)} tools"
        if agent_instance.tools_config
        else "Loaded agent from registry (no tools configured)"
    )


async def _initialize_sandbox_context(
    agent_instance: "CodeActPlaybookAgent",
    state: CodeActState,
) -> str:
    """
    Initialize sandbox with all required context:
    1. Default tools (LLM tools, web search, etc.)
    2. Selected tools from state
    3. Previous execution context from state

    Returns:
        unconnected_links: Links to apps requiring authentication
    """
    # 1. Load default tools and update sandbox
    (
        agent_instance.default_tool_defs,
        default_tool_context,
        unconnected_links,
    ) = await agent_instance._load_default_tools()
    await agent_instance.sandbox.update_context(default_tool_context)
    logger.debug(f"Loaded {len(default_tool_context)} default tool context items")

    # 2. Load selected tools (user-requested tools from previous interactions)
    tool_defs, loaded_tools_context = await agent_instance._load_tools(state["selected_tool_ids"])
    await agent_instance.sandbox.update_context(loaded_tools_context)
    logger.debug(f"Loaded {len(state['selected_tool_ids'])} selected tools")

    # 3. Restore previous execution context if available
    if state.get("context"):
        await agent_instance.sandbox.update_context(state.get("context"))
        logger.debug("Restored previous execution context from state")

    return unconnected_links


async def _execute_agent_script_if_needed(
    agent_instance: "CodeActPlaybookAgent",
    state: CodeActState,
) -> str | None:
    """
    Execute agent script on first message to load helper functions into context.

    This is only done for playbook agents that have a script defined in their instructions.
    The script typically contains helper functions that the agent can use.

    Returns:
        Updated serialized context if script was executed, None otherwise
    """
    # Only execute on first message with a loaded agent
    if len(state["messages"]) != 1 or not agent_instance.agent:
        return None

    script = agent_instance.agent.instructions.get("script")
    if not script:
        return None

    logger.info("Executing agent script to load helper functions into context")

    # Execute the script
    result = await agent_instance.sandbox.run(script)
    if result["exit_code"] != 0:
        logger.error(f"Error executing agent script: {result.get('error_message', 'Unknown error')}")
        # Continue anyway - agent may still be partially functional

    # Get updated context to persist the defined functions
    updated_context = await agent_instance.sandbox.get_context()
    logger.debug("Captured updated context after agent script execution")

    return updated_context


def _should_recover_state(state: CodeActState) -> bool:
    """
    Check if state needs recovery due to corrupted/incomplete tool calls.

    Returns:
        True if state recovery is needed
    """
    return detect_pending_tool(state["messages"])


def _create_routing_command(
    *,
    next_node: Literal["call_model", "execute_tools"] | None = None,
    plan: list[str] | None = None,
    context: str | None = None,
    template_id: str | None = None,
    agent_id: str | None = None,
    messages: list | None = None,
) -> Command[Literal["call_model", "execute_tools"]]:
    """
    Create a routing command with state updates.

    Args:
        next_node: The next node to route to
        plan: Agent plan to update in state
        context: Execution context to update in state
        template_id: Template ID to update in state
        messages: Messages to update in state

    Returns:
        Command object with routing and state updates
    """
    updates = {}
    if plan is not None:
        updates["plan"] = plan
    if context is not None:
        updates["context"] = context
    if template_id is not None:
        updates["template_id"] = template_id
    if agent_id is not None:
        updates["agent_id"] = agent_id
    if messages is not None:
        updates["messages"] = messages

    return (
        Command(goto=next_node, update=updates)
        if next_node and updates
        else Command(goto=next_node)
        if next_node
        else Command(update=updates)
        if updates
        else Command()
    )


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

    The route_entry node is responsible for:
    1. Loading agent from registry (if configured)
    2. Initializing sandbox with tools and context
    3. Building system instructions (via prompts module)
    4. Executing agent script on first message (for playbook agents)
    5. Recovering from corrupted state
    6. Routing to the appropriate next node

    Args:
        agent_instance: The CodeActPlaybookAgent instance

    Returns:
        The route_entry node function
    """

    async def route_entry(
        state: CodeActState, config: RunnableConfig
    ) -> Command[Literal["call_model", "execute_tools"]]:
        """
        Entry point for the CodeAct agent graph.

        Prepares the agent and sandbox, then routes to either:
        - call_model: Normal execution flow
        - execute_tools: (Not used in current routing logic)
        """
        logger.debug(f"route_entry: Entry point reached with {len(state['messages'])} messages")
        # Step 1: Load agent from registry if needed
        await _load_agent_from_registry(agent_instance, state)

        # Step 2: Initialize sandbox context (tools + previous state)
        unconnected_links = await _initialize_sandbox_context(agent_instance, state)

        # Step 3: Build system instructions (using prompts module)
        agent_instance.final_instructions = build_complete_system_instructions(
            base_prompt=agent_instance.instructions,
            agent=agent_instance.agent,
            tool_definitions=agent_instance.default_tool_defs,
            agent_type="agent" if state.get("agent_id") else "template",
        )

        # Step 4: Execute agent script if this is the first message
        updated_context = await _execute_agent_script_if_needed(agent_instance, state)
        agent_instance.is_initialized = True

        # Step 5: Check if state recovery is needed
        if _should_recover_state(state):
            logger.warning("Detected corrupted/incomplete tool call - recovering state")
            return _create_routing_command(
                next_node="call_model",
                messages=[RemoveMessage(id=REMOVE_ALL_MESSAGES), *state["messages"]],
            )

        # Step 6: Route based on agent state
        # If agent script was executed (first message with loaded agent)
        if updated_context is not None:
            plan = agent_instance.agent.instructions.get("plan")
            template_id = config["configurable"].get("template_id")
            agent_id = config["configurable"].get("agent_id")

            # If there are unconnected apps, prompt user to authenticate
            if unconnected_links:
                logger.info("Agent has unconnected apps - prompting user to authenticate")
                ai_msg = (
                    f"Please login to the following app(s) using the following links "
                    f"and let me know in order to proceed with using the pre-loaded agent:\n"
                    f"{unconnected_links}"
                )
                return _create_routing_command(
                    plan=plan,
                    context=updated_context,
                    template_id=template_id,
                    agent_id=agent_id,
                    messages=[AIMessage(content=ai_msg)],
                )

            # Normal flow with agent script loaded
            logger.debug("Routing to call_model with agent script context loaded")
            return _create_routing_command(
                next_node="call_model",
                plan=plan,
                context=updated_context,
                template_id=template_id,
                agent_id=agent_id,
            )

        # Normal flow - no agent script execution needed
        logger.debug("Routing to call_model (normal flow)")
        return _create_routing_command(next_node="call_model")

    return route_entry
