"""Tool collection, filtering, namespacing, and routing.

Handles:
- ToolDefinition: represents a single tool from an MCP server
- Tool filtering: include/exclude modes per server
- Namespacing: auto/always/never collision resolution
- Routing table: maps exposed tool names back to origin servers
"""

from __future__ import annotations

from dataclasses import dataclass, field

from hatchdx import HdxError
from hatchdx.agent.config import AgentConfig, ToolFilterConfig


# ---------------------------------------------------------------------------
# Errors
# ---------------------------------------------------------------------------


class ToolFilterError(HdxError):
    """Invalid tool filter configuration."""


class ToolCollisionError(HdxError):
    """Tool name collision with namespacing=never."""


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------


@dataclass
class ToolDefinition:
    """A single tool discovered from an MCP server."""

    name: str
    description: str
    server_name: str
    input_schema: dict = field(default_factory=dict)


@dataclass
class ToolRoute:
    """Maps an exposed tool name to its origin server and original name."""

    exposed_name: str
    original_name: str
    server_name: str


@dataclass
class CollectedTools:
    """Result of collecting and processing tools from all servers."""

    tools: list[ToolDefinition]
    routes: dict[str, ToolRoute]
    filtered_count: int = 0
    total_before_filter: int = 0

    @property
    def tool_count(self) -> int:
        return len(self.tools)

    def get_route(self, exposed_name: str) -> ToolRoute | None:
        return self.routes.get(exposed_name)

    def tools_by_server(self) -> dict[str, list[ToolDefinition]]:
        """Group tools by their origin server."""
        by_server: dict[str, list[ToolDefinition]] = {}
        for tool in self.tools:
            by_server.setdefault(tool.server_name, []).append(tool)
        return by_server


# ---------------------------------------------------------------------------
# Tool filtering
# ---------------------------------------------------------------------------


def validate_tool_filters(
    filters: dict[str, ToolFilterConfig],
    server_tools: dict[str, list[ToolDefinition]],
) -> None:
    """Validate that tool filters reference known servers and tools.

    Raises ToolFilterError if:
    - A filter references a server not in the agent config
    - A filter references a tool that doesn't exist on the server
    """
    for server_name, filter_config in filters.items():
        if server_name not in server_tools:
            available = ", ".join(sorted(server_tools.keys())) or "(none)"
            raise ToolFilterError(
                f"Tool filter references unknown server: '{server_name}'\n"
                f"Available servers: {available}"
            )

        known_tools = {t.name for t in server_tools[server_name]}

        for tool_name in filter_config.include:
            if tool_name not in known_tools:
                available = ", ".join(sorted(known_tools)) or "(none)"
                raise ToolFilterError(
                    f"Tool filter for '{server_name}' references unknown tool: '{tool_name}'\n"
                    f"Available tools: {available}"
                )

        for tool_name in filter_config.exclude:
            if tool_name not in known_tools:
                available = ", ".join(sorted(known_tools)) or "(none)"
                raise ToolFilterError(
                    f"Tool filter for '{server_name}' references unknown tool: '{tool_name}'\n"
                    f"Available tools: {available}"
                )


def apply_tool_filters(
    tools: list[ToolDefinition],
    filters: dict[str, ToolFilterConfig],
) -> list[ToolDefinition]:
    """Apply include/exclude filters to a list of tools.

    For each tool, checks if its server has a filter:
    - include mode: only tools in the include list pass through
    - exclude mode: tools in the exclude list are removed
    - no filter: tool passes through unchanged

    Returns a new list with filtered tools removed.
    """
    result = []
    for tool in tools:
        filter_config = filters.get(tool.server_name)
        if filter_config is None:
            result.append(tool)
            continue

        if filter_config.include:
            if tool.name in filter_config.include:
                result.append(tool)
        elif filter_config.exclude:
            if tool.name not in filter_config.exclude:
                result.append(tool)
        else:
            # Empty filter — pass through
            result.append(tool)

    return result


# ---------------------------------------------------------------------------
# Tool namespacing
# ---------------------------------------------------------------------------


def apply_namespacing(
    tools: list[ToolDefinition],
    mode: str,
) -> tuple[list[ToolDefinition], dict[str, ToolRoute]]:
    """Apply namespacing to tool names based on the configured mode.

    Modes:
    - "auto": flat names if unique, namespace on collision
    - "always": all tools namespaced as server_name:tool_name
    - "never": flat names, error on collision

    Returns:
        Tuple of (tools with updated names, routing table)

    Raises:
        ToolCollisionError: if mode="never" and names collide
    """
    if mode == "always":
        return _namespace_always(tools)
    elif mode == "never":
        return _namespace_never(tools)
    else:
        return _namespace_auto(tools)


