"""Base class for server-side interceptors."""

import abc
from asyncio import iscoroutine
from typing import Any, Callable, Tuple

import grpc
from grpc import aio as grpc_aio  # Needed for grpcio pre-1.33.2


class ServerInterceptor(grpc.ServerInterceptor, metaclass=abc.ABCMeta):
    """Base class for server-side interceptors.

    To implement an interceptor, subclass this class and override the intercept method.
    """

    @abc.abstractmethod
    def intercept(
        self,
        method: Callable,
        request_or_iterator: Any,
        context: grpc.ServicerContext,
        method_name: str,
    ) -> Any:  # pragma: no cover
        """Override this method to implement a custom interceptor.

        You should call method(request_or_iterator, context) to invoke the next handler
        (either the RPC method implementation, or the next interceptor in the list).

        Args:
            method: Either the RPC method implementation, or the next interceptor in
                the chain.
            request_or_iterator: The RPC request, as a protobuf message if it is a
                unary request, or an iterator of protobuf messages if it is a streaming
                request.
            context: The ServicerContext pass by gRPC to the service.
            method_name: A string of the form "/protobuf.package.Service/Method"

        Returns:
            This should return the result of method(request, context), which
            is typically the RPC method response, as a protobuf message, or an
            iterator of protobuf messages for streaming responses. The interceptor is
            free to modify this in some way, however.
        """
        return method(request_or_iterator, context)

    # Implementation of grpc.ServerInterceptor, do not override.
    def intercept_service(self, continuation, handler_call_details):
        """Implementation of grpc.ServerInterceptor.

        This is not part of the grpc_interceptor.ServerInterceptor API, but must have
        a public name. Do not override it, unless you know what you're doing.
        """
        next_handler = continuation(handler_call_details)
        # Returns None if the method isn't implemented.
        if next_handler is None:
            return

        handler_factory, next_handler_method = _get_factory_and_method(next_handler)

        def invoke_intercept_method(request_or_iterator, context):
            method_name = handler_call_details.method
            return self.intercept(
                next_handler_method,
                request_or_iterator,
                context,
                method_name,
            )

        return handler_factory(
            invoke_intercept_method,
            request_deserializer=next_handler.request_deserializer,
            response_serializer=next_handler.response_serializer,
        )


class AsyncServerMiddleware(grpc_aio.ServerInterceptor, metaclass=abc.ABCMeta):
    """Base class for asyncio server-side interceptors.

    To implement an interceptor, subclass this class and override the intercept method.
    """

    @abc.abstractmethod
    async def intercept(
        self,
        method: Callable,
        request_or_iterator: Any,
        context: grpc_aio.ServicerContext,
        method_name: str,
    ) -> Any:  # pragma: no cover
        """Override this method to implement a custom interceptor.

        You should await method(request_or_iterator, context) to invoke the next handler
        (either the RPC method implementation, or the next interceptor in the list).

        Args:
            method: Either the RPC method implementation, or the next interceptor in
                the chain.
            request_or_iterator: The RPC request, as a protobuf message if it is a
                unary request, or an iterator of protobuf messages if it is a streaming
                request.
            context: The ServicerContext pass by gRPC to the service.
            method_name: A string of the form "/protobuf.package.Service/Method"

        Returns:
            This should return the result of method(request_or_iterator, context),
            which is typically the RPC method response, as a protobuf message. The
            interceptor is free to modify this in some way, however.
        """
        response_or_iterator = method(request_or_iterator, context)
        if hasattr(response_or_iterator, "__aiter__"):
            return response_or_iterator
        else:
            return await response_or_iterator

    # Implementation of grpc.ServerInterceptor, do not override.
    async def intercept_service(self, continuation, handler_call_details):
        """Implementation of grpc.aio.ServerInterceptor.

        This is not part of the grpc_interceptor.AsyncServerInterceptor API, but must
        have a public name. Do not override it, unless you know what you're doing.
        """
        next_handler = await continuation(handler_call_details)
        if not next_handler:
            return
        handler_factory, next_handler_method = _get_factory_and_method(next_handler)

        if next_handler.response_streaming:

            async def invoke_intercept_method(request, context):
                method_name = handler_call_details.method

                coroutine_or_asyncgen = self.intercept(
                    next_handler_method,
                    request,
                    context,
                    method_name,
                )

                # Async server streaming handlers return async_generator, because they
                # use the async def + yield syntax. However, this is NOT a coroutine
                # and hence is not awaitable. This can be a problem if the interceptor
                # ignores the individual streaming response items and simply returns the
                # result of method(request, context). In that case the interceptor IS a
                # coroutine, and hence should be awaited. In both cases, we need
                # something we can iterate over so that THIS function is an
                # async_generator like the actual RPC method.
                if iscoroutine(coroutine_or_asyncgen):
                    asyncgen_or_none = await coroutine_or_asyncgen
                    # If a handler is using the read/write API, it will return None.
                    if not asyncgen_or_none:
                        return
                    asyncgen = asyncgen_or_none
                else:
                    asyncgen = coroutine_or_asyncgen

                async for r in asyncgen:
                    yield r

        else:

            async def invoke_intercept_method(request, context):
                method_name = handler_call_details.method
                return await self.intercept(
                    next_handler_method,
                    request,
                    context,
                    method_name,
                )

        return handler_factory(
            invoke_intercept_method,
            request_deserializer=next_handler.request_deserializer,
            response_serializer=next_handler.response_serializer,
        )

    async def call(self, method, request, context):
        """Call method in an async way"""
        response = method(request, context)
        if hasattr(response, "__aiter__"):
            return response
        else:
            return await response


