import collections
from timeit import default_timer
from typing import Iterable, AsyncIterable

import grpc

UNARY = "UNARY"
SERVER_STREAMING = "SERVER_STREAMING"
CLIENT_STREAMING = "CLIENT_STREAMING"
BIDI_STREAMING = "BIDI_STREAMING"
UNKNOWN = "UNKNOWN"


async def wrap_iterator_inc_counter(
    iterator,
    msg_counter,
    grpc_type,
    grpc_service_name,
    grpc_method_name,
    handled_metric=None,
    latency_metric=None,
    start_time=None,
):
    """Wraps an iterator and collect metrics."""

    def record_completion(exception=None):
        if handled_metric:
            if exception:
                if isinstance(exception, grpc.RpcError):
                    code_name = exception.code().name
                else:
                    code_name = "UNKNOWN"
            else:
                code_name = "OK"

            handled_metric.labels(
                grpc_type=grpc_type, grpc_service=grpc_service_name, grpc_method=grpc_method_name, grpc_code=code_name
            ).inc()

        if latency_metric and start_time:
            duration = max(default_timer() - start_time, 0)
            latency_metric.observe(duration)

    try:
        if isinstance(iterator, (AsyncIterable, collections.abc.AsyncIterable)):
            async for item in iterator:
                msg_counter.labels(
                    grpc_type=grpc_type, grpc_service=grpc_service_name, grpc_method=grpc_method_name
                ).inc()
                yield item
        elif isinstance(iterator, (Iterable, collections.abc.Iterable)):
            for item in iterator:
                msg_counter.labels(
                    grpc_type=grpc_type, grpc_service=grpc_service_name, grpc_method=grpc_method_name
                ).inc()
                yield item
        else:
            pass
        record_completion(None)

    except Exception as e:
        record_completion(e)
        raise e


def sync_wrap_iterator_inc_counter(
    iterator,
    msg_counter,
    grpc_type,
    grpc_service_name,
    grpc_method_name,
    handled_metric=None,  # NEW: Метрика статусов
    latency_metric=None,  # NEW: Метрика времени
    start_time=None,  # NEW: Время старта
):
    """Wraps an iterator and collect metrics."""

    def record_completion(exception=None):
        if handled_metric:
            if exception:
                if isinstance(exception, grpc.RpcError):
                    code_name = exception.code().name
                else:
                    code_name = "UNKNOWN"
            else:
                code_name = "OK"

            handled_metric.labels(
                grpc_type=grpc_type, grpc_service=grpc_service_name, grpc_method=grpc_method_name, grpc_code=code_name
            ).inc()

        if latency_metric and start_time:
            duration = max(default_timer() - start_time, 0)
            latency_metric.observe(duration)

    try:
        for item in iterator:
            msg_counter.labels(grpc_type=grpc_type, grpc_service=grpc_service_name, grpc_method=grpc_method_name).inc()
            yield item
        record_completion(None)

    except Exception as e:
        record_completion(e)
        raise e


def get_method_type(request_streaming, response_streaming):
    """
    Infers the method type from if the request or the response is streaming.

    # The Method type is coming from:
    # https://grpc.io/grpc-java/javadoc/io/grpc/MethodDescriptor.MethodType.html
    """
    if request_streaming and response_streaming:
        return BIDI_STREAMING
    elif request_streaming and not response_streaming:
        return CLIENT_STREAMING
    elif not request_streaming and response_streaming:
        return SERVER_STREAMING
    return UNARY


def split_method_call(handler_call_details):
    """
    Infers the grpc service and method name from the handler_call_details.
    """

    # e.g. /package.ServiceName/MethodName
    parts = (
        handler_call_details.method.split("/")
        if not isinstance(handler_call_details.method, bytes)
        else handler_call_details.method.decode("utf-8").split("/")
    )
    if len(parts) < 3:
        return "", "", False

    grpc_service_name, grpc_method_name = parts[1:3]
    return grpc_service_name, grpc_method_name, True
