from contextlib import AsyncExitStack
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import mcp
from mcp.client.stdio import stdio_client

from otto.agent import Tool


class GatewayError(RuntimeError):
    """Raised when gateway setup or MCP connectivity fails."""


@dataclass(frozen=True)
class Policy:
    blocked_paths: list[str]
    max_calls_per_session: int
    allowed_servers: list[str] | None


@dataclass
class ToolSource:
    """Internal tracking of where a tool came from."""

    name: str
    server: str | None


@dataclass
class _MCPServerHandle:
    name: str
    stack: AsyncExitStack
    session: Any


class Gateway:
    def __init__(self, policy: Policy):
        self._policy = policy
        self._tools: dict[str, Tool] = {}
        self._tool_sources: dict[str, ToolSource] = {}
        self._mcp_servers: dict[str, _MCPServerHandle] = {}
        self._call_count = 0

    async def add_mcp_server(self, name: str, command: list[str]) -> None:
        """Start MCP server over stdio, discover tools, and register them."""
        if self._policy.allowed_servers is not None and name not in self._policy.allowed_servers:
            raise GatewayError(f"Server '{name}' is not in allowed_servers")
        if not command:
            raise GatewayError(f"Failed to start MCP server '{name}': command is empty")

        stack = AsyncExitStack()
        try:
            params = mcp.StdioServerParameters(command=command[0], args=command[1:])
            read_stream, write_stream = await stack.enter_async_context(stdio_client(params))
            session = await stack.enter_async_context(mcp.ClientSession(read_stream, write_stream))
            await session.initialize()
            listed_tools = await session.list_tools()
        except Exception as exc:
            await stack.aclose()
            raise GatewayError(f"Failed to start MCP server '{name}': {exc}") from exc

        self._mcp_servers[name] = _MCPServerHandle(name=name, stack=stack, session=session)

        discovered_tools: list[Tool] = []
        for mcp_tool in listed_tools.tools:
            discovered_tools.append(
                Tool(
                    name=mcp_tool.name,
                    description=mcp_tool.description or "",
                    parameters=mcp_tool.inputSchema,
                    execute=self._mcp_tool_executor(session, mcp_tool.name),
                )
            )

        self._register_tools(discovered_tools, ToolSource(name="mcp", server=name))

    def add_tools(self, tools: list[Tool], source: str = "builtin") -> None:
        """Register tools from any source (built-ins, commands, etc.)."""
        self._register_tools(tools, ToolSource(name=source, server=None))

    def tools(self) -> list[Tool]:
        """Return all registered tools wrapped with policy enforcement."""
        wrapped: list[Tool] = []
        for tool in self._tools.values():
            wrapped.append(
                Tool(
                    name=tool.name,
                    description=tool.description,
                    parameters=tool.parameters,
                    execute=self._policy_wrapped_execute(tool),
                )
            )
        return wrapped

    async def close(self) -> None:
        """Shut down all MCP server connections and subprocesses."""
        for handle in self._mcp_servers.values():
            await handle.stack.aclose()
        self._mcp_servers.clear()

    def _register_tools(self, tools: list[Tool], source: ToolSource) -> None:
        for tool in tools:
            self._tools[tool.name] = tool
            self._tool_sources[tool.name] = source

    def _policy_wrapped_execute(self, tool: Tool):
        async def execute(**kwargs: Any) -> Any:
            denial_reason = self._policy_denial(kwargs)
            if denial_reason is not None:
                return f"Policy denied: {denial_reason}"

            self._call_count += 1
            result = await tool.execute(**kwargs)
            return result

        return execute

    def _policy_denial(self, arguments: dict[str, Any]) -> str | None:
        if self._policy.max_calls_per_session > 0:
            if self._call_count >= self._policy.max_calls_per_session:
                return "maximum tool calls per session exceeded"

        blocked_paths = [str(Path(path).expanduser()) for path in self._policy.blocked_paths]
        for value in _iter_string_values(arguments):
            if "/" not in value and "~" not in value:
                continue

            expanded_value = str(Path(value).expanduser())
            for blocked_path in blocked_paths:
                if not blocked_path:
                    continue
                if blocked_path in value or value.startswith(blocked_path):
                    return f"argument '{value}' matched blocked path '{blocked_path}'"
                if blocked_path in expanded_value or expanded_value.startswith(blocked_path):
                    return f"argument '{value}' matched blocked path '{blocked_path}'"

        return None

    def _mcp_tool_executor(self, session: Any, tool_name: str):
        async def execute(**kwargs: Any) -> str:
            result = await session.call_tool(tool_name, kwargs)
            return _mcp_result_to_text(result)

        return execute


def _iter_string_values(value: Any):
    if isinstance(value, str):
        yield value
        return

    if isinstance(value, dict):
        for item in value.values():
            yield from _iter_string_values(item)
        return

    if isinstance(value, (list, tuple, set)):
        for item in value:
            yield from _iter_string_values(item)


def _mcp_result_to_text(result: Any) -> str:
    if isinstance(result, str):
        return result

    content = getattr(result, "content", None)
    if isinstance(result, dict):
        content = result.get("content")

    if content:
        parts: list[str] = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
                continue
            if isinstance(item, dict):
                text = item.get("text")
                if text is not None:
                    parts.append(str(text))
                    continue
            else:
                text = getattr(item, "text", None)
                if text is not None:
                    parts.append(str(text))
                    continue
            parts.append(str(item))
        if parts:
            return "\n".join(parts)

    structured_content = getattr(result, "structuredContent", None)
    if isinstance(result, dict):
        structured_content = result.get("structuredContent")
    if structured_content is not None:
        return str(structured_content)

    return str(result)
