from prometheus_client import (
    CollectorRegistry,
    Counter,
    Gauge,
    Histogram,
    generate_latest,
)

from core.domain.errors import FailureReason
from core.domain.models import NormalizedCDR


class WorkerMetrics:
    def __init__(self) -> None:
        self.registry = CollectorRegistry()

        self._processed = Counter(
            "cdr_processed_total",
            "Total CDRs successfully processed and stored.",
            labelnames=(
                "service_type",
                "radio_access_type",
            ),
            registry=self.registry,
        )

        self._dlq = Counter(
            "cdr_dlq_total",
            "Total CDRs routed to the DLQ.",
            labelnames=(
                "reason",
            ),
            registry=self.registry,
        )

        self._kafka_to_clickhouse_latency = Histogram(
            "cdr_kafka_to_clickhouse_latency_seconds",
            "Time from Kafka message creation to ClickHouse write.",
            buckets=(
                0.05,
                0.1,
                0.25,
                0.5,
                1,
                2.5,
                5,
            ),
            registry=self.registry,
        )

        self._batch_flush_duration = Histogram(
            "cdr_batch_flush_duration_seconds",
            "Duration of ClickHouse batch flushes.",
            registry=self.registry,
        )

        self._batch_size = Histogram(
            "cdr_batch_size_at_flush",
            "Number of CDRs in each flushed batch.",
            buckets=(
                10,
                50,
                100,
                250,
                500,
                1000,
            ),
            registry=self.registry,
        )

        self._buffer_size = Gauge(
            "cdr_in_flight_buffer_size",
            "Current number of CDRs buffered for ClickHouse.",
            registry=self.registry,
        )

        self._consumer_lag = Gauge(
            "cdr_kafka_consumer_lag",
            "Kafka consumer lag by partition.",
            labelnames=(
                "partition",
            ),
            registry=self.registry,
        )

        # Ensure all required DLQ reasons appear even before
        # their first failure.
        for reason in FailureReason:
            self._dlq.labels(
                reason=reason.value
            )

    def record_processed(
        self,
        cdr: NormalizedCDR,
    ) -> None:
        self._processed.labels(
            service_type=cdr.service_type,
            radio_access_type=cdr.radio_access_type,
        ).inc()

    def record_dlq(
        self,
        reason: str,
    ) -> None:
        self._dlq.labels(
            reason=reason
        ).inc()

    def observe_kafka_to_clickhouse_latency(
        self,
        seconds: float,
    ) -> None:
        self._kafka_to_clickhouse_latency.observe(
            max(seconds, 0.0)
        )

    def observe_batch_flush_duration(
        self,
        seconds: float,
    ) -> None:
        self._batch_flush_duration.observe(
            seconds
        )

    def observe_batch_size(
        self,
        size: int,
    ) -> None:
        self._batch_size.observe(
            size
        )

    def set_buffer_size(
        self,
        size: int,
    ) -> None:
        self._buffer_size.set(
            size
        )

    def set_consumer_lag(
        self,
        partition: int,
        lag: int,
    ) -> None:
        self._consumer_lag.labels(
            partition=str(partition)
        ).set(
            max(lag, 0)
        )

    def render(self) -> bytes:
        return generate_latest(
            self.registry
        )