import inspect
from timeit import default_timer
from typing import Callable, Awaitable, Sequence

import grpc
from grpc.aio import (
    ServerInterceptor,
    UnaryUnaryClientInterceptor,
    UnaryStreamClientInterceptor,
    StreamUnaryClientInterceptor,
    StreamStreamClientInterceptor,
    ClientInterceptor,
)

from prometheus_client.registry import REGISTRY

from . import _utils
from .metrics import (
    init_server_metrics,
    init_client_metrics,
    ServerGrpcMetrics,
    DEFAULT_LATENCY_BUCKETS,
    ClientGrpcMetrics,
)

from ... import MetricsExporter


# TODO потребует в будущем рефактор DRY
class PromAsyncServerInterceptor(ServerInterceptor):
    def __init__(
        self,
        enable_handling_time_histogram=False,
        registry=REGISTRY,
        prefix: str = "",
        metrics: ServerGrpcMetrics | None = None,
        latency_buckets: Sequence[float] = DEFAULT_LATENCY_BUCKETS,
    ):
        self._metrics = metrics if metrics else init_server_metrics(registry, prefix, latency_buckets)
        self._enable_handling_time_histogram = enable_handling_time_histogram
        self._code_to_status_mapping = {x.value[0]: x for x in grpc.StatusCode}

    async def intercept_service(
        self,
        continuation: Callable[[grpc.HandlerCallDetails], Awaitable[grpc.RpcMethodHandler]],
        handler_call_details: grpc.HandlerCallDetails,
    ) -> grpc.RpcMethodHandler:

        handler = await continuation(handler_call_details)

        if not handler:
            return handler

        grpc_service_name, grpc_method_name, _ = _utils.split_method_call(handler_call_details)

        wrapper_fn = self._create_metrics_wrapper(grpc_service_name, grpc_method_name)

        return self._wrap_rpc_behavior(handler, wrapper_fn)

    def _create_metrics_wrapper(self, service_name: str, method_name: str):
        """Фабрика для создания функции-обертки с замкнутым контекстом имени метода."""

        def metrics_wrapper(behavior, request_streaming, response_streaming):
            grpc_type = _utils.get_method_type(request_streaming, response_streaming)

            if response_streaming:
                return self._stream_behavior(behavior, request_streaming, grpc_type, service_name, method_name)
            else:
                return self._unary_behavior(behavior, request_streaming, grpc_type, service_name, method_name)

        return metrics_wrapper

    def _unary_behavior(self, behavior, request_streaming, grpc_type, service, method):
        """Обработчик для Unary-ответов (Unary-Unary, Stream-Unary)."""

        async def implementation(request_or_iterator, servicer_context):
            start = default_timer()
            self._inc_started_counter(grpc_type, service, method)

            try:
                actual_request = self._prepare_request(
                    request_or_iterator, request_streaming, grpc_type, service, method
                )

                response = behavior(actual_request, servicer_context)
                if inspect.iscoroutine(response):
                    response = await response

                self._record_result(grpc_type, service, method, servicer_context)
                return response

            except Exception as e:
                self._record_exception(grpc_type, service, method, e)
                raise e
            finally:
                self._observe_histogram(grpc_type, service, method, start)

        return implementation

    def _stream_behavior(self, behavior, request_streaming, grpc_type, service, method):
        """Обработчик для Streaming-ответов (Unary-Stream, Stream-Stream)."""

        async def implementation(request_or_iterator, servicer_context):
            start = default_timer()
            self._inc_started_counter(grpc_type, service, method)

            try:
                actual_request = self._prepare_request(
                    request_or_iterator, request_streaming, grpc_type, service, method
                )

                response_iterator = behavior(actual_request, servicer_context)
                if inspect.iscoroutine(response_iterator):
                    response_iterator = await response_iterator

                response_iterator = self._prepare_response_stream(response_iterator, grpc_type, service, method)

                async for item in response_iterator:
                    yield item

                self._record_result(grpc_type, service, method, servicer_context)

            except Exception as e:
                self._record_exception(grpc_type, service, method, e)
                raise e
            finally:
                self._observe_histogram(grpc_type, service, method, start)

        return implementation

    def _prepare_request(self, request, is_streaming, grpc_type, service, method):
        """Оборачивает входящий поток запросов для подсчета метрик, если необходимо."""
        if is_streaming:
            return _utils.wrap_iterator_inc_counter(
                request,
                self._metrics.grpc_server_stream_msg_received,
                grpc_type,
                service,
                method,
            )
        return request

    def _prepare_response_stream(self, iterator, grpc_type, service, method):
        """Оборачивает исходящий поток ответов для подсчета метрик."""
        return _utils.wrap_iterator_inc_counter(
            iterator,
            self._metrics.grpc_server_stream_msg_sent,
            grpc_type,
            service,
            method,
        )

    def _inc_started_counter(self, grpc_type, service, method):
        self._metrics.grpc_server_started_counter.labels(
            grpc_type=grpc_type,
            grpc_service=service,
            grpc_method=method,
        ).inc()

    def _record_result(self, grpc_type, service, method, context):
        """Запись метрики успешной (или отмененной клиентом) обработки."""
        code = self._compute_status_code(context).name
        self.increase_grpc_server_handled_total_counter(grpc_type, service, method, code)

    def _record_exception(self, grpc_type, service, method, exc):
        """Запись метрики при возникновении исключения."""
        if isinstance(exc, grpc.RpcError):
            code = self._compute_error_code(exc).name
        else:
            code = grpc.StatusCode.UNKNOWN.name

        self.increase_grpc_server_handled_total_counter(grpc_type, service, method, code)

    def _observe_histogram(self, grpc_type, service, method, start_time):
        """Запись времени выполнения."""
        if self._enable_handling_time_histogram:
            elapsed = max(default_timer() - start_time, 0)
            self._metrics.grpc_server_handled_histogram.labels(
                grpc_type=grpc_type,
                grpc_service=service,
                grpc_method=method,
            ).observe(elapsed)

    def _compute_status_code(self, servicer_context):
        if servicer_context.cancelled():
            return grpc.StatusCode.CANCELLED
        if servicer_context.code() is None:
            return grpc.StatusCode.OK
        return self._code_to_status_mapping.get(servicer_context.code(), grpc.StatusCode.UNKNOWN)

    def _compute_error_code(self, grpc_exception):
        if isinstance(grpc_exception, grpc.RpcError):
            return grpc_exception.code()
        return grpc.StatusCode.UNKNOWN

    def increase_grpc_server_handled_total_counter(self, grpc_type, grpc_service_name, grpc_method_name, grpc_code):
        self._metrics.grpc_server_handled_total_counter.labels(
            grpc_type=grpc_type,
            grpc_service=grpc_service_name,
            grpc_method=grpc_method_name,
            grpc_code=grpc_code,
        ).inc()

    def _wrap_rpc_behavior(self, handler, fn):
        if handler is None:
            return None

        is_req_stream = handler.request_streaming
        is_res_stream = handler.response_streaming

        if is_req_stream and is_res_stream:
            behavior_fn = handler.stream_stream
            handler_factory = grpc.stream_stream_rpc_method_handler
        elif is_req_stream:
            behavior_fn = handler.stream_unary
            handler_factory = grpc.stream_unary_rpc_method_handler
        elif is_res_stream:
            behavior_fn = handler.unary_stream
            handler_factory = grpc.unary_stream_rpc_method_handler
        else:
            behavior_fn = handler.unary_unary
            handler_factory = grpc.unary_unary_rpc_method_handler

        return handler_factory(
            fn(behavior_fn, is_req_stream, is_res_stream),
            request_deserializer=handler.request_deserializer,
            response_serializer=handler.response_serializer,
        )


