"""MCP client implementation over Streamable HTTP (JSON-RPC 2.0)."""

from __future__ import annotations

import json

import httpx

from .errors import (
    DASHBOARD_URL,
    ConfigurationError,
    MCPConnectionError,
    MCPError,
    ToolNotFoundError,
)
from .utils import initiate_oauth, poll_oauth_token


class OAuthRequired(Exception):
    """Raised internally when an MCP server responds with 401."""


# ---------------------------------------------------------------------------
# MCPClient — single MCP server connection
# ---------------------------------------------------------------------------


class MCPClient:
    """Connects to a single MCP server over Streamable HTTP.

    Implements the JSON-RPC 2.0 envelope for MCP protocol methods:
    initialize, tools/list, tools/call.
    """

    def __init__(self, server_url: str, token: str | None = None, *, server_label: str = ""):
        self._url = server_url.rstrip("/")
        self._token = token
        self._request_id = 0
        self._session_id: str | None = None
        self._label = server_label

    def _next_id(self) -> int:
        self._request_id += 1
        return self._request_id

    def _headers(self) -> dict[str, str]:
        h = {
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }
        if self._token:
            h["Authorization"] = f"Bearer {self._token}"
        if self._session_id:
            h["Mcp-Session-Id"] = self._session_id
        return h

    async def _send(self, method: str, params: dict | None = None) -> dict:
        """Send a JSON-RPC request and return the result."""
        payload: dict = {
            "jsonrpc": "2.0",
            "id": self._next_id(),
            "method": method,
        }
        if params is not None:
            payload["params"] = params

        try:
            async with httpx.AsyncClient(timeout=60) as client:
                resp = await client.post(
                    self._url, json=payload, headers=self._headers()
                )
        except (httpx.ConnectError, httpx.TimeoutException) as exc:
            label_hint = f" '{self._label}'" if self._label else ""
            raise MCPConnectionError(
                f"Could not connect to MCP server{label_hint}. "
                f"The server may be unreachable — check your network or the connector URL in the dashboard at {DASHBOARD_URL}."
            ) from exc

        # Persist session ID if the server provides one
        if "mcp-session-id" in resp.headers:
            self._session_id = resp.headers["mcp-session-id"]

        if resp.status_code == 401:
            raise OAuthRequired()

        try:
            resp.raise_for_status()
        except httpx.HTTPStatusError as exc:
            label_hint = f" '{self._label}'" if self._label else ""
            raise MCPConnectionError(
                f"Could not reach MCP server{label_hint} (HTTP {resp.status_code}). "
                f"The server may be down — check its status or verify the connector URL in the dashboard at {DASHBOARD_URL}."
            ) from exc

        content_type = resp.headers.get("content-type", "")
        if "text/event-stream" in content_type:
            return _parse_sse(resp.text, server_label=self._label)

        data = resp.json()
        if "error" in data:
            err = data["error"]
            raise MCPError(
                err.get("message", "Unknown MCP error"),
                code=err.get("code", -1),
                server_label=self._label,
            )
        return data.get("result", {})

    # -- MCP protocol methods ------------------------------------------------

    async def initialize(self) -> dict:
        """Send the MCP `initialize` handshake."""
        return await self._send(
            "initialize",
            {
                "protocolVersion": "2025-03-26",
                "capabilities": {},
                "clientInfo": {"name": "akashic", "version": "0.1.0"},
            },
        )

    async def list_tools(self) -> list[dict]:
        """Send `tools/list` and return the tool definitions."""
        result = await self._send("tools/list")
        return result.get("tools", [])

    async def call_tool(self, name: str, arguments: dict) -> dict:
        """Send `tools/call` and return the result."""
        return await self._send("tools/call", {"name": name, "arguments": arguments})


# ---------------------------------------------------------------------------
# MCPManager — orchestrates multiple MCPClient instances
# ---------------------------------------------------------------------------


