# File: src/mcp/server/__init__.py
# Purpose: Provide lightweight MCP server primitives and JSON-RPC loop helpers.
# Context: Used by local MCP servers and tests to register tools and exchange
#          protocol-compliant messages without external dependencies.

"""Lightweight JSON-RPC utilities for minimal MCP server scenarios.

Purpose:
    Supply decorator-based helpers and a stdio event loop so tests and local
    scripts can simulate MCP behaviour without the production stack.
Context:
    Loaded by integration tests that require MCP-compatible responses while
    keeping dependencies minimal for speed and isolation.
Side Effects:
    Installs signal handlers and writes JSON-RPC frames to the configured I/O
    streams during execution.
"""

from __future__ import annotations

import asyncio
import json
import os
import signal
import sys
from collections.abc import Awaitable, Callable, Iterable, Sequence
from dataclasses import asdict, is_dataclass
from typing import IO, Any

ToolListHandler = Callable[[], Awaitable[Sequence[object]]]
ToolCallHandler = Callable[[str, dict[str, Any]], Awaitable[Sequence[object]]]
JsonDict = dict[str, Any]


def _tool_to_dict(tool: object) -> dict[str, object]:
    """Convert an MCP tool representation into JSON-serialisable data.

    Purpose:
        Normalise tool metadata produced by tests or the production stack.

    Args:
        tool: Dataclass or object exposing name/description/inputSchema attributes.

    Returns:
        dict[str, object]: Mapping filtered to exclude ``None`` values.
    Side Effects:
        None.
    """

    if is_dataclass(tool):
        data = asdict(tool)
    else:
        data = {
            "name": getattr(tool, "name", None),
            "description": getattr(tool, "description", None),
            "inputSchema": getattr(tool, "inputSchema", None),
        }
    return {key: value for key, value in data.items() if value is not None}


