"""Middleware MCP server implementation for MCP Chain."""

from typing import Any, Dict
from .types import DictMCPServer, MetadataTransformer, RequestResponseTransformer


class MiddlewareMCPServer:
    """An MCP Server that works with dict-based transformers and can chain other dict-based servers."""

    def __init__(
        self,
        downstream_server: DictMCPServer,
        metadata_transformer: MetadataTransformer = None,
        request_transformer: RequestResponseTransformer = None,
    ):
        self._downstream = downstream_server
        self._metadata_transformer = metadata_transformer or (
            lambda next_server, metadata_dict: next_server.get_metadata()
        )
        self._request_transformer = request_transformer or (
            lambda next_server, request_dict: next_server.handle_request(request_dict)
        )

    def get_metadata(self) -> Dict[str, Any]:
        """Returns server metadata as dict."""
        if self._downstream is None:
            raise ValueError("No downstream server configured")

        # Call metadata transformer with empty dict (metadata transformers don't use the metadata_dict parameter)
        return self._metadata_transformer(self._downstream, {})

    def handle_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """Handles dict request and returns dict response."""
        if self._downstream is None:
            raise ValueError("No downstream server configured")

        # Call request transformer with the request dict
        return self._request_transformer(self._downstream, request)

    def then(self, *args):
        """Chain another transformer or server."""
        if self._downstream is None:
            raise ValueError("No downstream server configured")

        # Check if downstream has a then method
        if hasattr(self._downstream, "then"):
            # Delegate to child's then method
            child_result = self._downstream.then(*args)

            # Create a new MiddlewareMCPServer that's a copy of ourselves
            # but with the child_result as the new downstream
            return MiddlewareMCPServer(
                downstream_server=child_result,
                metadata_transformer=self._metadata_transformer,
                request_transformer=self._request_transformer,
            )
        else:
            # Downstream doesn't have `then` method - this should be an error
            # since we can't chain further on a terminal server
            raise ValueError(
                "Cannot chain on downstream server: it has no `then` method"
            )
