import argparse
import inspect
import json
import logging
import os
import sys
from typing import Any, Callable

from callite.client.rpc_client import RPCClient
from callite.mcp.schema import json_type_to_python_type, mcp_schema_to_parameters

try:
    from mcp.server.fastmcp import FastMCP
except ImportError:
    FastMCP = None  # type: ignore


logger = logging.getLogger(__name__)


class MCPBridge:
    """Universal MCP gateway that bridges AI agents to callite RPC services.

    Discovers callite services via their ``__describe__`` endpoint and
    dynamically registers them as MCP tools, resources, and prompts.

    Example::

        bridge = MCPBridge("redis://localhost:6379", ["data_service"])
        bridge.run()
    """

    def __init__(
        self,
        redis_url: str,
        services: list[str],
        *,
        name: str = "callite-bridge",
        execution_timeout: int = 30,
        queue_prefix: str = '/callite',
    ):
        """
        Args:
            redis_url: Redis connection URL.
            services: List of callite service names to bridge.
            name: MCP server display name.
            execution_timeout: Timeout in seconds for callite RPC calls.
            queue_prefix: Redis key prefix (must match the services).
        """
        if FastMCP is None:
            raise ImportError(
                "MCPBridge requires the 'mcp' package. "
                "Install it with: pip install callite[mcp]"
            )

        self._redis_url = redis_url
        self._service_names = services
        self._execution_timeout = execution_timeout
        self._queue_prefix = queue_prefix
        self._clients: dict[str, RPCClient] = {}
        self._mcp = FastMCP(name)

    def run(self, transport: str = "stdio") -> None:
        """Discover services, register MCP primitives, and start the MCP server.

        Args:
            transport: MCP transport type ("stdio" or "streamable-http").
        """
        self._discover_services()
        self._mcp.run(transport=transport)

    def _discover_services(self) -> None:
        """Connect to each service and fetch its tool/prompt/resource metadata."""
        for service_name in self._service_names:
            try:
                client = RPCClient(
                    self._redis_url,
                    service_name,
                    execution_timeout=self._execution_timeout,
                    queue_prefix=self._queue_prefix,
                )
                self._clients[service_name] = client

                description = client.execute('__describe__')
                tools = description.get('tools', [])
                prompts = description.get('prompts', [])
                logger.info(f"Discovered service '{service_name}': "
                            f"{len(tools)} tools, {len(prompts)} prompts")

                for tool_def in tools:
                    self._register_mcp_tool(service_name, client, tool_def)

                for prompt_def in prompts:
                    self._register_mcp_prompt(service_name, client, prompt_def)

                # Use native protocol methods for resources and resource templates
                self._discover_resources(service_name, client)
                self._discover_resource_templates(service_name, client)

            except Exception:
                logger.exception(f"Failed to discover service '{service_name}'")

    def _register_mcp_tool(self, service_name: str, client: RPCClient, tool_def: dict) -> None:
        """Dynamically register a callite method as an MCP tool.

        If the tool definition carries an ``input_schema`` (set by
        ``register_proxy``) and no ``parameters``, the raw JSON Schema is
        converted on-the-fly so the handler gets a proper signature.
        """
        method_name = tool_def['name']
        description = tool_def.get('description', '')
        parameters = tool_def.get('parameters', {})

        # Proxied tools may only have input_schema — convert to parameters
        if not parameters and tool_def.get('input_schema'):
            parameters = mcp_schema_to_parameters(tool_def['input_schema'])

        tool_display = f"{service_name}_{method_name}"

        handler = _make_tool_handler(client, method_name, parameters, description)

        self._mcp.tool(name=tool_display, description=description)(handler)
        logger.debug(f"Registered MCP tool: {tool_display}")

    def _discover_resources(self, service_name: str, client: RPCClient) -> None:
        """Call list_resources on the service and register results as MCP resources."""
        try:
            resources = client.execute('list_resources')
        except Exception:
            logger.debug(f"Service '{service_name}' does not support list_resources")
            return
        for resource_def in resources or []:
            self._register_mcp_resource(service_name, client, resource_def)

    def _discover_resource_templates(self, service_name: str, client: RPCClient) -> None:
        """Call list_resource_templates on the service and register results as MCP resource templates."""
        try:
            templates = client.execute('list_resource_templates')
        except Exception:
            logger.debug(f"Service '{service_name}' does not support list_resource_templates")
            return
        for template_def in templates or []:
            self._register_mcp_resource_template(service_name, client, template_def)

    def _register_mcp_resource(self, service_name: str, client: RPCClient, resource_def: dict) -> None:
        """Register a resource from the service's list_resources response as an MCP resource."""
        uri = resource_def.get('uri', '')
        method_name = resource_def.get('name', '')
        description = resource_def.get('description', '')

        if not uri or not method_name:
            return

        def make_resource_handler(c, m, u):
            def handler() -> str:
                result = c.execute('read_resource', u)
                contents = result.get('contents', []) if isinstance(result, dict) else []
                if contents:
                    return contents[0].get('text', _format_result(result))
                return _format_result(result)
            handler.__name__ = f"resource_{service_name}_{m}"
            handler.__doc__ = description
            return handler

        fn = make_resource_handler(client, method_name, uri)
        self._mcp.resource(uri, description=description)(fn)
        logger.debug(f"Registered MCP resource: {uri}")

    def _register_mcp_resource_template(self, service_name: str, client: RPCClient, template_def: dict) -> None:
        """Register a resource template from the service's list_resource_templates response."""
        uri_template = template_def.get('uriTemplate', '')
        method_name = template_def.get('name', '')
        description = template_def.get('description', '')

        if not uri_template or not method_name:
            return

        def make_template_handler(c, u_template):
            def handler(uri: str) -> str:
                result = c.execute('read_resource', uri)
                contents = result.get('contents', []) if isinstance(result, dict) else []
                if contents:
                    return contents[0].get('text', _format_result(result))
                return _format_result(result)
            handler.__name__ = f"template_{service_name}_{method_name}"
            handler.__doc__ = description
            return handler

        fn = make_template_handler(client, uri_template)
        self._mcp.resource(uri_template, description=description)(fn)
        logger.debug(f"Registered MCP resource template: {uri_template}")

    def _register_mcp_prompt(self, service_name: str, client: RPCClient, prompt_def: dict) -> None:
        """Register a callite prompt template as an MCP prompt."""
        prompt_name = prompt_def['name']
        description = prompt_def.get('description', '')
        arguments = prompt_def.get('arguments', [])

        prompt_display = f"{service_name}_{prompt_name}"
        rpc_name = f"__prompt_{prompt_name}"

        handler = _make_prompt_handler(client, rpc_name, arguments, description)

        self._mcp.prompt(name=prompt_display, description=description)(handler)
        logger.debug(f"Registered MCP prompt: {prompt_display}")

    def close(self) -> None:
        """Close all RPCClient connections."""
        for client in self._clients.values():
            try:
                client.close()
            except Exception:
                pass
        self._clients.clear()


