"""
This module contains Handler classes. A Handler manages a collection of Nexus service
handlers. It receives and responds to incoming Nexus requests, dispatching to the
corresponding operation handler.
"""

# A description of the dispatch logic follows.

# There are two cases:

# Case 1: Every user service handler class has a corresponding service definition
# ===============================================================================

# I.e., there are service definitions that look like::

#     @service
#     class MyServiceDefinition:
#         my_op: nexusrpc.Operation[I, O]


# and every service handler class looks like::

#     @service_handler(service=MyServiceDefinition)
#     class MyServiceHandler:
#         @sync_operation
#         def my_op(self, ...)


# Import time
# -----------

# 1. The @service decorator builds a ServiceDefinition instance and attaches it to
#    MyServiceDefinition.

#    The ServiceDefinition contains `name` and a map of OperationDefinition instances, keyed
#    by OperationDefinition.name (this is the publicly advertised name).

#    OperationDefinition contains `name`, `method_name`, and input and output types, all
#    non-None.

#    Operation is the sugar version of OperationDefinition. It is identical to
#    OperationDefinition except that `method_name`, and the input and output types may be
#    None.

# 2. The @sync_operation decorator builds an Operation instance and attaches it to a factory
#    method that is attached to the my_op method object.

# 3. The @service_handler decorator acquires the ServiceDefinition instance from
#    MyServiceDefinition and attaches it to the MyServiceHandler class.


# Handler-registration time
# -------------------------

# 4. Handler.__init__ is called with [MyServiceHandler()]

# 5. A ServiceHandler instance is built from the user service handler class. This comprises
#    a ServiceDefinition and a map {op.name: OperationHandler}. The map is built by taking
#    every operation definition in the service definition and locating the operation handler
#    factory method whose *method name* matches the method name of the operation definition.

# 6. Finally we build a map {service_definition.name: ServiceHandler} using the service
#    definition in each ServiceHandler.

# Request-handling time
# ---------------------

# Now suppose a request has arrived for service s and operation o.

# 7. The Handler does self.service_handlers[s], yielding an instance of ServiceHandler.

# 8. The ServiceHandler does self.operation_handlers[o], yielding an instance of
#    OperationHandler

# Therefore we require that Handler.service_handlers and ServiceHandler.operation_handlers
# are keyed by the publicly advertised service and operation name respectively. This was
# achieved at steps (3) and (2) respectively.


# Case 2: There exists a user service handler class without a corresponding service definition
# ============================================================================================

# I.e., at least one user service handler class looks like::

#     @service_handler
#     class MyServiceHandler:
#         @sync_operation
#         def my_op(...)

# This follows Case 1 with the following differences at import time:

# - Step (1) does not occur.
# - At step (3) the ServiceDefinition is synthesized by the @service_handler decorator
#   from MyServiceHandler. An error will be thrown if the user did not supply input and
#   output type annotations on their operation handler methods (they do not have to do
#   this under Case 1).

from __future__ import annotations

import asyncio
import concurrent.futures
from abc import ABC, abstractmethod
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Callable, Optional, Union, cast

from typing_extensions import Self, TypeGuard

from nexusrpc._common import HandlerError, HandlerErrorType
from nexusrpc._serializer import LazyValueT
from nexusrpc._service import ServiceDefinition
from nexusrpc._util import get_service_definition, is_async_callable

from ._common import (
    CancelOperationContext,
    OperationContext,
    StartOperationContext,
    StartOperationResultAsync,
    StartOperationResultSync,
)
from ._operation_handler import (
    MiddlewareSafeOperationHandler,
    OperationHandler,
    collect_operation_handler_factories_by_method_name,
)


