from timeit import default_timer
from typing import Sequence

import grpc
from prometheus_client.registry import REGISTRY, CollectorRegistry

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


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

    def intercept_service(self, continuation, handler_call_details):
        grpc_service_name, grpc_method_name, ok = _utils.split_method_call(handler_call_details)

        if not ok:
            grpc_service_name, grpc_method_name = "unknown", "unknown"

        def metrics_wrapper(behavior, request_streaming, response_streaming):
            def new_behavior(request_or_iterator, servicer_context):
                grpc_type = _utils.get_method_type(request_streaming, response_streaming)
                self._metrics.grpc_server_started_counter.labels(
                    grpc_type=grpc_type,
                    grpc_service=grpc_service_name,
                    grpc_method=grpc_method_name,
                ).inc()

                start_time = default_timer()

                if request_streaming:
                    request_or_iterator = _utils.sync_wrap_iterator_inc_counter(
                        request_or_iterator,
                        self._metrics.grpc_server_stream_msg_received,
                        grpc_type,
                        grpc_service_name,
                        grpc_method_name,
                    )

                try:
                    response_or_iterator = behavior(request_or_iterator, servicer_context)
                    if response_streaming:
                        response_or_iterator = _utils.sync_wrap_iterator_inc_counter(
                            response_or_iterator,
                            self._metrics.grpc_server_stream_msg_sent,
                            grpc_type,
                            grpc_service_name,
                            grpc_method_name,
                        )

                        return self._monitor_response_completion(
                            response_or_iterator,
                            servicer_context,
                            grpc_type,
                            grpc_service_name,
                            grpc_method_name,
                            start_time,
                        )
                    else:
                        self._record_completion(
                            servicer_context, grpc_type, grpc_service_name, grpc_method_name, start_time
                        )
                        return response_or_iterator

                except Exception as e:
                    self._record_completion(
                        servicer_context, grpc_type, grpc_service_name, grpc_method_name, start_time, e
                    )
                    raise e

            return new_behavior

        return self._wrap_rpc_behavior(continuation(handler_call_details), metrics_wrapper)

    def _monitor_response_completion(self, iterator, context, grpc_type, service, method, start_time):
        """
        Генератор, который крутит переданный итератор.
        Он нужен, чтобы поймать момент окончания стрима (успех или ошибку) и записать Latency/Status.
        """
        try:
            for item in iterator:
                yield item

            self._record_completion(context, grpc_type, service, method, start_time)

        except Exception as e:
            self._record_completion(context, grpc_type, service, method, start_time, e)
            raise e

    def _record_completion(self, context, grpc_type, service, method, start_time, exception=None):
        """Записывает handled_total (Status Code) и Histogram (Latency)."""

        if exception:
            if isinstance(exception, grpc.RpcError):
                code = exception.code()
                code_name = code.name if code else _utils.UNKNOWN
            else:
                code_name = _utils.UNKNOWN
        else:
            code = context.code()
            code_name = code.name if code else "OK"

        self._metrics.grpc_server_handled_total_counter.labels(
            grpc_type=grpc_type,
            grpc_service=service,
            grpc_method=method,
            grpc_code=code_name,
        ).inc()

        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 _wrap_rpc_behavior(self, handler, fn):
        if handler is None:
            return None

        if handler.request_streaming and handler.response_streaming:
            behavior_fn = handler.stream_stream
            handler_factory = grpc.stream_stream_rpc_method_handler
        elif handler.request_streaming and not handler.response_streaming:
            behavior_fn = handler.stream_unary
            handler_factory = grpc.stream_unary_rpc_method_handler
        elif not handler.request_streaming and handler.response_streaming:
            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, handler.request_streaming, handler.response_streaming),
            request_deserializer=handler.request_deserializer,
            response_serializer=handler.response_serializer,
        )


class PromClientInterceptor(
    grpc.UnaryUnaryClientInterceptor,
    grpc.UnaryStreamClientInterceptor,
    grpc.StreamUnaryClientInterceptor,
    grpc.StreamStreamClientInterceptor,
):
    """
    Intercept gRPC client requests.
    """

    def __init__(
        self,
        enable_client_handling_time_histogram: bool = False,
        enable_client_stream_receive_time_histogram: bool = False,
        enable_client_stream_send_time_histogram: bool = False,
        registry=REGISTRY,
        prefix: str = "",
        metrics: ClientGrpcMetrics | None = None,
        latency_buckets: Sequence[float] = DEFAULT_LATENCY_BUCKETS,
    ):
        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 if metrics else init_client_metrics(registry, prefix, latency_buckets)

    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 = continuation(client_call_details, request)
        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))

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

        return handler

    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 = 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.sync_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

    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.sync_wrap_iterator_inc_counter(
            request_iterator, iterator_metric, grpc_type, grpc_service_name, grpc_method_name
        )

        start = default_timer()
        handler = continuation(client_call_details, request_iterator)

        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))

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

        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))

        return handler

    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()

        response_iterator = continuation(
            client_call_details,
            _utils.sync_wrap_iterator_inc_counter(
                request_iterator,
                self._metrics.grpc_client_stream_msg_sent,
                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
            )

        response_iterator = _utils.sync_wrap_iterator_inc_counter(
            iterator=response_iterator,
            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,
        )

        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.ServerInterceptor:
    """
    Создает интерцептор. Рекомендуемый способ использования.
    """
    return PromServerInterceptor(
        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,
) -> PromClientInterceptor:
    """
    Создает клиентский интерцептор.
    """
    return PromClientInterceptor(
        registry=metric_exporter.get_registry(),
        prefix=metric_exporter.prefix,
        enable_client_handling_time_histogram=enable_client_handling_time_histogram,
        enable_client_stream_receive_time_histogram=enable_client_stream_receive_time_histogram,
        enable_client_stream_send_time_histogram=enable_client_stream_send_time_histogram,
        metrics=metrics,
        latency_buckets=latency_buckets,
    )
