"""
Call model node for the CodeAct agent.

This node handles LLM invocation with tool binding, caching, and fallback logic.
"""

import copy
from typing import TYPE_CHECKING, Literal, cast

from langchain_anthropic import ChatAnthropic
from langchain_core.messages import AIMessage
from langgraph.types import Command
from loguru import logger

from universal_mcp.agents.codeact0.state import CodeActState
from universal_mcp.agents.codeact0.utils import (
    build_anthropic_cache_message,
    invoke_with_retries,
    remove_thinking,
)
from universal_mcp.agents.llm import load_chat_model

# LLM constants
FALLBACK_TEMPERATURE = 0.3
FALLBACK_THINKING = False
CACHE_TTL = "5m"

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


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

    Args:
        agent_instance: The CodeActPlaybookAgent instance

    Returns:
        The call_model node function
    """

    async def call_model(state: CodeActState) -> Command[Literal["execute_tools", "route_entry"]]:
        """This node now only ever binds the four meta-tools to the LLM."""
        logger.debug(f"call_model: Processing {len(state['messages'])} messages")
        if not agent_instance.is_initialized:
            return Command(goto="route_entry")

        messages = build_anthropic_cache_message(agent_instance.final_instructions) + state["messages"]

        # Force cache
        if isinstance(agent_instance.model_instance, ChatAnthropic):
            if isinstance(messages[-1].content, str):
                pass
            else:
                last = copy.deepcopy(messages[-1])
                last.content[-1]["cache_control"] = {"type": "ephemeral", "ttl": CACHE_TTL}
                messages[-1] = last
                logger.debug("call_model: Applied cache_control to last message")

        try:
            logger.debug(
                f"call_model: Invoking primary model with retries (max_retries={agent_instance.config.retry.max_retries})"
            )
            response = await invoke_with_retries(
                agent_instance.model_with_tools, messages, max_attempts=agent_instance.config.retry.max_retries
            )
            logger.info("call_model: Primary model invocation successful")
        except Exception as e:
            logger.warning(f"Primary model failed or returned empty response: {e}. Falling back to backup model.")
            if (
                not hasattr(agent_instance, "fallback_model_with_tools")
                or agent_instance.fallback_model_with_tools is None
            ):
                logger.info(f"call_model: Creating fallback model: {agent_instance.config.llm.fallback_model}")
                agent_instance.fallback_model_instance = load_chat_model(
                    agent_instance.config.llm.fallback_model,
                    thinking=FALLBACK_THINKING,
                    temperature=FALLBACK_TEMPERATURE,
                )
                agent_instance.fallback_model_with_tools = agent_instance._bind_tools_to_model(
                    agent_instance.fallback_model_instance
                )

            # Use cleaned-up messages (no thinking) when talking to the fallback model
            fallback_messages = remove_thinking(messages)
            logger.debug("call_model: Invoking fallback model")
            response = cast(
                AIMessage,
                await agent_instance.fallback_model_with_tools.with_retry().ainvoke(fallback_messages),
            )
            response = cast(AIMessage, remove_thinking(response))
            logger.info("call_model: Fallback model invocation successful")

        if response.tool_calls:
            logger.debug(
                f"call_model: Response contains {len(response.tool_calls)} tool calls, routing to execute_tools"
            )
            return Command(goto="execute_tools", update={"messages": [response]})
        else:
            logger.debug("call_model: Response has no tool calls, returning to user via route_exit")
            return Command(goto="route_exit", update={"messages": [response]})

    return call_model
