from typing import Any, Optional

from opentelemetry.metrics import Counter, Histogram, MeterProvider
from opentelemetry.sdk.metrics import MeterProvider as SdkMeterProvider
from opentelemetry.sdk.metrics.export import (
    ConsoleMetricExporter,
    MetricReader,
    PeriodicExportingMetricReader,
)

from .protocol import Attributes

_METER_NAME = "metaflow.observability"
_SCHEMA_URL = "https://opentelemetry.io/schemas/1.24.0"


class OtelExporter:
    """Exports Metaflow step metrics via OpenTelemetry.

    By default creates a ``MeterProvider`` with a Prometheus exporter.
    Pass a custom ``meter_provider`` to use any OTel-compatible backend
    (useful for testing with ``InMemoryMetricReader``).
    """

    def __init__(self, meter_provider: Optional[MeterProvider] = None) -> None:
        if meter_provider is None:
            meter_provider = _default_provider()
        self._meter = meter_provider.get_meter(_METER_NAME, schema_url=_SCHEMA_URL)
        self._histograms: dict[str, Histogram] = {}
        self._gauges: dict[str, Any] = {}  # Gauge type varies across OTel SDK versions
        self._counters: dict[str, Counter] = {}

    def record_duration(self, name: str, seconds: float, attributes: Attributes) -> None:
        self._histogram(name).record(seconds, attributes=attributes)

    def record_gauge(self, name: str, value: float, attributes: Attributes) -> None:
        self._gauge(name).set(value, attributes=attributes)

    def record_counter(self, name: str, value: int, attributes: Attributes) -> None:
        self._counter(name).add(value, attributes=attributes)

    # ------------------------------------------------------------------
    # Instrument accessors (create-once, reuse)
    # ------------------------------------------------------------------

    def _histogram(self, name: str) -> Histogram:
        if name not in self._histograms:
            self._histograms[name] = self._meter.create_histogram(name)
        return self._histograms[name]

    def _gauge(self, name: str) -> Any:
        if name not in self._gauges:
            self._gauges[name] = self._meter.create_gauge(name)
        return self._gauges[name]

    def _counter(self, name: str) -> Counter:
        if name not in self._counters:
            self._counters[name] = self._meter.create_counter(name)
        return self._counters[name]


def _default_provider() -> SdkMeterProvider:
    reader: MetricReader
    try:
        from opentelemetry.exporter.prometheus import PrometheusMetricReader

        reader = PrometheusMetricReader()
    except ImportError:
        reader = PeriodicExportingMetricReader(ConsoleMetricExporter())

    return SdkMeterProvider(metric_readers=[reader])