class _PromAsyncUnaryUnaryClientInterceptor(UnaryUnaryClientInterceptor):
    def __init__(
        self,
        metrics: ClientGrpcMetrics,
        enable_client_handling_time_histogram: bool = False,
        enable_client_stream_receive_time_histogram: bool = False,
        enable_client_stream_send_time_histogram: bool = False,
    ):
        self._enable_client_handling_time_histogram = enable_client_handling_time_histogram
        self._enable_client_stream_receive_time_histogram = enable_client_stream_receive_time_histogram
        self._enable_client_stream_send_time_histogram = enable_client_stream_send_time_histogram
        self._metrics = metrics

    async def intercept_unary_unary(self, continuation, client_call_details, request):
        grpc_service_name, grpc_method_name, _ = _utils.split_method_call(client_call_details)
        grpc_type = _utils.UNARY

        self._metrics.grpc_client_started_counter.labels(
            grpc_type=grpc_type, grpc_service=grpc_service_name, grpc_method=grpc_method_name
        ).inc()

        start = default_timer()
        handler = await continuation(client_call_details, request)

        await handler

        if self._enable_client_handling_time_histogram:
            self._metrics.grpc_client_handled_histogram.labels(
                grpc_type=grpc_type,
                grpc_service=grpc_service_name,
                grpc_method=grpc_method_name,
            ).observe(max(default_timer() - start, 0))

        code = await handler.code()

        self._metrics.grpc_client_handled_counter.labels(
            grpc_type=grpc_type, grpc_service=grpc_service_name, grpc_method=grpc_method_name, grpc_code=code.name
        ).inc()

        return handler


