from dataclasses import dataclass
from typing import Sequence

from prometheus_client import Histogram, Counter, CollectorRegistry

from ...utils import add_prefix_to_metric_name

DEFAULT_LATENCY_BUCKETS = (
    0.005,
    0.01,
    0.025,
    0.05,
    0.075,
    0.1,
    0.25,
    0.5,
    0.75,
    1.0,
    2.5,
    5.0,
    7.5,
    10.0,
)


@dataclass(frozen=True, kw_only=True)
class ServerGrpcMetrics:
    grpc_server_started_counter: Counter
    grpc_server_stream_msg_received: Counter
    grpc_server_stream_msg_sent: Counter
    grpc_server_handled_histogram: Histogram
    grpc_server_handled_total_counter: Counter


@dataclass(frozen=True, kw_only=True)
class ClientGrpcMetrics:
    grpc_client_started_counter: Counter
    grpc_client_handled_counter: Counter
    grpc_client_stream_msg_received: Counter
    grpc_client_stream_msg_sent: Counter
    grpc_client_handled_histogram: Histogram
    grpc_client_stream_recv_histogram: Histogram
    grpc_client_stream_send_histogram: Histogram


def init_server_metrics(
    registry: CollectorRegistry,
    prefix: str = "",
    latency_buckets: Sequence[float] = DEFAULT_LATENCY_BUCKETS,
) -> ServerGrpcMetrics:
    return ServerGrpcMetrics(
        grpc_server_started_counter=Counter(
            add_prefix_to_metric_name("grpc_server_started_total", prefix),
            "Total number of RPCs started on the server.",
            ["grpc_type", "grpc_service", "grpc_method"],
            registry=registry,
        ),
        grpc_server_stream_msg_received=Counter(
            add_prefix_to_metric_name("grpc_server_msg_received_total", prefix),
            "Total number of RPC stream messages received on the server.",
            ["grpc_type", "grpc_service", "grpc_method"],
            registry=registry,
        ),
        grpc_server_stream_msg_sent=Counter(
            add_prefix_to_metric_name("grpc_server_msg_sent_total", prefix),
            "Total number of gRPC stream messages sent by the server.",
            ["grpc_type", "grpc_service", "grpc_method"],
            registry=registry,
        ),
        grpc_server_handled_histogram=Histogram(
            add_prefix_to_metric_name("grpc_server_handling_seconds", prefix),
            "Histogram of response latency (seconds) of gRPC that had been application-level handled by the server.",
            ["grpc_type", "grpc_service", "grpc_method"],
            registry=registry,
            buckets=latency_buckets,
        ),
        grpc_server_handled_total_counter=Counter(
            add_prefix_to_metric_name("grpc_server_handled_total", prefix),
            "Total number of RPCs completed on the server, regardless of success or failure.",
            ["grpc_type", "grpc_service", "grpc_method", "grpc_code"],
            registry=registry,
        ),
    )


def init_client_metrics(
    registry: CollectorRegistry,
    prefix: str = "",
    latency_buckets: Sequence[float] = DEFAULT_LATENCY_BUCKETS,
) -> ClientGrpcMetrics:
    return ClientGrpcMetrics(
        grpc_client_started_counter=Counter(
            add_prefix_to_metric_name("grpc_client_started_total", prefix),
            "Total number of RPCs started on the client",
            ["grpc_type", "grpc_service", "grpc_method"],
            registry=registry,
        ),
        grpc_client_handled_counter=Counter(
            add_prefix_to_metric_name("grpc_client_handled_total", prefix),
            "Total number of RPCs completed on the client, regardless of success or failure.",
            ["grpc_type", "grpc_service", "grpc_method", "grpc_code"],
            registry=registry,
        ),
        grpc_client_stream_msg_received=Counter(
            add_prefix_to_metric_name("grpc_client_msg_received_total", prefix),
            "Total number of RPC stream messages received by the client.",
            ["grpc_type", "grpc_service", "grpc_method"],
            registry=registry,
        ),
        grpc_client_stream_msg_sent=Counter(
            add_prefix_to_metric_name("grpc_client_msg_sent_total", prefix),
            "Total number of gRPC stream messages sent by the client.",
            ["grpc_type", "grpc_service", "grpc_method"],
            registry=registry,
        ),
        grpc_client_handled_histogram=Histogram(
            add_prefix_to_metric_name("grpc_client_handling_seconds", prefix),
            "Histogram of response latency (seconds) of the gRPC until it is finished by the application.",
            ["grpc_type", "grpc_service", "grpc_method"],
            registry=registry,
            buckets=latency_buckets,
        ),
        grpc_client_stream_recv_histogram=Histogram(
            add_prefix_to_metric_name("grpc_client_msg_recv_handling_seconds", prefix),
            "Histogram of response latency (seconds) of the gRPC single message receive.",
            ["grpc_type", "grpc_service", "grpc_method"],
            registry=registry,
            buckets=latency_buckets,
        ),
        grpc_client_stream_send_histogram=Histogram(
            add_prefix_to_metric_name("grpc_client_msg_send_handling_seconds", prefix),
            "Histogram of response latency (seconds) of the gRPC single message send.",
            ["grpc_type", "grpc_service", "grpc_method"],
            registry=registry,
            buckets=latency_buckets,
        ),
    )