def _namespace_always(
    tools: list[ToolDefinition],
) -> tuple[list[ToolDefinition], dict[str, ToolRoute]]:
    """Namespace all tools as server_name:tool_name."""
    result = []
    routes: dict[str, ToolRoute] = {}

    for tool in tools:
        exposed = f"{tool.server_name}:{tool.name}"
        result.append(
            ToolDefinition(
                name=exposed,
                description=tool.description,
                server_name=tool.server_name,
                input_schema=tool.input_schema,
            )
        )
        routes[exposed] = ToolRoute(
            exposed_name=exposed,
            original_name=tool.name,
            server_name=tool.server_name,
        )

    return result, routes


def _namespace_never(
    tools: list[ToolDefinition],
) -> tuple[list[ToolDefinition], dict[str, ToolRoute]]:
    """Use flat names, error on collision."""
    name_map: dict[str, list[str]] = {}
    for tool in tools:
        name_map.setdefault(tool.name, []).append(tool.server_name)

    # Check for collisions
    collisions = {name: servers for name, servers in name_map.items() if len(servers) > 1}
    if collisions:
        parts = []
        for name, servers in sorted(collisions.items()):
            parts.append(f"  '{name}' exists in: {', '.join(servers)}")
        raise ToolCollisionError(
            "Tool name collision with tool_namespacing=never:\n"
            + "\n".join(parts)
            + "\n\nUse tool_namespacing: auto or always to resolve collisions, "
            "or use tool_filter to limit exposed tools."
        )

    routes: dict[str, ToolRoute] = {}
    for tool in tools:
        routes[tool.name] = ToolRoute(
            exposed_name=tool.name,
            original_name=tool.name,
            server_name=tool.server_name,
        )

    return tools, routes


def _namespace_auto(
    tools: list[ToolDefinition],
) -> tuple[list[ToolDefinition], dict[str, ToolRoute]]:
    """Flat names if unique, namespace on collision."""
    # Build name → [server] map
    name_map: dict[str, list[str]] = {}
    for tool in tools:
        name_map.setdefault(tool.name, []).append(tool.server_name)

    colliding_names = {name for name, servers in name_map.items() if len(servers) > 1}

    result = []
    routes: dict[str, ToolRoute] = {}

    for tool in tools:
        if tool.name in colliding_names:
            exposed = f"{tool.server_name}:{tool.name}"
        else:
            exposed = tool.name

        result.append(
            ToolDefinition(
                name=exposed,
                description=tool.description,
                server_name=tool.server_name,
                input_schema=tool.input_schema,
            )
        )
        routes[exposed] = ToolRoute(
            exposed_name=exposed,
            original_name=tool.name,
            server_name=tool.server_name,
        )

    return result, routes


# ---------------------------------------------------------------------------
# Top-level tool collection
# ---------------------------------------------------------------------------


def collect_and_filter_tools(
    server_tools: dict[str, list[ToolDefinition]],
    config: AgentConfig,
) -> CollectedTools:
    """Collect tools from all servers, apply filters and namespacing.

    This is the main entry point for tool processing. It:
    1. Validates tool filters against known tools
    2. Flattens all server tools into a single list
    3. Applies include/exclude filters
    4. Applies namespacing (auto/always/never)
    5. Builds the routing table

    Args:
        server_tools: mapping of server_name → list of ToolDefinition
        config: the agent configuration

    Returns:
        CollectedTools with filtered, namespaced tools and routing table
    """
    # 1. Validate filters reference real servers/tools
    if config.tool_filter:
        validate_tool_filters(config.tool_filter, server_tools)

    # 2. Flatten
    all_tools = []
    for tools in server_tools.values():
        all_tools.extend(tools)

    total_before = len(all_tools)

    # 3. Apply filters
    filtered_tools = apply_tool_filters(all_tools, config.tool_filter)

    # 4. Apply namespacing
    namespaced_tools, routes = apply_namespacing(
        filtered_tools,
        config.settings.tool_namespacing,
    )

    return CollectedTools(
        tools=namespaced_tools,
        routes=routes,
        filtered_count=total_before - len(filtered_tools),
        total_before_filter=total_before,
    )