class _PromAsyncUnaryStreamClientInterceptor(UnaryStreamClientInterceptor):
    def __init__(
        self,
        metrics: ClientGrpcMetrics,
        enable_client_handling_time_histogram: bool = False,
        enable_client_stream_receive_time_histogram: bool = False,
        enable_client_stream_send_time_histogram: bool = False,
    ):
        self._enable_client_handling_time_histogram = enable_client_handling_time_histogram
        self._enable_client_stream_receive_time_histogram = enable_client_stream_receive_time_histogram
        self._enable_client_stream_send_time_histogram = enable_client_stream_send_time_histogram
        self._metrics = metrics

    async def intercept_unary_stream(self, continuation, client_call_details, request):
        grpc_service_name, grpc_method_name, _ = _utils.split_method_call(client_call_details)
        grpc_type = _utils.SERVER_STREAMING

        self._metrics.grpc_client_started_counter.labels(
            grpc_type=grpc_type,
            grpc_service=grpc_service_name,
            grpc_method=grpc_method_name,
        ).inc()

        start = default_timer()
        handler = await continuation(client_call_details, request)

        latency_metric = None
        if self._enable_client_handling_time_histogram:
            latency_metric = self._metrics.grpc_client_handled_histogram.labels(
                grpc_type=grpc_type,
                grpc_service=grpc_service_name,
                grpc_method=grpc_method_name,
            )

        handler = _utils.wrap_iterator_inc_counter(
            iterator=handler,
            msg_counter=self._metrics.grpc_client_stream_msg_received,
            grpc_type=grpc_type,
            grpc_service_name=grpc_service_name,
            grpc_method_name=grpc_method_name,
            handled_metric=self._metrics.grpc_client_handled_counter,
            latency_metric=latency_metric,
            start_time=start,
        )
        if self._enable_client_stream_receive_time_histogram:
            self._metrics.grpc_client_stream_recv_histogram.labels(
                grpc_type=grpc_type,
                grpc_service=grpc_service_name,
                grpc_method=grpc_method_name,
            ).observe(max(default_timer() - start, 0))

        return handler


class _PromAsyncStreamUnaryClientInterceptor(StreamUnaryClientInterceptor):
    def __init__(
        self,
        metrics: ClientGrpcMetrics,
        enable_client_handling_time_histogram: bool = False,
        enable_client_stream_receive_time_histogram: bool = False,
        enable_client_stream_send_time_histogram: bool = False,
    ):
        self._enable_client_handling_time_histogram = enable_client_handling_time_histogram
        self._enable_client_stream_receive_time_histogram = enable_client_stream_receive_time_histogram
        self._enable_client_stream_send_time_histogram = enable_client_stream_send_time_histogram
        self._metrics = metrics

    async def intercept_stream_unary(self, continuation, client_call_details, request_iterator):
        grpc_service_name, grpc_method_name, _ = _utils.split_method_call(client_call_details)
        grpc_type = _utils.CLIENT_STREAMING
        self._metrics.grpc_client_started_counter.labels(
            grpc_type=grpc_type,
            grpc_service=grpc_service_name,
            grpc_method=grpc_method_name,
        ).inc()
        iterator_metric = self._metrics.grpc_client_stream_msg_sent

        request_iterator = _utils.wrap_iterator_inc_counter(
            request_iterator,
            iterator_metric,
            grpc_type,
            grpc_service_name,
            grpc_method_name,
        )
        start = default_timer()
        handler = await continuation(client_call_details, request_iterator)

        await handler

        if self._enable_client_handling_time_histogram:
            self._metrics.grpc_client_handled_histogram.labels(
                grpc_type=grpc_type,
                grpc_service=grpc_service_name,
                grpc_method=grpc_method_name,
            ).observe(max(default_timer() - start, 0))

        if self._enable_client_stream_send_time_histogram:
            self._metrics.grpc_client_stream_send_histogram.labels(
                grpc_type=grpc_type,
                grpc_service=grpc_service_name,
                grpc_method=grpc_method_name,
            ).observe(max(default_timer() - start, 0))

        code = await handler.code()
        self._metrics.grpc_client_handled_counter.labels(
            grpc_type=grpc_type, grpc_service=grpc_service_name, grpc_method=grpc_method_name, grpc_code=code.name
        ).inc()

        return handler