def _get_factory_and_method(
    rpc_handler: grpc.RpcMethodHandler,
) -> Tuple[Callable, Callable]:
    if rpc_handler.unary_unary:
        return grpc.unary_unary_rpc_method_handler, rpc_handler.unary_unary
    elif rpc_handler.unary_stream:
        return grpc.unary_stream_rpc_method_handler, rpc_handler.unary_stream
    elif rpc_handler.stream_unary:
        return grpc.stream_unary_rpc_method_handler, rpc_handler.stream_unary
    elif rpc_handler.stream_stream:
        return grpc.stream_stream_rpc_method_handler, rpc_handler.stream_stream
    else:  # pragma: no cover
        raise RuntimeError("RPC handler implementation does not exist")


class MethodName:
    """Represents a gRPC method name.

    gRPC methods are defined by three parts, represented by the three attributes.

    Attributes:
        package: This is defined by the `package foo.bar;` designation in the protocol
            buffer definition, or it could be defined by the protocol buffer directory
            structure, depending on the language
            (see https://developers.google.com/protocol-buffers/docs/proto3#packages).
        service: This is the service name in the protocol buffer definition (e.g.,
            `service SearchService { ... }`.
        method: This is the method name. (e.g., `rpc Search(...) returns (...);`).
    """

    def __init__(self, package: str, service: str, method: str):
        self.package = package
        self.service = service
        self.method = method

    def __repr__(self) -> str:
        """Object-like representation."""
        return f"MethodName(package='{self.package}', service='{self.service}'," f" method='{self.method}')"

    @property
    def fully_qualified_service(self):
        """Return the service name prefixed with the package.

        Example:
            >>> MethodName("foo.bar", "SearchService", "Search").fully_qualified_service
            'foo.bar.SearchService'
        """
        return f"{self.package}.{self.service}" if self.package else self.service


def parse_method_name(method_name: str) -> MethodName:
    """Parse a method name into package, service and endpoint components.

    Arguments:
        method_name: A string of the form "/foo.bar.SearchService/Search", as passed to
            ServerInterceptor.intercept().

    Returns:
        A MethodName object.

    Example:
        >>> parse_method_name("/foo.bar.SearchService/Search")
        MethodName(package='foo.bar', service='SearchService', method='Search')
    """
    _, package_and_service, method = method_name.split("/")
    *maybe_package, service = package_and_service.rsplit(".", maxsplit=1)
    package = maybe_package[0] if maybe_package else ""
    return MethodName(package, service, method)
