from __future__ import annotations

from collections.abc import Mapping
from typing import Any

from ._core import RequestMetricsFactory


class RequestMetricsCollector:
    """Prometheus collector that pulls typed samples from RequestMetricsFactory."""

    def __init__(
        self,
        factory: RequestMetricsFactory,
        *,
        base_labels: Mapping[str, str] | None = None,
        registry: Any | None = None,
    ) -> None:
        from prometheus_client import REGISTRY

        self.factory = factory
        self.base_labels = dict(base_labels or {})
        self.registry = registry or REGISTRY
        self.registry.register(self)

    def collect(self):
        from prometheus_client.core import CounterMetricFamily, GaugeMetricFamily

        families: dict[tuple[str, tuple[str, ...], str], object] = {}
        for sample in self.factory.prometheus_samples(self.base_labels):
            label_names = tuple(sorted(sample.labels.keys()))
            key = (sample.name, label_names, sample.metric_type)
            family = families.get(key)
            if family is None:
                if sample.metric_type == "counter":
                    family = CounterMetricFamily(sample.name, sample.name, labels=list(label_names))
                else:
                    family = GaugeMetricFamily(sample.name, sample.name, labels=list(label_names))
                families[key] = family
            label_values = [sample.labels[name] for name in label_names]
            family.add_metric(label_values, sample.value)

        yield from families.values()


def prometheus_strfmt(
    factory: RequestMetricsFactory,
    *,
    base_labels: Mapping[str, str] | None = None,
) -> str:
    """Return Prometheus text exposition from the Rust encoder path."""
    return factory.prometheus_strfmt(dict(base_labels or {}))