class AbstractHandler(ABC):
    """
    .. warning::
        This API is experimental and may change in the future.
    """

    @abstractmethod
    def start_operation(
        self,
        ctx: StartOperationContext,
        input: LazyValueT,
    ) -> Union[
        StartOperationResultSync[Any],
        StartOperationResultAsync,
        Awaitable[
            Union[
                StartOperationResultSync[Any],
                StartOperationResultAsync,
            ]
        ],
    ]: ...

    @abstractmethod
    def cancel_operation(
        self, ctx: CancelOperationContext, token: str
    ) -> Union[None, Awaitable[None]]:
        """Handle a Cancel Operation request.

        Args:
            ctx: The operation context.
            token: The operation token.
        """
        ...


class BaseServiceCollectionHandler(AbstractHandler, ABC):
    """
    A Nexus handler, managing a collection of Nexus service handlers.

    Operation requests are delegated to a :py:class:`ServiceHandler` based on the service
    name in the operation context.

    .. warning::
        This API is experimental and may change in the future.
    """

    def __init__(
        self,
        user_service_handlers: Sequence[Any],
        executor: Optional[concurrent.futures.Executor] = None,
    ):
        """Initialize a :py:class:`Handler` instance from user service handler instances.

        The user service handler instances must have been decorated with the
        :py:func:`@nexusrpc.handler.service_handler` decorator.

        Args:
            user_service_handlers: A sequence of user service handlers.
            executor: A concurrent.futures.Executor in which to run non-`async def` operation handlers.
        """
        self.executor = _Executor(executor) if executor else None
        self.service_handlers = self._register_service_handlers(user_service_handlers)

    def _register_service_handlers(
        self, user_service_handlers: Sequence[Any]
    ) -> Mapping[str, ServiceHandler]:
        service_handlers = {}
        for sh in user_service_handlers:
            if isinstance(sh, type):
                raise TypeError(
                    f"Expected a service instance, but got a class: {type(sh)}. "
                    "Nexus service handlers must be supplied as instances, not classes."
                )
            # Users may register ServiceHandler instances directly.
            if not isinstance(sh, ServiceHandler):
                # It must be a user service handler instance (i.e. an instance of a class
                # decorated with @nexusrpc.handler.service_handler).
                sh = ServiceHandler.from_user_instance(sh)
            if sh.service.name in service_handlers:
                raise RuntimeError(
                    f"Service '{sh.service.name}' has already been registered."
                )
            service_handlers[sh.service.name] = sh
        return service_handlers

    def _get_service_handler(self, service_name: str) -> ServiceHandler:
        """Return a service handler, given the service name."""
        service = self.service_handlers.get(service_name)
        if service is None:
            raise HandlerError(
                f"No handler for service '{service_name}'.",
                type=HandlerErrorType.NOT_FOUND,
            )
        return service


