"""
Per-agent tool registry with whitelisting.

This module provides the AgentToolRegistry class that wraps RLM-Toolkit's
ToolRegistry with per-agent filtering. Each agent only sees the tools
defined in their YAML configuration.

Requirements: AGENT-03 (per-agent tool whitelisting)
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set

if TYPE_CHECKING:
    from rlm_toolkit.tools import Tool, ToolRegistry


class _DictToolRegistry:
    """Simple dict-based tool registry when rlm_toolkit is not available.

    This provides a minimal ToolRegistry-compatible interface for cases
    where RLM-Toolkit is not installed or when a lightweight registry
    is needed.
    """

    def __init__(self):
        self._tools: Dict[str, Any] = {}

    def register(self, tool: Any) -> None:
        """Register a tool."""
        self._tools[tool.name] = tool

    def get(self, name: str) -> Optional[Any]:
        """Get a tool by name."""
        return self._tools.get(name)

    def list_tools(self) -> List[str]:
        """List all registered tool names."""
        return list(self._tools.keys())


class AgentToolRegistry:
    """Per-agent tool registry with whitelisting (AGENT-03).

    This class wraps a global tool registry with per-agent filtering.
    Each agent only sees tools defined in their YAML configuration,
    and tool-specific config can be applied per-agent.

    Example:
        ```python
        registry = AgentToolRegistry()

        # Register global tools
        registry.register_global_tool(read_file_tool)
        registry.register_global_tool(shell_tool)

        # Configure agent's allowed tools
        registry.configure_agent(
            "code-reviewer",
            allowed_tools=["read_file", "shell"],
            tool_configs={"shell": {"timeout": 60}}
        )

        # Get tool for agent (only if whitelisted)
        tool = registry.get_tool("code-reviewer", "read_file")  # Returns tool
        tool = registry.get_tool("code-reviewer", "write_file")  # Returns None
        ```

    Attributes:
        global_registry: Global tool registry
        _agent_tools: Per-agent tool whitelists
        _tool_configs: Per-agent tool configurations
    """

    def __init__(self, global_registry: Optional["ToolRegistry"] = None):
        """Initialize the agent tool registry.

        Args:
            global_registry: Optional tool registry.
                            If not provided, a simple dict-based registry is used.
        """
        # Use provided registry or create a simple dict-based one
        if global_registry is None:
            global_registry = _DictToolRegistry()

        self.global_registry = global_registry
        self._agent_tools: Dict[str, Set[str]] = {}
        self._tool_configs: Dict[str, Dict[str, dict]] = {}

    def register_global_tool(self, tool: "Tool") -> None:
        """Register a tool in the global registry.

        Args:
            tool: Tool instance to register
        """
        self.global_registry.register(tool)

    def configure_agent(
        self,
        agent_id: str,
        allowed_tools: List[str],
        tool_configs: Optional[Dict[str, dict]] = None,
    ) -> None:
        """Configure which tools an agent can use.

        This method sets up the whitelist of tools available to a specific
        agent and optionally configures tool-specific settings.

        Args:
            agent_id: Unique agent identifier
            allowed_tools: List of tool names the agent can use
            tool_configs: Optional dict mapping tool names to config dicts
        """
        self._agent_tools[agent_id] = set(allowed_tools)
        self._tool_configs[agent_id] = tool_configs or {}

    def get_tool(self, agent_id: str, tool_name: str) -> Optional["Tool"]:
        """Get tool for agent if allowed.

        Returns the tool from the global registry only if the tool is
        in the agent's whitelist. Applies any tool-specific configuration.

        Args:
            agent_id: Agent identifier
            tool_name: Name of the tool to retrieve

        Returns:
            Tool instance if allowed and exists, None otherwise
        """
        allowed = self._agent_tools.get(agent_id, set())

        if tool_name not in allowed:
            return None

        tool = self.global_registry.get(tool_name)
        if not tool:
            return None

        # Apply tool-specific config if any
        config = self._tool_configs.get(agent_id, {}).get(tool_name, {})
        if config:
            return self._apply_config(tool, config)

        return tool

    def list_tools(self, agent_id: str) -> List[str]:
        """List tools available to agent.

        Args:
            agent_id: Agent identifier

        Returns:
            Sorted list of tool names available to this agent
        """
        return sorted(self._agent_tools.get(agent_id, set()))

    def is_allowed(self, agent_id: str, tool_name: str) -> bool:
        """Check if tool is allowed for agent.

        Args:
            agent_id: Agent identifier
            tool_name: Name of the tool to check

        Returns:
            True if the tool is in the agent's whitelist
        """
        return tool_name in self._agent_tools.get(agent_id, set())

    def _apply_config(self, tool: "Tool", config: dict) -> "Tool":
        """Apply configuration to tool instance.

        Creates a copy of the tool with the specified configuration
        applied as attribute overrides.

        Args:
            tool: Tool instance to configure
            config: Dictionary of attribute name -> value

        Returns:
            Tool with configuration applied
        """
        # Create a copy if needed, apply config attributes
        for key, value in config.items():
            if hasattr(tool, key):
                setattr(tool, key, value)
        return tool

    def clear_agent(self, agent_id: str) -> None:
        """Clear tool configuration for an agent.

        Args:
            agent_id: Agent identifier to clear
        """
        self._agent_tools.pop(agent_id, None)
        self._tool_configs.pop(agent_id, None)

    def get_all_registered_tools(self) -> List[str]:
        """List all tools in the global registry.

        Returns:
            List of all registered tool names
        """
        return self.global_registry.list_tools()