def _make_tool_handler(
    client: RPCClient, method_name: str, parameters: dict[str, Any], description: str
) -> Callable[..., str]:
    """Create a wrapper function for a callite tool with proper type annotations.

    The returned function has a dynamic ``__signature__`` and ``__annotations__``
    so that FastMCP can generate the correct JSON Schema.
    """
    params: list[inspect.Parameter] = []
    annotations: dict[str, Any] = {'return': str}

    for pname, pinfo in parameters.items():
        py_type = json_type_to_python_type(pinfo.get('type', 'string'))
        default = pinfo.get('default', inspect.Parameter.empty)

        params.append(
            inspect.Parameter(
                pname,
                kind=inspect.Parameter.KEYWORD_ONLY,
                default=default,
                annotation=py_type,
            )
        )
        annotations[pname] = py_type

    sig = inspect.Signature(parameters=params)

    def handler(**kwargs: Any) -> str:
        result = client.execute(method_name, **kwargs)
        return _format_result(result)

    handler.__name__ = method_name  # type: ignore[attr-defined]
    handler.__qualname__ = method_name  # type: ignore[attr-defined]
    handler.__doc__ = description
    handler.__signature__ = sig  # type: ignore[attr-defined]
    handler.__annotations__ = annotations

    return handler


def _make_prompt_handler(
    client: RPCClient, rpc_name: str, arguments: list[dict[str, Any]], description: str
) -> Callable[..., str]:
    """Create a wrapper function for a callite prompt with proper type annotations."""
    params: list[inspect.Parameter] = []
    annotations: dict[str, Any] = {'return': str}

    for arg in arguments:
        pname = arg['name']
        required = arg.get('required', True)
        default = inspect.Parameter.empty if required else ''

        params.append(
            inspect.Parameter(
                pname,
                kind=inspect.Parameter.KEYWORD_ONLY,
                default=default,
                annotation=str,
            )
        )
        annotations[pname] = str

    sig = inspect.Signature(parameters=params)

    def handler(**kwargs: Any) -> str:
        result = client.execute(rpc_name, **kwargs)
        return _format_result(result)

    handler.__name__ = rpc_name  # type: ignore[attr-defined]
    handler.__qualname__ = rpc_name  # type: ignore[attr-defined]
    handler.__doc__ = description
    handler.__signature__ = sig  # type: ignore[attr-defined]
    handler.__annotations__ = annotations

    return handler


def _format_result(result: Any) -> str:
    """Convert an arbitrary callite response into a string for MCP."""
    if isinstance(result, str):
        return result
    if isinstance(result, (dict, list)):
        return json.dumps(result, default=str)
    return str(result)


def main():
    """CLI entry point for running the MCP-Callite bridge."""
    parser = argparse.ArgumentParser(
        prog='mcp-callite-bridge',
        description='MCP gateway that bridges AI agents to callite RPC services',
    )
    parser.add_argument(
        '--redis',
        default=os.getenv('REDIS_URL', 'redis://localhost:6379/0'),
        help='Redis connection URL (default: redis://localhost:6379/0)',
    )
    parser.add_argument(
        '--services',
        required=True,
        help='Comma-separated list of callite service names',
    )
    parser.add_argument(
        '--transport',
        default='stdio',
        choices=['stdio', 'streamable-http'],
        help='MCP transport type (default: stdio)',
    )
    parser.add_argument(
        '--name',
        default='callite-bridge',
        help='MCP server display name (default: callite-bridge)',
    )
    parser.add_argument(
        '--timeout',
        type=int,
        default=30,
        help='RPC execution timeout in seconds (default: 30)',
    )
    parser.add_argument(
        '--queue-prefix',
        default='/callite',
        help='Redis key prefix (default: /callite)',
    )

    args = parser.parse_args()

    log_level = os.getenv('LOG_LEVEL', 'INFO')
    logging.basicConfig(
        level=getattr(logging, log_level.upper(), logging.INFO),
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        stream=sys.stderr,
    )

    service_list = [s.strip() for s in args.services.split(',') if s.strip()]
    if not service_list:
        parser.error('--services must contain at least one service name')

    bridge = MCPBridge(
        redis_url=args.redis,
        services=service_list,
        name=args.name,
        execution_timeout=args.timeout,
        queue_prefix=args.queue_prefix,
    )

    try:
        bridge.run(transport=args.transport)
    finally:
        bridge.close()


if __name__ == '__main__':
    main()
