import json
import logging
import pickle
import threading
import time
from typing import Any, Callable

import redis
from tenacity import retry

from callite.mcp.schema import extract_tool_metadata, mcp_schema_to_parameters
from callite.rpctypes.response import Response
from callite.shared.redis_connection import RedisConnection

logger = logging.getLogger(__name__)


class RPCServer(RedisConnection):
    _log_level_default = 'INFO'

    def __init__(self, conn_url: str, service: str, *args, **kwargs):
        super().__init__(conn_url, service, *args, **kwargs)
        self._registered_methods: dict[str, dict[str, Any]] = {}
        self._tool_metadata: dict[str, dict[str, Any]] = {}
        self._prompt_metadata: dict[str, dict[str, Any]] = {}
        self._registered_prompts: dict[str, Callable[..., Any]] = {}
        self._protocol_methods: set[str] = set()
        self._xread_groupname = kwargs.get('xread_groupname', 'generic')

        # Auto-register the __describe__ discovery endpoint and MCP protocol methods
        self.register_method(self._describe, '__describe__', returns=True)
        self._register_protocol_methods()

        self._subscribe_redis_thread = threading.Thread(target=self._subscribe_redis, daemon=True)
        self._subscribe_redis_thread.start()
        self._logger.debug("Server started")

    @property
    def _request_stream(self) -> str:
        return f'{self._queue_prefix}/request/{self._service}'

    def subscribe(self, handler: Callable, method_name: str | None = None) -> None:
        self.register_method(handler, method_name, False)

    def register(self, handler: Callable, method_name: str | None = None) -> Callable:
        return self.register_method(handler, method_name, True)

    def register_tool(self, handler: Callable[..., Any] | None = None, *, description: str | None = None, method_name: str | None = None) -> Callable[..., Any]:
        """Register a method with rich metadata for MCP tool discovery.

        Can be used as a plain decorator or as a decorator factory with arguments:

            @server.register_tool
            def my_func(x: int): ...

            @server.register_tool(description="Does something useful")
            def my_func(x: int): ...

        Args:
            handler: The function to register (when used without parentheses).
            description: Human-readable description for MCP tool listing.
            method_name: Override the method name (defaults to handler.__name__).
        """
        def decorator(fn: Callable) -> Callable:
            name = method_name or fn.__name__
            metadata = extract_tool_metadata(fn, description=description, method_name=name)
            self._tool_metadata[name] = metadata
            return self.register_method(fn, name, returns=True)

        if handler is not None:
            return decorator(handler)
        return decorator

    def register_prompt(self, handler: Callable[..., Any] | None = None, *, description: str | None = None, method_name: str | None = None) -> Callable[..., Any]:
        """Register a prompt template for MCP prompt discovery.

        The handler function should accept keyword arguments and return a string
        that serves as the prompt text.

            @server.register_prompt(description="Analyze data")
            def analyze(data: str) -> str:
                return f"Please analyze: {data}"

        Args:
            handler: The function to register (when used without parentheses).
            description: Human-readable description for MCP prompt listing.
            method_name: Override the prompt name (defaults to handler.__name__).
        """
        def decorator(fn: Callable) -> Callable:
            name = method_name or fn.__name__
            metadata = extract_tool_metadata(fn, description=description, method_name=name)

            prompt_args = []
            for param_name, param_info in metadata['parameters'].items():
                prompt_args.append({
                    'name': param_name,
                    'description': param_info.get('description', ''),
                    'required': param_info.get('required', True),
                })

            self._prompt_metadata[name] = {
                'name': name,
                'description': metadata['description'],
                'arguments': prompt_args,
            }

            rpc_name = f'__prompt_{name}'
            self._registered_prompts[name] = fn
            self.register_method(fn, rpc_name, returns=True)
            return fn

        if handler is not None:
            return decorator(handler)
        return decorator

    def register_proxy(self, proxy, *, prefix: str | None = None) -> list[str]:
        """Register tools from an external MCP server via :class:`~callite.mcp.proxy.MCPProxy`.

        Connects to the proxy (if not already connected), discovers its tools,
        and registers a callite RPC method for each one.  The method forwards
        every incoming request to the external subprocess.

        Args:
            proxy: An :class:`~callite.mcp.proxy.MCPProxy` instance.
            prefix: Optional prefix prepended to each tool name
                (e.g. ``"sqlite"`` → ``"sqlite_list_tables"``).

        Returns:
            List of registered method names.
        """
        self._logger.info("Registering tools from MCP proxy")
        tools = proxy.get_tools()
        registered: list[str] = []
        self._logger.debug(f"Discovered {len(tools)} tools from proxy")

        for tool in tools:
            method_name = f"{prefix}_{tool.name}" if prefix else tool.name

            # Convert the MCP inputSchema to callite's parameter format
            parameters = mcp_schema_to_parameters(
                getattr(tool, 'inputSchema', None)
            )

            metadata: dict[str, Any] = {
                'name': method_name,
                'description': getattr(tool, 'description', '') or '',
                'parameters': parameters,
            }
            # Preserve the raw JSON Schema for richer bridge pass-through
            if getattr(tool, 'inputSchema', None):
                metadata['input_schema'] = tool.inputSchema

            self._tool_metadata[method_name] = metadata

            # Build a synchronous wrapper that delegates to the proxy
            def _make_handler(p, tool_name):
                def handler(*_args, **kwargs):
                    return p.call_tool(tool_name, kwargs)
                handler.__name__ = tool_name
                return handler

            handler = _make_handler(proxy, tool.name)
            self.register_method(handler, method_name, returns=True)
            registered.append(method_name)
            self._logger.debug("Registered proxy tool: %s", method_name)

        # Register MCP protocol-level methods
        protocol_methods: list[tuple[str, Any]] = [
            ("list_tools", proxy.list_tools),
            ("list_prompts", proxy.list_prompts),
            ("get_prompt", proxy.get_prompt),
            ("list_resources", proxy.list_resources),
            ("read_resource", proxy.read_resource),
            ("list_resource_templates", proxy.list_resource_templates),
            ("send_ping", proxy.send_ping),
        ]
        for proto_name, proto_fn in protocol_methods:
            method_name = f"{prefix}_{proto_name}" if prefix else proto_name
            self.register_method(proto_fn, method_name, returns=True)
            self._protocol_methods.add(method_name)
            registered.append(method_name)
            self._logger.debug("Registered proxy protocol method: %s", method_name)

        self._logger.info(
            "Registered %d tools from MCP proxy (%s)",
            len(registered),
            proxy.command,
        )
        return registered

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

    def _register_protocol_methods(self) -> None:
        """Auto-register MCP protocol-level methods on the server."""
        protocol: list[tuple[str, Callable]] = [
            ('list_tools', self._list_tools),
            ('list_prompts', self._list_prompts),
            ('get_prompt', self._get_prompt),
            ('list_resources', self._list_resources),
            ('read_resource', self._read_resource),
            ('list_resource_templates', self._list_resource_templates),
            ('send_ping', self._send_ping),
        ]
        for name, fn in protocol:
            self.register_method(fn, name, returns=True)
            self._protocol_methods.add(name)

    def _list_tools(self) -> list[dict[str, Any]]:
        """Return all registered tools as a list of MCP-compatible dicts."""
        tools = []
        for method_name, method_info in self._registered_methods.items():
            if method_name.startswith('__') or not method_info['returns']:
                continue
            if method_name in self._protocol_methods:
                continue
            if method_name in self._tool_metadata:
                tools.append(self._tool_metadata[method_name])
            else:
                metadata = extract_tool_metadata(method_info['func'], method_name=method_name)
                tools.append(metadata)
        return tools

    def _list_prompts(self) -> list[dict[str, Any]]:
        """Return all registered prompts as a list of MCP-compatible dicts."""
        return list(self._prompt_metadata.values())

    def _get_prompt(self, name: str, arguments: dict[str, Any] | None = None) -> dict[str, Any]:
        """Fetch a prompt by name, rendering it with the supplied arguments.

        Args:
            name: The prompt name to retrieve.
            arguments: Optional keyword arguments passed to the prompt handler.

        Returns:
            A dict with ``name``, ``description``, and ``messages`` keys.
        """
        if name not in self._registered_prompts:
            raise KeyError(f"Prompt '{name}' not found")
        fn = self._registered_prompts[name]
        text = fn(**(arguments or {}))
        prompt_meta = self._prompt_metadata.get(name, {})
        return {
            'name': name,
            'description': prompt_meta.get('description', ''),
            'messages': [{'role': 'user', 'content': {'type': 'text', 'text': str(text)}}],
        }

    def _list_resources(self) -> list[dict[str, Any]]:
        """Return resources derived from parameterless get_/read_ tools.

        A tool is surfaced as a resource when its name starts with ``get_``
        or ``read_`` and it has no required parameters.
        """
        resources = []
        for method_name, method_info in self._registered_methods.items():
            if method_name.startswith('__') or not method_info['returns']:
                continue
            if method_name in self._protocol_methods:
                continue
            if not method_name.startswith(('get_', 'read_')):
                continue
            metadata = self._tool_metadata.get(method_name) or extract_tool_metadata(
                method_info['func'], method_name=method_name
            )
            has_required = any(
                p.get('required', True)
                for p in metadata.get('parameters', {}).values()
            )
            if has_required:
                continue
            uri = f"callite://{self._service}/{method_name}"
            resources.append({
                'uri': uri,
                'name': method_name,
                'description': metadata.get('description', ''),
                'mimeType': 'text/plain',
            })
        return resources

    def _read_resource(self, uri: str) -> dict[str, Any]:
        """Read a resource identified by its URI.

        Args:
            uri: The resource URI (e.g. ``callite://my_service/get_status``).

        Returns:
            A dict with ``uri`` and ``contents`` keys.
        """
        prefix = f"callite://{self._service}/"
        if not uri.startswith(prefix):
            raise ValueError(f"Unknown resource URI: {uri!r}")
        method_name = uri[len(prefix):]
        if method_name not in self._registered_methods:
            raise KeyError(f"Resource method '{method_name}' not registered")
        result = self._registered_methods[method_name]['func']()
        text = result if isinstance(result, str) else json.dumps(result, default=str)
        return {
            'uri': uri,
            'contents': [{'uri': uri, 'mimeType': 'text/plain', 'text': text}],
        }

    def _list_resource_templates(self) -> list[dict[str, Any]]:
        """Return resource templates derived from get_/read_ tools that have parameters."""
        templates = []
        for method_name, method_info in self._registered_methods.items():
            if method_name.startswith('__') or not method_info['returns']:
                continue
            if method_name in self._protocol_methods:
                continue
            if not method_name.startswith(('get_', 'read_')):
                continue
            metadata = self._tool_metadata.get(method_name) or extract_tool_metadata(
                method_info['func'], method_name=method_name
            )
            params = metadata.get('parameters', {})
            if not params:
                continue
            # Build a URI template from required parameters, e.g. get_user/{user_id}
            required_params = [p for p, info in params.items() if info.get('required', True)]
            if not required_params:
                continue
            uri_template = f"callite://{self._service}/{method_name}/" + "/".join(
                f"{{{p}}}" for p in required_params
            )
            templates.append({
                'uriTemplate': uri_template,
                'name': method_name,
                'description': metadata.get('description', ''),
                'mimeType': 'text/plain',
            })
        return templates

    def _send_ping(self) -> dict[str, str]:
        """Respond to a protocol ping with service info."""
        return {'status': 'ok', 'service': self._service}

    def register_method(self, handler: Callable, method_name: str | None = None, returns: bool = True) -> Callable:
        method_name = method_name or handler.__name__
        self._logger.debug(f"Registering method {method_name}")
        self._registered_methods[method_name] = {'func': handler, 'returns': returns}
        return handler

    def _describe(self) -> dict:
        """Discovery endpoint returning metadata about registered tools and prompts."""
        tools = []
        for method_name, method_info in self._registered_methods.items():
            if method_name.startswith('__'):
                continue
            if not method_info['returns']:
                continue
            if method_name in self._protocol_methods:
                continue

            if method_name in self._tool_metadata:
                tools.append(self._tool_metadata[method_name])
            else:
                metadata = extract_tool_metadata(method_info['func'], method_name=method_name)
                tools.append(metadata)

        return {
            'service': self._service,
            'tools': tools,
            'prompts': list(self._prompt_metadata.values()),
        }

    def run_forever(self) -> None:
        while self._running:
            time.sleep(1000000)

    @retry
    def _subscribe_redis(self):
        while self._running:
            if not self._check_connection():
                self._connect()
            self._ensure_consumer_group()
            self._logger.debug("Checking for messages")
            messages = self._read_messages_from_redis()
            self._logger.debug(f"Received {len(messages)} messages")
            self._process_messages(messages)

    def _check_connection(self):
        try:
            self._rds.ping()
            return True
        except redis.exceptions.ConnectionError:
            return False

    def _ensure_consumer_group(self):
        try:
            self._rds.xgroup_create(self._request_stream, self._xread_groupname, mkstream=True)
        except redis.exceptions.ResponseError as e:
            if "BUSYGROUP" not in str(e) and "already exists" not in str(e):
                raise

    def _read_messages_from_redis(self):
        messages = self._rds.xreadgroup(
            self._xread_groupname, self._connection_id,
            {self._request_stream: '>'}, count=1, block=1000
        )
        self._logger.debug(f"{len(messages)} messages received from {self._request_stream}")
        return messages

    def _process_messages(self, messages):
        for _, message_list in messages:
            for _message in message_list:
                message_id, message_data = _message
                request = pickle.loads(message_data[b'data'])
                self._logger.info(f"Processing message {message_id} with data: {request}")
                threading.Thread(target=self._process_single, args=(request, message_id), daemon=True).start()

    def _process_single(self, request, message_id):
        self._rds.xack(self._request_stream, self._xread_groupname, message_id)

        if request.method not in self._registered_methods:
            self._logger.error(f"Method {request.method} not registered")
            error_response = Response(self._service, message_id, status='error', error=f"Method {request.method} not registered")
            payload = pickle.dumps({'data': error_response, 'request_id': request.request_id})
            self._rds.publish(f'{self._queue_prefix}/response/{request.client_id}', payload)
            return

        method_info = self._registered_methods[request.method]
        if method_info['returns']:
            response = self._call_method(request.method, message_id, *request.args, **request.kwargs)
            payload = pickle.dumps({'data': response, 'request_id': request.request_id})
            self._rds.publish(f'{self._queue_prefix}/response/{request.client_id}', payload)
            self._logger.info(f"Processed message {message_id} and response published to {self._queue_prefix}/response/{request.request_id}")
        else:
            self._call_method(request.method, message_id, *request.args, **request.kwargs)
            self._logger.info(f"Processed message {message_id} without response")

    def _call_method(self, method: str, message_id, *args, **kwargs) -> Response:
        try:
            message_id = message_id.decode('utf-8') if isinstance(message_id, bytes) else message_id
            data = self._registered_methods[method]['func'](*args, **kwargs)
            response = Response(self._service, message_id)
            response.data = data
            return response
        except Exception as e:
            self._logger.error(e)
            return Response(self._service, message_id, status='error', error=str(e))