class _PromAsyncStreamStreamClientInterceptor(StreamStreamClientInterceptor):
    def __init__(
        self,
        metrics: ClientGrpcMetrics,
        enable_client_handling_time_histogram: bool = False,
        enable_client_stream_receive_time_histogram: bool = False,
        enable_client_stream_send_time_histogram: bool = False,
    ):
        self._enable_client_handling_time_histogram = enable_client_handling_time_histogram
        self._enable_client_stream_receive_time_histogram = enable_client_stream_receive_time_histogram
        self._enable_client_stream_send_time_histogram = enable_client_stream_send_time_histogram
        self._metrics = metrics

    async def intercept_stream_stream(self, continuation, client_call_details, request_iterator):
        grpc_service_name, grpc_method_name, _ = _utils.split_method_call(client_call_details)
        grpc_type = _utils.BIDI_STREAMING
        self._metrics.grpc_client_started_counter.labels(
            grpc_type=grpc_type,
            grpc_service=grpc_service_name,
            grpc_method=grpc_method_name,
        ).inc()

        start = default_timer()

        iterator_sent_metric = self._metrics.grpc_client_stream_msg_sent

        response_iterator = await continuation(
            client_call_details,
            _utils.wrap_iterator_inc_counter(
                request_iterator,
                iterator_sent_metric,
                grpc_type,
                grpc_service_name,
                grpc_method_name,
            ),
        )

        latency_metric = None
        if self._enable_client_handling_time_histogram:
            latency_metric = self._metrics.grpc_client_handled_histogram.labels(
                grpc_type=grpc_type,
                grpc_service=grpc_service_name,
                grpc_method=grpc_method_name,
            )

        if self._enable_client_stream_send_time_histogram:
            self._metrics.grpc_client_stream_send_histogram.labels(
                grpc_type=grpc_type,
                grpc_service=grpc_service_name,
                grpc_method=grpc_method_name,
            ).observe(max(default_timer() - start, 0))

        iterator_received_metric = self._metrics.grpc_client_stream_msg_received

        response_iterator = _utils.wrap_iterator_inc_counter(
            iterator=response_iterator,
            msg_counter=iterator_received_metric,
            grpc_type=grpc_type,
            grpc_service_name=grpc_service_name,
            grpc_method_name=grpc_method_name,
            handled_metric=self._metrics.grpc_client_handled_counter,
            latency_metric=latency_metric,
            start_time=start,
        )

        if self._enable_client_stream_receive_time_histogram:
            self._metrics.grpc_client_stream_recv_histogram.labels(
                grpc_type=grpc_type,
                grpc_service=grpc_service_name,
                grpc_method=grpc_method_name,
            ).observe(max(default_timer() - start, 0))

        return response_iterator


def create_server_metrics_interceptor(
    metric_exporter: MetricsExporter,
    enable_handling_time_histogram: bool = True,
    metrics: ServerGrpcMetrics | None = None,
    latency_buckets: Sequence[float] = DEFAULT_LATENCY_BUCKETS,
) -> grpc.aio.ServerInterceptor:
    """
    Создает интерцептор. Рекомендуемый способ использования.
    """
    return PromAsyncServerInterceptor(
        registry=metric_exporter.get_registry(),
        prefix=metric_exporter.prefix,
        enable_handling_time_histogram=enable_handling_time_histogram,
        metrics=metrics,
        latency_buckets=latency_buckets,
    )


def create_client_metrics_interceptor(
    metric_exporter: MetricsExporter,
    enable_client_handling_time_histogram: bool = True,
    enable_client_stream_receive_time_histogram: bool = True,
    enable_client_stream_send_time_histogram: bool = True,
    metrics: ClientGrpcMetrics | None = None,
    latency_buckets: Sequence[float] = DEFAULT_LATENCY_BUCKETS,
) -> list[ClientInterceptor]:
    """
    Создает клиентский интерцептор.
    """
    if metrics is None:
        metrics = init_client_metrics(metric_exporter.get_registry(), metric_exporter.prefix, latency_buckets)

    return [
        _PromAsyncUnaryUnaryClientInterceptor(
            metrics,
            enable_client_handling_time_histogram,
            enable_client_stream_receive_time_histogram,
            enable_client_stream_send_time_histogram,
        ),
        _PromAsyncUnaryStreamClientInterceptor(
            metrics,
            enable_client_handling_time_histogram,
            enable_client_stream_receive_time_histogram,
            enable_client_stream_send_time_histogram,
        ),
        _PromAsyncStreamUnaryClientInterceptor(
            metrics,
            enable_client_handling_time_histogram,
            enable_client_stream_receive_time_histogram,
            enable_client_stream_send_time_histogram,
        ),
        _PromAsyncStreamStreamClientInterceptor(
            metrics,
            enable_client_handling_time_histogram,
            enable_client_stream_receive_time_histogram,
            enable_client_stream_send_time_histogram,
        ),
    ]
