from datetime import datetime
from typing import Any

from langchain_anthropic import ChatAnthropic
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.graph import START, StateGraph
from loguru import logger
from universal_mcp.tools.utils import tool_config_to_list
from universal_mcp.types import ToolFormat

from universal_mcp.agentr import AgentrRegistry
from universal_mcp.agents.base import BaseAgent
from universal_mcp.agents.codeact0.config import CodeActConfig, ContextSchema, get_config
from universal_mcp.agents.codeact0.llm_tool import smart_print
from universal_mcp.agents.codeact0.nodes import (
    create_call_model_node,
    create_execute_tools_node,
    create_route_entry_node,
    create_route_exit_node,
)
from universal_mcp.agents.codeact0.prompts import (
    build_tool_definitions,
)
from universal_mcp.agents.codeact0.sandbox import (
    Sandbox,
    execute_python_code,
)
from universal_mcp.agents.codeact0.state import (
    CodeActState,
)
from universal_mcp.agents.codeact0.tools import (
    create_agent_builder_tools,
    create_meta_tools,
    create_sandbox_tools,
    get_valid_tools,
    write_agent_creation_checklist,
)
from universal_mcp.agents.llm import load_chat_model

# LLM constants
AGENT_BUILDER_TEMPERATURE = 0.3
AGENT_BUILDER_THINKING = False
CACHE_TTL = "5m"