class Handler(BaseServiceCollectionHandler):
    """
    A Nexus handler manages a collection of Nexus service handlers.

    Operation requests are dispatched to a :py:class:`ServiceHandler` based on the
    service name in the operation context.

    This class supports user operation handlers that are either `async def` or `def`. If
    `def` user operation handlers are to be supported, an executor must be provided.

    The methods of this class itself are `async def`. There is currently no alternative
    Handler class with `def` methods.

    .. warning::
        This API is experimental and may change in the future.

    Example:
        .. code-block:: python

            import concurrent.futures
            from nexusrpc.handler import Handler

            # Create service handler instances
            my_service = MyServiceHandler()

            # Create handler with async operations only
            handler = Handler([my_service])

            # Create handler that supports both async and sync operations
            executor = concurrent.futures.ThreadPoolExecutor(max_workers=10)
            handler = Handler([my_service], executor=executor)

            # Use handler to process requests
            result = await handler.start_operation(ctx, input_lazy_value)
    """

    def __init__(
        self,
        user_service_handlers: Sequence[Any],
        executor: Optional[concurrent.futures.Executor] = None,
        middleware: Sequence[OperationHandlerMiddleware] | None = None,
    ):
        self._middleware = cast(Sequence[OperationHandlerMiddleware], middleware or [])
        super().__init__(user_service_handlers, executor=executor)
        if not self.executor:
            self._validate_all_operation_handlers_are_async()

    async def start_operation(
        self,
        ctx: StartOperationContext,
        input: LazyValueT,
    ) -> Union[
        StartOperationResultSync[Any],
        StartOperationResultAsync,
    ]:
        """Handle a Start Operation request.

        Args:
            ctx: The operation context.
            input: The input to the operation, as a LazyValue.
        """
        service_handler = self._get_service_handler(ctx.service)
        op_handler = self._get_operation_handler(ctx, service_handler, ctx.operation)

        op_defn = service_handler.service.operation_definitions[ctx.operation]
        deserialized_input = await input.consume(as_type=op_defn.input_type)
        return await op_handler.start(ctx, deserialized_input)

    async def cancel_operation(self, ctx: CancelOperationContext, token: str) -> None:
        """Handle a Cancel Operation request.

        Args:
            ctx: The operation context.
            token: The operation token.
        """
        service_handler = self._get_service_handler(ctx.service)
        op_handler = self._get_operation_handler(ctx, service_handler, ctx.operation)
        return await op_handler.cancel(ctx, token)

    def _get_operation_handler(
        self, ctx: OperationContext, service_handler: ServiceHandler, operation: str
    ) -> MiddlewareSafeOperationHandler:
        """
        Get the specified handler for the specified operation from the given service_handler and apply all middleware.
        """
        op_handler: MiddlewareSafeOperationHandler = _EnsuredAwaitableOperationHandler(
            self.executor, service_handler.get_operation_handler(operation)
        )

        for middleware in reversed(self._middleware):
            op_handler = middleware.intercept(ctx, op_handler)

        return op_handler

    def _validate_all_operation_handlers_are_async(self) -> None:
        for service_handler in self.service_handlers.values():
            for op_handler in service_handler.operation_handlers.values():
                _ = self._assert_async_callable(op_handler.start)
                _ = self._assert_async_callable(op_handler.cancel)

    def _assert_async_callable(
        self, method: Callable[..., Any]
    ) -> TypeGuard[Callable[..., Awaitable[Any]]]:
        if not is_async_callable(method):
            raise RuntimeError(
                f"Operation handler method {method} is not an `async def` method, "
                f"but you have not supplied an executor to nexusrpc.handler.Handler. "
            )
        return True


@dataclass(frozen=True)
class ServiceHandler:
    """Internal representation of an instance of a user's service handler class.

    A service handler is a class decorated with
    :py:func:`@nexusrpc.handler.service_handler` that defines operation handler methods
    using decorators such as :py:func:`@nexusrpc.handler.operation_handler`.

    Instances of this class are created automatically from user service implementation
    instances on creation of a Handler instance, at Nexus handler start time. While the
    user's class defines operation handlers as factory methods to be called at handler
    start time, this class contains the :py:class:`OperationHandler` instances
    themselves.

    You may create instances of this class manually and pass them to the Handler
    constructor, for example when programmatically creating Nexus service
    implementations.
    """

    service: ServiceDefinition
    operation_handlers: dict[str, OperationHandler[Any, Any]]

    @classmethod
    def from_user_instance(cls, user_instance: Any) -> Self:
        """Create a :py:class:`ServiceHandler` from a user service instance."""

        service = get_service_definition(user_instance.__class__)
        if not isinstance(service, ServiceDefinition):
            raise RuntimeError(
                f"Service '{user_instance}' does not have a service definition. "
                f"Use the :py:func:`@nexusrpc.handler.service_handler` decorator on your class to define "
                f"a Nexus service implementation."
            )

        # Construct a map of operation handlers keyed by the op name from the service
        # definition (i.e. by the name by which the operation can be requested)
        factories_by_method_name = collect_operation_handler_factories_by_method_name(
            user_instance.__class__, service
        )
        op_handlers = {
            name: factories_by_method_name[defn.method_name](user_instance)
            for name, defn in service.operation_definitions.items()
        }
        return cls(
            service=service,
            operation_handlers=op_handlers,
        )

    def get_operation_handler(self, operation_name: str) -> OperationHandler[Any, Any]:
        """Return an operation handler, given the operation name."""
        if operation_name not in self.service.operation_definitions:
            raise HandlerError(
                f"Nexus service definition '{self.service.name}' has no operation "
                f"'{operation_name}'. There are {len(self.service.operation_definitions)} operations "
                f"in the definition.",
                type=HandlerErrorType.NOT_FOUND,
            )
        operation_handler = self.operation_handlers.get(operation_name)
        if operation_handler is None:
            raise HandlerError(
                f"Nexus service implementation '{self.service.name}' has no handler for "
                f"operation '{operation_name}'. There are {len(self.operation_handlers)} "
                f"available operation handlers.",
                type=HandlerErrorType.NOT_FOUND,
            )
        return operation_handler