class MCPManager:
    """Manages connections to multiple MCP servers.

    Handles OAuth initiation, tool listing, and routing tool calls
    to the correct server.
    """

    def __init__(self):
        self._clients: dict[str, MCPClient] = {}  # label → client
        self._raw_tools: dict[str, list[dict]] = {}  # label → raw MCP tools
        self._tool_map: dict[str, str] = {}  # tool_name → label
        self._tools: list[dict] = []
        self._pending_auth: list[dict] = []

    @property
    def tools(self) -> list[dict]:
        """All MCP tools across all servers in OpenAI function-call format."""
        return self._tools

    @property
    def pending_auth(self) -> list[dict]:
        """Auth URLs still awaiting user completion."""
        return self._pending_auth

    def reset_tools(self) -> None:
        """Clear tool registrations but keep clients and raw tool caches.

        Call this before a new ``connect()`` round so that different
        ``allowedTools`` filters can be applied without re-initializing
        MCP connections.
        """
        self._tool_map.clear()
        self._tools.clear()

    async def connect(
        self, connectors: list[dict], api_key: str
    ) -> tuple[list[dict], list[dict]]:
        """Try to connect to all MCP servers.

        For servers that respond 401, initiates OAuth and adds them to
        ``pending_auth`` instead.

        Returns:
            (tool_schemas, pending_auth)  where tool_schemas is in OpenAI
            function-call format and pending_auth is a list of
            ``{serverLabel, authorizationUrl, sessionId}``.
        """
        self._pending_auth = []

        for conn in connectors:
            label = conn.get("serverLabel", conn.get("server_label", ""))
            url = conn.get("mcpUrl", conn.get("server_url", ""))
            connector_id = conn.get("id")
            allowed = conn.get("allowedTools")

            # Reuse existing client + cached raw tools if already connected
            if label in self._clients and label in self._raw_tools:
                self._add_tools(label, self._raw_tools[label], allowed=allowed)
                continue

            client = MCPClient(url, server_label=label)
            try:
                await client.initialize()
                mcp_tools = await client.list_tools()
                self._clients[label] = client
                self._raw_tools[label] = mcp_tools
                self._add_tools(label, mcp_tools, allowed=allowed)
            except OAuthRequired:
                if not connector_id:
                    raise ConfigurationError(
                        f"Server '{label}' requires sign-in but its connector is missing an ID. "
                        f"This is a configuration issue — re-add the connector in the Akashic dashboard at {DASHBOARD_URL}."
                    )
                result = await initiate_oauth(api_key, connector_id)
                self._pending_auth.append(
                    {
                        "serverLabel": label,
                        "authorizationUrl": result["authorizationUrl"],
                        "sessionId": result["sessionId"],
                        "server_url": url,
                        "allowedTools": allowed,
                    }
                )

        return self._tools, self._pending_auth

    async def complete_auth(self, api_key: str) -> None:
        """Poll for all pending OAuth sessions, then connect and load tools."""
        for auth in self._pending_auth:
            token = await poll_oauth_token(api_key, auth["sessionId"])
            label = auth["serverLabel"]
            client = MCPClient(auth["server_url"], token=token, server_label=label)
            await client.initialize()
            mcp_tools = await client.list_tools()
            self._clients[label] = client
            self._raw_tools[label] = mcp_tools
            self._add_tools(label, mcp_tools, allowed=auth.get("allowedTools"))

        self._pending_auth = []

    async def call_tool(self, name: str, arguments: dict) -> dict:
        """Route a tool call to the MCPClient that owns it."""
        label = self._tool_map.get(name)
        if not label:
            raise ToolNotFoundError(
                f"Tool '{name}' is not registered. "
                f"Make sure you called register() before execute(), "
                f"and that the tool exists in your agent's connectors."
            )
        client = self._clients.get(label)
        if not client:
            raise ToolNotFoundError(
                f"Tool '{name}' was registered but its MCP server is no longer connected. "
                f"Call register() to reconnect."
            )
        return await client.call_tool(name, arguments)

    # -- internal helpers ----------------------------------------------------

    def _add_tools(
        self, label: str, mcp_tools: list[dict], allowed: list[str] | None = None
    ) -> None:
        """Convert MCP tool schemas to OpenAI format and register them.

        Args:
            label: Server label used for routing tool calls.
            mcp_tools: Raw tool definitions from MCP ``tools/list``.
            allowed: If not None, only include tools whose name is in this list.
                     None means all tools are included (backwards-compatible default).
        """
        for tool in mcp_tools:
            name = tool["name"]
            if allowed is not None and name not in allowed:
                continue
            self._tool_map[name] = label
            self._tools.append(
                {
                    "type": "function",
                    "name": name,
                    "description": tool.get("description", ""),
                    "parameters": tool.get(
                        "inputSchema", {"type": "object", "properties": {}}
                    ),
                }
            )


# ---------------------------------------------------------------------------
# SSE helpers
# ---------------------------------------------------------------------------


def _parse_sse(text: str, *, server_label: str = "") -> dict:
    """Extract the JSON-RPC result from an SSE response body."""
    for line in text.split("\n"):
        if not line.startswith("data:"):
            continue
        payload = line[5:].lstrip()
        if not payload:
            continue
        data = json.loads(payload)
        if "result" in data:
            return data["result"]
        if "error" in data:
            err = data["error"]
            raise MCPError(
                err.get("message", "Unknown MCP error"),
                code=err.get("code", -1),
                server_label=server_label,
            )
    raise MCPError(
        "The MCP server returned an unreadable response. "
        f"This is likely a server-side issue — try again or check the connector status in the dashboard at {DASHBOARD_URL}.",
        server_label=server_label,
    )