class CodeActPlaybookAgent(BaseAgent):
    def __init__(
        self,
        name: str,
        instructions: str,
        model: str,
        memory: BaseCheckpointSaver | None = None,
        registry: AgentrRegistry | None = None,
        agent_builder_registry: Any = None,
        sandbox_timeout: int | None = None,
        config: CodeActConfig | None = None,
        **kwargs,
    ):
        # Get config instance
        self.config = config or get_config()
        self.is_initialized = False

        # Use config for sandbox timeout if not explicitly provided
        if sandbox_timeout is None:
            sandbox_timeout = self.config.sandbox.timeout

        logger.info(f"Initializing CodeActPlaybookAgent: name={name}, model={model}, sandbox_timeout={sandbox_timeout}")
        super().__init__(
            name=name,
            instructions=instructions,
            model=model,
            memory=memory,
            **kwargs,
        )
        logger.debug(f"Loading primary and agent builder model: {model}")
        self.model_instance = load_chat_model(model)
        self.agent_builder_model_instance = load_chat_model(
            model,
            thinking=AGENT_BUILDER_THINKING,
            disable_streaming=True,
            tags=("quiet",),
            temperature=AGENT_BUILDER_TEMPERATURE,
        )
        self.registry = registry
        self.agent_builder_registry = agent_builder_registry
        self.agent = None

        self.tools_config = {}
        logger.debug(
            f"Tools config loaded: {len(self.tools_config)} tools" if self.tools_config else "No tools configured"
        )
        self.sandbox_timeout = sandbox_timeout
        self.final_instructions = ""
        logger.debug(f"Initializing sandbox with timeout={sandbox_timeout}s")
        self.sandbox = Sandbox(sandbox_timeout=self.sandbox_timeout)
        self.meta_tools = create_meta_tools(self.registry)
        self.sandbox_tools = create_sandbox_tools(self.registry)
        self.eval_mode = kwargs.get("eval_mode", False)
        self.run_start_time = datetime.now()
        logger.info(f"CodeActPlaybookAgent initialized successfully. eval_mode={self.eval_mode}")

    async def _load_tools(self, tools: list[str]):
        """Load tools into the registry and return the tool definitions and context
        Args:
            tools: List of tool names to load
        Returns:
            tool_defs: List of tool definitions
            tool_context: Dictionary of tool context
        """
        logger.debug(f"Loading {len(tools)} tools: {tools}")
        await self.registry.load_tools(tools)
        exported_tools: list = await self.registry.export_tools(tools, ToolFormat.NATIVE)  # List of callabales
        tool_defs, tool_context = build_tool_definitions(exported_tools)
        logger.info(f"Successfully loaded {len(tool_defs)} tool definitions")
        # Ideally we should add the context to sandbox directly, but we don't have a way to do that yet.
        return tool_defs, tool_context

    async def _load_default_tools(self):
        logger.info("Loading default tools")
        all_tools = []
        # Sandbox specific tools
        all_tools.append(smart_print)
        logger.debug("Added sandbox-specific tool: smart_print")
        # Load llm tools
        llm_tools_list = ["llm__generate_text", "llm__classify_data", "llm__extract_data", "llm__call_llm"]
        logger.debug(f"Loading {len(llm_tools_list)} LLM tools")
        exported_llm_tools: list = await self.registry.export_tools(llm_tools_list, ToolFormat.NATIVE)
        all_tools.extend(exported_llm_tools)
        # Load web search
        web_search = self.sandbox_tools.get("web_search")
        all_tools.append(web_search)
        logger.debug("Added web_search tool")
        # Load filesystem tools
        all_tools.append(self.sandbox_tools["fetch_as_markdown"])
        all_tools.append(self.sandbox_tools["upload_file"])
        logger.debug("Added sandbox tools: web_search, fetch_as_markdown, upload_file")
        unconnected_links = []
        # Playbook specific tools
        if self.tools_config:
            logger.debug(f"Loading {len(self.tools_config)} playbook-specific tools")
            tool_cofig = tool_config_to_list(self.tools_config)
            playbook_tools: list = await self.registry.export_tools(tool_cofig)
            all_tools.extend(playbook_tools)
            valid_tools, _, _, unconnected_links = await get_valid_tools(tool_cofig, self.registry)
            if unconnected_links:
                unconnected_links = "\n".join(unconnected_links)
        # Build and export
        tool_defs, tool_context = build_tool_definitions(all_tools)
        logger.info(f"Default tools loaded successfully: {len(all_tools)} total tools, {len(tool_defs)} definitions")
        return tool_defs, tool_context, unconnected_links

    def _bind_tools_to_model(self, model_instance):
        if not hasattr(self, "agent_facing_tools"):
            logger.debug("Creating agent-facing tools for first time")
            agent_builder_tools = create_agent_builder_tools()
            self.agent_facing_tools = [
                execute_python_code,
                agent_builder_tools["plan_agent"],
                agent_builder_tools["schedule_code_and_save_agent"],
                self.meta_tools["search_functions"],
                self.meta_tools["load_functions"],
                write_agent_creation_checklist,
            ]
            logger.debug(f"Created {len(self.agent_facing_tools)} agent-facing tools")

        if isinstance(model_instance, ChatAnthropic):
            logger.debug("Binding tools to ChatAnthropic model with cache_control")
            return model_instance.bind_tools(
                tools=self.agent_facing_tools,
                tool_choice="auto",
                cache_control={"type": "ephemeral", "ttl": CACHE_TTL},
            )
        else:
            logger.debug(f"Binding tools to {type(model_instance).__name__} model")
            return model_instance.bind_tools(
                tools=self.agent_facing_tools,
                tool_choice="auto",
            )

    async def _build_graph(self):  # noqa: PLR0915
        """Build the graph for the CodeAct Playbook Agent."""
        logger.info("Building CodeActPlaybookAgent graph")
        await self.registry.ainit()
        self.model_with_tools = self._bind_tools_to_model(self.model_instance)

        # Create node functions from factory functions
        call_model = create_call_model_node(self)
        execute_tools = create_execute_tools_node(self)
        route_entry = create_route_entry_node(self)
        route_exit = create_route_exit_node(self)

        agent = StateGraph(state_schema=CodeActState, context_schema=ContextSchema)
        agent.add_node(call_model)
        agent.add_node(execute_tools)
        agent.add_node(route_entry)
        agent.add_node(route_exit)
        agent.add_edge(START, "route_entry")
        logger.info("Graph built successfully with nodes: call_model, execute_tools, route_entry")
        return agent.compile(checkpointer=self.memory).with_config(
            {"recursion_limit": self.config.graph.recursion_limit}
        )

    def __enter__(self):
        """Support synchronous context manager protocol."""
        logger.debug(f"Entering context manager for agent: {self.name}")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Support synchronous context manager protocol."""
        logger.debug(f"Exiting context manager for agent: {self.name}")
        # No cleanup needed currently, but could close resources here
        return False  # Don't suppress exceptions

    async def __aenter__(self):
        """Support asynchronous context manager protocol."""
        logger.debug(f"Entering async context manager for agent: {self.name}")
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Support asynchronous context manager protocol."""
        logger.debug(f"Exiting async context manager for agent: {self.name}")
        # No cleanup needed currently, but could close resources here
        return False  # Don't suppress exceptions
