"""Tool management for LLMling agents."""

from __future__ import annotations

import asyncio
from collections.abc import Callable, Sequence
from contextlib import AsyncExitStack, contextmanager
from dataclasses import dataclass, field, fields
from datetime import datetime
import os
from typing import TYPE_CHECKING, Any, Literal, Self

from llmling import BaseRegistry, LLMCallableTool, ToolError
from psygnal import Signal

from llmling_agent.log import get_logger
from llmling_agent.mcp_server.client import MCPClient
from llmling_agent.models.mcp_server import MCPServerConfig, SSEMCPServer, StdioMCPServer
from llmling_agent.tools.base import ToolInfo


if TYPE_CHECKING:
    from collections.abc import Iterator
    from types import TracebackType

    from llmling_agent.agent import AnyAgent
    from llmling_agent.common_types import AnyCallable, ToolSource, ToolType
    from llmling_agent.models.context import AgentContext


logger = get_logger(__name__)

MAX_LEN_DESCRIPTION = 1000


class ToolManager(BaseRegistry[str, ToolInfo]):
    """Manages tool registration, enabling/disabling and access.

    Inherits from BaseRegistry providing:
    - Dict-like access: manager["tool_name"] -> ToolInfo
    - Async startup/shutdown: await manager.startup()
    - Event observation: manager.add_observer(observer)
    - Registration: manager.register("tool_name", tool)
    - Listing: manager.list_items()
    - State check: manager.is_empty, manager.has_item()
    - Async iteration: async for name, tool in manager: ...
    """

    @dataclass(frozen=True)
    class ToolStateReset:
        """Emitted when tool states are reset."""

        previous_tools: dict[str, bool]
        new_tools: dict[str, bool]
        timestamp: datetime = field(default_factory=datetime.now)

    tool_states_reset = Signal(ToolStateReset)

    def __init__(
        self,
        tools: Sequence[ToolInfo | ToolType | dict[str, Any]] | None = None,
        *,
        tool_choice: bool | str | list[str] = True,
        context: AgentContext[Any] | None = None,
    ):
        """Initialize tool manager.

        Args:
            tools: Initial tools to register
            tool_choice: Control tool usage:
                - True: Allow all tools
                - False: No tools
                - str: Use specific tool
                - list[str]: Allow specific tools
            context: Tool context
        """
        super().__init__()
        self.tool_choice = tool_choice
        self.context = context
        self._mcp_clients: dict[str, MCPClient] = {}
        self.exit_stack = AsyncExitStack()

        # Register initial tools
        for tool in tools or []:
            t = self._validate_item(tool)
            self.register(t.name, t)

    def __prompt__(self) -> str:
        enabled_tools = [t.name for t in self.values() if t.enabled]
        if not enabled_tools:
            return "No tools available"
        return f"Available tools: {', '.join(enabled_tools)}"

    async def __aenter__(self) -> Self:
        """Enter async context and set up MCP servers."""
        try:
            # Setup MCP servers if configured
            if self.context and self.context.config and self.context.config.mcp_servers:
                await self.setup_mcp_servers(self.context.config.get_mcp_servers())
        except Exception as e:
            # Clean up on error
            await self.__aexit__(type(e), e, e.__traceback__)
            msg = "Failed to initialize tool manager"
            logger.exception(msg, exc_info=e)
            raise RuntimeError(msg) from e
        else:
            return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ):
        """Exit async context."""
        try:
            # Clean up MCP clients through exit stack
            await self.exit_stack.aclose()
            self._mcp_clients.clear()
        except Exception as e:
            msg = "Error during tool manager cleanup"
            logger.exception(msg, exc_info=e)
            raise RuntimeError(msg) from e

    def reset_states(self):
        """Reset all tools to their default enabled states."""
        for info in self.values():
            info.enabled = True

    @property
    def _error_class(self) -> type[ToolError]:
        """Error class for tool operations."""
        return ToolError

    def _validate_item(self, item: ToolInfo | ToolType | dict[str, Any]) -> ToolInfo:
        """Validate and convert items before registration."""
        match item:
            case ToolInfo():
                return item
            case LLMCallableTool():
                return ToolInfo(callable=item)
            case str():
                item = LLMCallableTool.from_callable(item)
                return ToolInfo(callable=item)
            case _ if callable(item):
                tool = LLMCallableTool.from_callable(item)
                return ToolInfo(callable=tool)
            case {"callable": callable_item, **config} if callable(
                callable_item
            ) or isinstance(callable_item, LLMCallableTool):
                # First convert callable to LLMCallableTool if needed
                tool = (
                    callable_item
                    if isinstance(callable_item, LLMCallableTool)
                    else LLMCallableTool.from_callable(callable_item)
                )

                # Get valid fields from ToolInfo dataclass (excluding 'callable')
                valid_keys = {f.name for f in fields(ToolInfo)} - {"callable"}
                tool_config = {k: v for k, v in config.items() if k in valid_keys}

                return ToolInfo(callable=tool, **tool_config)  # type: ignore

            case _:
                typ = type(item)
                msg = f"Item must be ToolInfo, LLMCallableTool, or callable. Got {typ}"
                raise ToolError(msg)

    def enable_tool(self, tool_name: str):
        """Enable a previously disabled tool."""
        if tool_name not in self:
            msg = f"Tool not found: {tool_name}"
            raise ToolError(msg)
        tool_info = self[tool_name]
        tool_info.enabled = True
        self.events.changed(tool_name, tool_info)
        logger.debug("Enabled tool: %s", tool_name)

    def disable_tool(self, tool_name: str):
        """Disable a tool."""
        if tool_name not in self:
            msg = f"Tool not found: {tool_name}"
            raise ToolError(msg)
        tool_info = self[tool_name]
        tool_info.enabled = False
        self.events.changed(tool_name, tool_info)
        logger.debug("Disabled tool: %s", tool_name)

    def is_tool_enabled(self, tool_name: str) -> bool:
        """Check if a tool is currently enabled."""
        return self[tool_name].enabled if tool_name in self else False

    def list_tools(self) -> dict[str, bool]:
        """Get a mapping of all tools and their enabled status."""
        return {name: info.enabled for name, info in self.items()}

    def get_tools(
        self,
        state: Literal["all", "enabled", "disabled"] = "all",
    ) -> list[LLMCallableTool]:
        """Get tool objects based on filters."""
        tools = list(self.values())
        match self.tool_choice:
            case str():
                tools = [self[self.tool_choice]]
            case list():
                tools = [self[name] for name in self.tool_choice]
            case True | None:
                tools = tools
            case _:
                tools = []
        filtered_tools = [info.callable for info in tools if info.matches_filter(state)]
        # Sort by priority (if any have non-default priority)
        if any(self[t.name].priority != 100 for t in filtered_tools):  # noqa: PLR2004
            filtered_tools.sort(key=lambda t: self[t.name].priority)

        return filtered_tools

    def get_tool_names(
        self, state: Literal["all", "enabled", "disabled"] = "all"
    ) -> set[str]:
        """Get tool names based on state."""
        return {name for name, info in self.items() if info.matches_filter(state)}

    def register_tool(
        self,
        tool: LLMCallableTool | AnyCallable | str,
        *,
        name_override: str | None = None,
        description_override: str | None = None,
        enabled: bool = True,
        source: ToolSource = "runtime",
        priority: int = 100,
        requires_confirmation: bool = False,
        requires_capability: str | None = None,
        metadata: dict[str, str] | None = None,
    ) -> ToolInfo:
        """Register a new tool with custom settings.

        Args:
            tool: Tool to register (callable, LLMCallableTool, or import path)
            enabled: Whether tool is initially enabled
            name_override: Optional name override for the tool
            description_override: Optional description override for the tool
            source: Tool source (runtime/agent/builtin/dynamic)
            priority: Execution priority (lower = higher priority)
            requires_confirmation: Whether tool needs confirmation
            requires_capability: Optional capability needed to use tool
            metadata: Additional tool metadata

        Returns:
            Created ToolInfo instance
        """
        # First convert to basic ToolInfo
        match tool:
            case LLMCallableTool():
                llm_tool = tool
                llm_tool.name = name_override or llm_tool.name
                llm_tool.description = description_override or llm_tool.description
            case _:
                llm_tool = LLMCallableTool.from_callable(
                    tool,
                    name_override=name_override,
                    description_override=description_override,
                )

        if llm_tool.description and len(llm_tool.description) > MAX_LEN_DESCRIPTION:
            msg = f"Too long description for {tool}"
            raise ToolError(msg)
        tool_info = ToolInfo(
            llm_tool,
            enabled=enabled,
            source=source,
            priority=priority,
            requires_confirmation=requires_confirmation,
            requires_capability=requires_capability,
            metadata=metadata or {},
        )
        # Register the tool
        self.register(tool_info.name, tool_info)
        return tool_info

    def register_worker(
        self,
        worker: AnyAgent[Any, Any],
        *,
        name: str | None = None,
        reset_history_on_run: bool = True,
        pass_message_history: bool = False,
        share_context: bool = False,
        parent: AnyAgent[Any, Any] | None = None,
    ) -> ToolInfo:
        """Register an agent as a worker tool.

        Args:
            worker: Agent to register as worker
            name: Optional name override for the worker tool
            reset_history_on_run: Whether to clear history before each run
            pass_message_history: Whether to pass parent's message history
            share_context: Whether to pass parent's context/deps
            parent: Optional parent agent for history/context sharing
        """
        tool = worker.to_agent_tool(
            parent=parent,
            name=name,
            reset_history_on_run=reset_history_on_run,
            pass_message_history=pass_message_history,
            share_context=share_context,
        )
        msg = "Registering worker %s as tool %s"
        logger.debug(msg, worker.name, tool.name)
        return self.register_tool(tool, source="agent", metadata={"agent": worker.name})

    def reset(self):
        """Reset tool states."""
        old_tools = self.list_tools()
        self.reset_states()
        new_tools = self.list_tools()

        event = self.ToolStateReset(old_tools, new_tools)
        self.tool_states_reset.emit(event)

    async def setup_mcp_servers(self, servers: list[MCPServerConfig]):
        """Set up multiple MCP server integrations.

        Args:
            servers: List of MCP server configurations
        """
        try:
            for server in servers:
                if not server.enabled:
                    continue

                match server:
                    case StdioMCPServer():
                        # Create client with stdio mode
                        client = MCPClient(stdio_mode=True)
                        client = await self.exit_stack.enter_async_context(client)
                        env = os.environ.copy()
                        # Update with server-specific environment if provided
                        if server.environment:
                            env.update(server.environment)
                        # Ensure UTF-8 encoding
                        env["PYTHONIOENCODING"] = "utf-8"
                        await client.connect(server.command, args=server.args, env=env)
                        # Store client
                        client_id = f"{server.command}_{' '.join(server.args)}"
                        self._mcp_clients[client_id] = client

                        # Register tools
                        self.register_mcp_tools(client)

                    case SSEMCPServer():
                        # SSE client without stdio mode
                        client = MCPClient(stdio_mode=False)
                        client = await self.exit_stack.enter_async_context(client)

                        await client.connect(
                            command="",  # Not used for SSE
                            args=[],  # Not used for SSE
                            url=server.url,
                            env=server.environment,
                        )

                        # Store client
                        client_id = f"sse_{server.url}"
                        self._mcp_clients[client_id] = client

                        # Register tools
                        self.register_mcp_tools(client)

        except Exception as e:
            msg = "Failed to setup MCP servers"
            logger.exception(msg, exc_info=e)
            raise RuntimeError(msg) from e

    async def cleanup(self):
        """Clean up resources including all MCP clients."""
        try:
            # Store tools to remove
            to_remove = [
                name
                for name, info in self.items()
                if info.metadata.get("mcp_tool") is not None
            ]

            try:
                # Clean up exit stack (which includes MCP clients)
                await self.exit_stack.aclose()
            except RuntimeError as e:
                if "different task" in str(e):
                    # Handle task context mismatch
                    current_task = asyncio.current_task()
                    if current_task:
                        loop = asyncio.get_running_loop()
                        await loop.create_task(self.exit_stack.aclose())
                else:
                    raise

            self._mcp_clients.clear()

            # Remove MCP tools
            for name in to_remove:
                del self._items[name]

        except Exception as e:
            msg = "Error during tool manager cleanup"
            logger.exception(msg, exc_info=e)
            raise RuntimeError(msg) from e

    @contextmanager
    def temporary_tools(
        self,
        tools: ToolType | Sequence[ToolType],
        *,
        exclusive: bool = False,
    ) -> Iterator[list[ToolInfo]]:
        """Temporarily register tools.

        Args:
            tools: Tool(s) to register
            exclusive: Whether to temporarily disable all other tools

        Yields:
            List of registered tool infos

        Example:
            ```python
            with tool_manager.temporary_tools([tool1, tool2], exclusive=True) as tools:
                # Only tool1 and tool2 are available
                await agent.run(prompt)
            # Original tool states are restored
            ```
        """
        # Normalize inputs to lists
        tools_list: list[ToolType] = (
            [tools] if not isinstance(tools, Sequence) else list(tools)
        )

        # Store original tool states if exclusive
        original_states: dict[str, bool] = {}
        if exclusive:
            original_states = {name: tool.enabled for name, tool in self.items()}
            # Disable all existing tools
            for t in self.values():
                t.enabled = False

        # Register all tools
        registered_tools: list[ToolInfo] = []
        try:
            for tool in tools_list:
                tool_info = self.register_tool(tool)
                registered_tools.append(tool_info)
            yield registered_tools

        finally:
            # Remove temporary tools
            for tool_info in registered_tools:
                del self[tool_info.name]

            # Restore original tool states if exclusive
            if exclusive:
                for name_, was_enabled in original_states.items():
                    if t := self.get(name_):
                        t.enabled = was_enabled

    def register_mcp_tools(
        self,
        mcp_client: MCPClient,
    ) -> list[ToolInfo]:
        """Register MCP tools with tool manager."""
        registered = []

        for mcp_tool in mcp_client._available_tools:
            # Create properly typed callable from schema
            tool_callable = mcp_client.create_tool_callable(mcp_tool)

            # The function already has proper typing, so no schema override needed
            llm_tool = LLMCallableTool.from_callable(tool_callable)

            metadata = {"mcp_tool": mcp_tool.name}
            tool_info = self.register_tool(llm_tool, source="mcp", metadata=metadata)
            registered.append(tool_info)

            logger.debug("Registered MCP tool: %s", mcp_tool.name)

        return registered

    def tool(
        self,
        name: str | None = None,
        *,
        description: str | None = None,
        enabled: bool = True,
        source: ToolSource = "runtime",
        priority: int = 100,
        requires_confirmation: bool = False,
        requires_capability: str | None = None,
        metadata: dict[str, str] | None = None,
    ) -> Callable[[AnyCallable], AnyCallable]:
        """Decorator to register a function as a tool.

        Args:
            name: Optional override for tool name (defaults to function name)
            description: Optional description override
            enabled: Whether tool is initially enabled
            source: Tool source type
            priority: Execution priority (lower = higher)
            requires_confirmation: Whether tool needs confirmation
            requires_capability: Optional required capability
            metadata: Additional tool metadata

        Returns:
            Decorator function that registers the tool

        Example:
            @tool_manager.register(
                name="search_docs",
                description="Search documentation",
                requires_confirmation=True
            )
            async def search(query: str) -> str:
                '''Search the docs.'''
                return "Results..."
        """

        def decorator(func: AnyCallable) -> AnyCallable:
            self.register_tool(
                func,
                name_override=name,
                description_override=description,
                enabled=enabled,
                source=source,
                priority=priority,
                requires_confirmation=requires_confirmation,
                requires_capability=requires_capability,
                metadata=metadata,
            )
            return func

        return decorator
