"""Middleware for gRPC server"""
import abc
import time
from typing import Any, Callable

import grpc
import ulid

from app import logger

log = logger.logger

CORRELATION_ID_KEY = "x-gitlab-correlation-id"
METRIC_API_CALLS = "spamcheck_api_calls"
METRIC_API_FAILURES = "spamcheck_api_failures"

# Metadata fields to include in log messages, if present
_metadata_log_fields = {"user-agent", "retry"}


# pylint: disable=abstract-method
class SpamCheckContext(grpc.ServicerContext):
    """Custom implementation of a context object passed to method implementations."""

    # there is no init in the super class
    # pylint: disable=super-init-not-called
    def __init__(self, context: grpc.ServicerContext):
        self._original_context = context
        self.correlation_id = ulid.new()

    def abort(self, code, details):
        self._original_context.abort(code, details)

    def abort_with_status(self, status):
        self._original_context.abort_with_status(status)

    def add_callback(self, callback):
        self._original_context.add_callback(callback)

    def auth_context(self):
        return self._original_context.auth_context()

    def cancel(self):
        self._original_context.cancel()

    def invocation_metadata(self):
        return self._original_context.invocation_metadata()

    def is_active(self):
        return self._original_context.is_active()

    def peer(self):
        return self._original_context.peer()

    def peer_identities(self):
        return self._original_context.peer_identities()

    def peer_identity_key(self):
        return self._original_context.peer_identity_key()

    def send_initial_metadata(self, initial_metadata):
        self._original_context.send_initial_metadata(initial_metadata)

    def set_code(self, code):
        self._original_context.set_code(code)

    def set_details(self, details):
        self._original_context.set_details(details)

    def set_trailing_metadata(self, trailing_metadata):
        self._original_context.abort(trailing_metadata)

    def time_remaining(self):
        return self._original_context.time_remaining()


# ServerInterceptor code copied from:
# https://github.com/d5h-foss/grpc-interceptor/blob/master/src/grpc_interceptor/server.py#L9
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: Any,
        context: grpc.ServicerContext,
        method_name: str,
    ) -> Any:
        """Override this method to implement a custom interceptor.
        You should call method(request, 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: The RPC request, as a protobuf message.
            context: The ServicerContext pass by gRPC to the service.
            method_name: A string of the form "/protobuf.package.Service/Method"
        Returns:
            This should generally return the result of method(request, context), which
            is typically the RPC method response, as a protobuf message. The
            interceptor is free to modify this in some way, however.
        """
        return method(request, 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)

        # Only intercept if it's unary:
        if next_handler.request_streaming or next_handler.response_streaming:
            return next_handler

        def invoke_intercept_method(request, context):
            next_interceptor_or_implementation = next_handler.unary_unary
            method_name = handler_call_details.method
            return self.intercept(
                next_interceptor_or_implementation,
                request,
                context,
                method_name,
            )

        return grpc.unary_unary_rpc_method_handler(
            invoke_intercept_method,
            request_deserializer=next_handler.request_deserializer,
            response_serializer=next_handler.response_serializer,
        )


# pylint: disable=too-few-public-methods
class CorrelationIDInterceptor(ServerInterceptor):
    """Add correlation_id to gRPC requests."""

    def intercept(self, method, request, context, method_name):
        ctx = SpamCheckContext(context)
        for item in context.invocation_metadata():
            item = item._asdict()
            if item["key"].lower() == CORRELATION_ID_KEY:
                ctx.correlation_id = item["value"]
                break
        return method(request, ctx)


# pylint: disable=too-few-public-methods
class LoggingInterceptor(ServerInterceptor):
    """Logging interceptor for gRPC requests."""

    def intercept(self, method, request, context, method_name):
        metadata = {
            "correlation_id": context.correlation_id,
            "failed": False,
            "method": method_name,
            "metric": METRIC_API_CALLS,
            "peer": context.peer(),
        }
        for item in context.invocation_metadata():
            item = item._asdict()
            if item["key"] in _metadata_log_fields:
                metadata[item["key"]] = item["value"]

        start_time = time.time()
        try:
            out = method(request, context)
        # Catch and log any unhandled excecptions processing requests
        # pylint: disable=broad-except
        except Exception as ex:
            metadata["failed"] = True
            fields = {
                "metric": METRIC_API_FAILURES,
                "correlation_id": context.correlation_id,
            }
            log.error(str(ex), extra=fields)
            context.set_code(grpc.StatusCode.INTERNAL)
            context.set_details("server error")
            raise
        finally:
            metadata["request_latency_milliseconds"] = (time.time() - start_time) * 1000
            log.info("Processed Request", extra=metadata)
        return out
