import pytest
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import InMemoryMetricReader

from metaflow_extensions.observability.plugins.export.otel import OtelExporter
from metaflow_extensions.observability.plugins.export.protocol import MetricsExporter


def _make_exporter() -> tuple:
    reader = InMemoryMetricReader()
    provider = MeterProvider(metric_readers=[reader])
    exporter = OtelExporter(meter_provider=provider)
    return exporter, reader


def _collect_metrics(reader: InMemoryMetricReader) -> dict:
    """Return {metric_name: [data_point, ...]} from the reader."""
    result = {}
    for metric in reader.get_metrics_data().resource_metrics:
        for scope in metric.scope_metrics:
            for m in scope.metrics:
                result[m.name] = list(m.data.data_points)
    return result


class TestOtelExporter:
    def test_implements_metrics_exporter_protocol(self):
        exporter, _ = _make_exporter()
        assert isinstance(exporter, MetricsExporter)

    def test_record_duration_creates_histogram(self):
        exporter, reader = _make_exporter()

        exporter.record_duration("step.duration", 3.5, {"step_name": "train"})

        metrics = _collect_metrics(reader)
        assert "step.duration" in metrics
        point = metrics["step.duration"][0]
        assert point.sum == pytest.approx(3.5)

    def test_record_gauge_stores_last_value(self):
        exporter, reader = _make_exporter()

        exporter.record_gauge("step.cpu.avg_percent", 72.5, {"step_name": "train"})

        metrics = _collect_metrics(reader)
        assert "step.cpu.avg_percent" in metrics
        point = metrics["step.cpu.avg_percent"][0]
        assert point.value == pytest.approx(72.5)

    def test_record_counter_accumulates(self):
        exporter, reader = _make_exporter()

        exporter.record_counter("step.failures", 1, {"step_name": "train"})
        exporter.record_counter("step.failures", 1, {"step_name": "train"})

        metrics = _collect_metrics(reader)
        assert "step.failures" in metrics
        point = metrics["step.failures"][0]
        assert point.value == 2

    def test_attributes_attached_to_data_point(self):
        exporter, reader = _make_exporter()

        exporter.record_duration(
            "step.duration",
            1.0,
            {"step_name": "process", "flow_name": "MyFlow"},
        )

        metrics = _collect_metrics(reader)
        point = metrics["step.duration"][0]
        assert point.attributes["step_name"] == "process"
        assert point.attributes["flow_name"] == "MyFlow"

    def test_multiple_gauges_recorded_independently(self):
        exporter, reader = _make_exporter()

        exporter.record_gauge("step.cpu.avg_percent", 40.0, {"step_name": "s"})
        exporter.record_gauge("step.memory.avg_mb", 256.0, {"step_name": "s"})

        metrics = _collect_metrics(reader)
        assert "step.cpu.avg_percent" in metrics
        assert "step.memory.avg_mb" in metrics