class _Executor:
    """An executor for synchronous functions."""

    # Users are require to pass in a `concurrent.futures.Executor` in order to use
    # non-`async def`s. This is what `run_in_executor` is documented to require.
    # This means that nexusrpc initially has a hard-coded dependency on the asyncio
    # event loop; if necessary we can add the ability to pass in an event loop
    # implementation at the level of Handler.

    def __init__(self, executor: concurrent.futures.Executor):
        self._executor = executor

    def submit_to_event_loop(
        self, fn: Callable[..., Any], *args: Any
    ) -> Awaitable[Any]:
        return asyncio.get_event_loop().run_in_executor(self._executor, fn, *args)

    def submit(
        self, fn: Callable[..., Any], *args: Any
    ) -> concurrent.futures.Future[Any]:
        return self._executor.submit(fn, *args)


class OperationHandlerMiddleware(ABC):
    """
    Middleware for operation handlers.

    This should be extended by any operation handler middelware.

    .. warning::
        This API is experimental and may change in the future.
    """

    @abstractmethod
    def intercept(
        self,
        ctx: OperationContext,  # type: ignore[reportUnusedParameter]
        next: MiddlewareSafeOperationHandler,
    ) -> MiddlewareSafeOperationHandler:
        """
        Method called for intercepting operation handlers.

        Args:
            ctx: The :py:class:`OperationContext` that will be passed to the operation handler.
            next: The underlying operation handler that this middleware
                should delegate to.

        Returns:
            The new middleware that will be used to invoke
            :py:attr:`OperationHandler.start` or :py:attr:`OperationHandler.cancel`.
        """
        ...


class _EnsuredAwaitableOperationHandler(MiddlewareSafeOperationHandler):
    """
    An :py:class:`AwaitableOperationHandler` that wraps an :py:class:`OperationHandler` and uses an :py:class:`_Executor` to ensure
    that the :py:attr:`start` and :py:attr:`cancel` methods are awaitable.
    """

    def __init__(
        self,
        executor: _Executor | None,
        op_handler: OperationHandler[Any, Any],
    ):
        self._executor = executor
        self._op_handler = op_handler

    async def start(
        self, ctx: StartOperationContext, input: Any
    ) -> StartOperationResultSync[Any] | StartOperationResultAsync:
        """
        Start the operation using the wrapped :py:class:`OperationHandler`.
        """
        if is_async_callable(self._op_handler.start):
            return await self._op_handler.start(ctx, input)
        else:
            assert self._executor
            return await self._executor.submit_to_event_loop(
                self._op_handler.start, ctx, input
            )

    async def cancel(self, ctx: CancelOperationContext, token: str) -> None:
        """
        Cancel an operation using the wrapped :py:class:`OperationHandler`.
        """
        if is_async_callable(self._op_handler.cancel):
            return await self._op_handler.cancel(ctx, token)
        else:
            assert self._executor
            return self._executor.submit(self._op_handler.cancel, ctx, token).result()