class Server:
    """Manage MCP tool registration and serve JSON-RPC messages over stdio.

    Purpose:
        Provide decorator-friendly handlers and a simple stdio event loop for
        integration tests.

    Args:
        name: Identifier advertised to MCP clients during initialization.
    Side Effects:
        Reads environment variables to populate protocol metadata.
    """

    def __init__(self, name: str) -> None:
        """Initialise registry state and server metadata.

        Purpose:
            Capture handler placeholders and environment-driven defaults.

        Args:
            name: Human-readable name for the MCP server instance.

        Returns:
            None.
        Side Effects:
            Reads ``MCP_PROTOCOL_VERSION`` and ``MCP_SERVER_VERSION`` from the environment.
        """

        self.name = name
        self._list_tools_handler: ToolListHandler | None = None
        self._call_tool_handler: ToolCallHandler | None = None
        self._protocol_version: str = os.getenv("MCP_PROTOCOL_VERSION", "2024-11-05")
        self._server_info: dict[str, str] = {
            "name": self.name,
            "version": os.getenv("MCP_SERVER_VERSION", "0.1.0"),
        }
        self._capabilities: dict[str, Any] = {
            "tools": {
                "listChanged": False,
            },
        }

    def list_tools(self) -> Callable[[ToolListHandler], ToolListHandler]:
        """Register the async list-tools handler via decorator syntax.

        Purpose:
            Mirror the production decorator so tests can supply coroutines.

        Args:
            None.

        Returns:
            Callable[[ToolListHandler], ToolListHandler]: Decorator that stores and returns the coroutine.
        Side Effects:
            Mutates the internal ``_list_tools_handler`` reference.
        """

        def decorator(fn: ToolListHandler) -> ToolListHandler:
            """Attach the provided coroutine as the list-tools handler.

            Purpose:
                Capture the coroutine for responding to ``tools/list`` requests.

            Args:
                fn: Coroutine returning a sequence of tool-like objects.

            Returns:
                ToolListHandler: The original coroutine for decorator chaining.
            Side Effects:
                Updates ``_list_tools_handler`` with the provided coroutine.
            """

            self._list_tools_handler = fn
            return fn

        return decorator

    def call_tool(self) -> Callable[[ToolCallHandler], ToolCallHandler]:
        """Register the async call-tool handler via decorator syntax.

        Purpose:
            Allow tests to provide the execution coroutine for MCP tools.

        Args:
            None.

        Returns:
            Callable[[ToolCallHandler], ToolCallHandler]: Decorator storing the coroutine.
        Side Effects:
            Mutates the internal ``_call_tool_handler`` reference.
        """

        def decorator(fn: ToolCallHandler) -> ToolCallHandler:
            """Attach the provided coroutine as the call-tool handler.

            Purpose:
                Enable ``tools/call`` requests to delegate to the supplied coroutine.

            Args:
                fn: Coroutine accepting the tool name and JSON-compatible arguments.

            Returns:
                ToolCallHandler: The original coroutine for decorator chaining.
            Side Effects:
                Updates ``_call_tool_handler`` with the provided coroutine.
            """

            self._call_tool_handler = fn
            return fn

        return decorator

    def create_initialization_options(self) -> dict[str, Any]:
        """Build the MCP initialization options response payload.

        Purpose:
            Provide protocol metadata compatible with MCP clients.

        Args:
            None.

        Returns:
            dict[str, Any]: Initialization options including protocol and capabilities.
        Side Effects:
            None.
        """

        return {
            "protocolVersion": self._protocol_version,
            "capabilities": {
                key: (value.copy() if isinstance(value, dict) else value)
                for key, value in self._capabilities.items()
            },
        }

    async def run(
        self,
        read_stream: IO[str] | None = None,
        write_stream: IO[str] | None = None,
        init_opts: dict[str, Any] | None = None,
    ) -> None:
        """Serve JSON-RPC 2.0 frames until completion.

        Purpose:
            Emulate the production MCP stdio loop for integration tests.

        Args:
            read_stream: Optional input stream; defaults to ``sys.stdin``.
            write_stream: Optional output stream; defaults to ``sys.stdout``.
            init_opts: Optional initialization options supplied by the client.

        Returns:
            None.
        Side Effects:
            Installs signal handlers, performs blocking reads via an executor,
            writes JSON to the output stream, and may terminate the process.
        """

        self._setup_signal_handlers()
        reader = read_stream or sys.stdin
        writer = write_stream or sys.stdout

        while True:
            line = await self._read_line(reader)
            if not line:
                await asyncio.sleep(0.05)
                continue

            message = self._parse_message(line)
            if message is None:
                continue

            response, terminate = await self._dispatch_message(message, init_opts)
            if response is None:
                continue

            self._emit_response(response, writer)
            if terminate:
                os._exit(0)
                return

    def _setup_signal_handlers(self) -> None:
        """Install termination handlers compatible with the MCP contract.

        Purpose:
            Exit quickly when the harness sends SIGTERM or SIGINT.

        Args:
            None.

        Returns:
            None.
        Side Effects:
            Calls ``signal.signal`` and may raise if signals are unsupported.
        """

        try:
            signal.signal(signal.SIGTERM, lambda *_: os._exit(0))
            signal.signal(signal.SIGINT, lambda *_: os._exit(0))
        except Exception:
            # Signals may be unsupported on some platforms.
            pass

    async def _read_line(self, reader: IO[str]) -> str:
        """Read a single line from the provided stream without blocking the loop.

        Purpose:
            Delegate blocking I/O to a thread executor to keep asyncio responsive.

        Args:
            reader: Stream object exposing ``readline``.

        Returns:
            str: Line read from the stream (may be an empty string).
        Side Effects:
            Performs file I/O via ``reader.readline``.
        """

        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, reader.readline)

    @staticmethod
    def _parse_message(payload: str) -> JsonDict | None:
        """Parse a JSON-RPC message from raw text.

        Purpose:
            Ignore malformed input gracefully while extracting valid JSON.

        Args:
            payload: Raw line pulled from the input stream.

        Returns:
            dict[str, Any] | None: Parsed JSON message or ``None`` on failure.
        Side Effects:
            None.
        """

        try:
            data = json.loads(payload)
        except Exception:
            return None
        if not isinstance(data, dict):
            return None
        return data

    async def _dispatch_message(
        self,
        message: JsonDict,
        init_opts: dict[str, Any] | None,
    ) -> tuple[JsonDict | None, bool]:
        """Route the incoming message to the appropriate handler.

        Purpose:
            Keep ``run`` concise by separating message dispatch logic.

        Args:
            message: Parsed JSON-RPC request or notification.
            init_opts: Optional initialization overrides supplied to ``run``.

        Returns:
            tuple[JsonDict | None, bool]: Response (if any) and termination flag.
        Side Effects:
            Propagates side effects from downstream handlers.
        """

        message_id = message.get("id")
        method = message.get("method")

        if self._should_ignore_notification(message_id, method):
            return None, False

        if method == "initialize":
            response = self._handle_initialize(message_id, message, init_opts)
            return response, False

        if method == "tools/list":
            response = await self._handle_tools_list(message_id)
            return response, False

        if method == "tools/call":
            response = await self._handle_tools_call(message_id, message)
            return response, True

        response = self._handle_default(message_id)
        return response, False

    @staticmethod
    def _should_ignore_notification(message_id: object, method: object) -> bool:
        """Determine whether the notification should be ignored.

        Purpose:
            Skip ``initialized`` notifications that lack an ``id``.

        Args:
            message_id: Value extracted from the JSON-RPC message.
            method: Method name from the incoming message.

        Returns:
            bool: ``True`` if the message should be ignored, ``False`` otherwise.
        Side Effects:
            None.
        """

        return message_id is None and isinstance(method, str) and method.endswith("initialized")

    def _handle_initialize(
        self,
        message_id: object,
        message: JsonDict,
        init_opts: dict[str, Any] | None,
    ) -> JsonDict:
        """Build the JSON-RPC response for ``initialize`` requests.

        Purpose:
            Merge client-supplied options with server defaults.

        Args:
            message_id: Identifier used to correlate the response.
            message: Full JSON-RPC request payload.
            init_opts: Optional overrides provided to ``run``.

        Returns:
            dict[str, Any]: Initialization JSON-RPC response payload.
        Side Effects:
            None.
        """

        params = message.get("params") or {}
        base_options = self.create_initialization_options()
        merged_options = self._merge_initialization_options(base_options, init_opts)
        requested_protocol = params.get("protocolVersion")
        if isinstance(requested_protocol, str) and requested_protocol.strip():
            merged_options["protocolVersion"] = requested_protocol

        return {
            "jsonrpc": "2.0",
            "id": message_id,
            "result": {
                **merged_options,
                "serverInfo": self._server_info.copy(),
            },
        }

    def _merge_initialization_options(
        self,
        base: dict[str, Any],
        overrides: dict[str, Any] | None,
    ) -> dict[str, Any]:
        """Combine base initialization options with caller overrides.

        Purpose:
            Respect caller preferences while preserving server defaults.

        Args:
            base: Default options produced by ``create_initialization_options``.
            overrides: Optional values supplied by the client or tests.

        Returns:
            dict[str, Any]: Shallow merge of base and override options.
        Side Effects:
            None.
        """

        if not overrides:
            return base
        merged: dict[str, Any] = dict(base)
        for key, value in overrides.items():
            if key == "capabilities" and isinstance(value, dict):
                base_caps = merged.get("capabilities")
                if isinstance(base_caps, dict):
                    merged["capabilities"] = {**base_caps, **value}
                else:
                    merged["capabilities"] = value
            else:
                merged[key] = value
        return merged

    async def _handle_tools_list(self, message_id: object) -> JsonDict:
        """Produce the response payload for ``tools/list`` requests.

        Purpose:
            Invoke the registered list handler and normalise tool data.

        Args:
            message_id: Identifier from the original JSON-RPC request.

        Returns:
            dict[str, Any]: JSON-RPC response containing the available tools.
        Side Effects:
            Executes the registered list handler when present.
        """

        tools: list[dict[str, Any]] = []
        if self._list_tools_handler is not None:
            try:
                raw_tools = await self._list_tools_handler()
                tools = [_tool_to_dict(tool) for tool in raw_tools or []]
            except Exception:
                tools = []
        return {"jsonrpc": "2.0", "id": message_id, "result": {"tools": tools}}

    async def _handle_tools_call(self, message_id: object, message: JsonDict) -> JsonDict:
        """Produce the response payload for ``tools/call`` requests.

        Purpose:
            Execute the registered tool handler and normalise returned content.

        Args:
            message_id: Identifier from the original JSON-RPC request.
            message: Parsed JSON-RPC message containing name and arguments.

        Returns:
            dict[str, Any]: JSON-RPC response with normalised content entries.
        Side Effects:
            Executes the registered call handler and may capture exceptions.
        """

        params = message.get("params") or {}
        name = params.get("name")
        arguments = params.get("arguments") or {}
        content: Sequence[object] | None = None

        if self._call_tool_handler is not None and isinstance(name, str):
            try:
                content = await self._call_tool_handler(name, dict(arguments))
            except Exception as error:
                content = [{"type": "text", "text": json.dumps({"error": str(error)})}]

        normalised = self._normalise_content(content or [])
        return {
            "jsonrpc": "2.0",
            "id": message_id,
            "result": {"content": normalised},
        }

    @staticmethod
    def _normalise_content(content: Iterable[object]) -> list[dict[str, Any]]:
        """Normalise tool execution results into JSON-compatible dictionaries.

        Purpose:
            Ensure downstream tests receive predictable content structures.

        Args:
            content: Iterable returned by the tool handler.

        Returns:
            list[dict[str, Any]]: JSON-compatible content entries.
        Side Effects:
            None.
        """

        normalised: list[dict[str, Any]] = []
        for item in content:
            if is_dataclass(item):
                normalised.append(asdict(item))
            elif isinstance(item, dict):
                normalised.append(dict(item))
            else:
                normalised.append({"type": "text", "text": str(item)})
        return normalised

    @staticmethod
    def _handle_default(message_id: object) -> JsonDict:
        """Return a benign acknowledgement for unrecognised methods.

        Purpose:
            Keep the JSON-RPC loop resilient to unexpected inputs.

        Args:
            message_id: Identifier from the incoming message.

        Returns:
            dict[str, Any]: JSON-RPC response acknowledging the request.
        Side Effects:
            None.
        """

        return {"jsonrpc": "2.0", "id": message_id, "result": {"ok": True}}

    @staticmethod
    def _emit_response(response: JsonDict, writer: IO[str]) -> None:
        """Write a JSON-RPC response to the configured output stream.

        Purpose:
            Centralise response serialisation for stdout and custom streams.

        Args:
            response: JSON-serialisable payload to emit.
            writer: Output stream used to deliver the response.

        Returns:
            None.
        Side Effects:
            Writes to the provided stream and flushes buffers as needed.
        """

        payload = json.dumps(response, ensure_ascii=False)
        if writer is sys.stdout:
            print(payload, flush=True)
        else:
            writer.write(payload + "\n")
            writer.flush()


# Helper re-export for tests and utilities.
server_tool_to_dict = _tool_to_dict
