from unittest.mock import MagicMock

import pytest

from metaflow_extensions.observability.plugins.observer import MetricsObserver
from metaflow_extensions.observability.plugins.samplers.cpu import CpuMemoryStats
from metaflow_extensions.observability.plugins.samplers.disk import DiskStats
from metaflow_extensions.observability.plugins.samplers.gpu import GpuDeviceStats, GpuStats


def _make_cpu_stats(**overrides):
    defaults = dict(
        avg_cpu_percent=50.0,
        max_cpu_percent=90.0,
        p95_cpu_percent=85.0,
        avg_memory_mb=256.0,
        max_memory_mb=300.0,
    )
    return CpuMemoryStats(**{**defaults, **overrides})


def _make_disk_stats(**overrides):
    defaults = dict(
        read_bytes=10_000_000,
        write_bytes=5_000_000,
        avg_read_throughput_mbps=10.0,
        avg_write_throughput_mbps=5.0,
    )
    return DiskStats(**{**defaults, **overrides})


def _make_gpu_stats(devices=None):
    if devices is None:
        devices = []
    return GpuStats(devices=devices)


def _make_observer(cpu_stats=None, disk_stats=None, gpu_stats=None, exporter=None):
    cpu_sampler = MagicMock()
    cpu_sampler.stop.return_value = cpu_stats or _make_cpu_stats()

    disk_sampler = MagicMock()
    disk_sampler.stop.return_value = disk_stats or _make_disk_stats()

    gpu_sampler = MagicMock()
    gpu_sampler.stop.return_value = gpu_stats or _make_gpu_stats()

    if exporter is None:
        exporter = MagicMock()

    observer = MetricsObserver(
        cpu_sampler=cpu_sampler,
        disk_sampler=disk_sampler,
        gpu_sampler=gpu_sampler,
        exporter=exporter,
    )
    return observer, cpu_sampler, disk_sampler, gpu_sampler, exporter


_BASE_ATTRS = {"step_name": "train", "flow_name": "MyFlow", "run_id": "123"}


class TestMetricsObserver:
    def test_starts_all_samplers_on_step_start(self):
        observer, cpu, disk, gpu, _ = _make_observer()

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)

        cpu.start.assert_called_once()
        disk.start.assert_called_once()
        gpu.start.assert_called_once()

    def test_stops_all_samplers_on_step_end(self):
        observer, cpu, disk, gpu, _ = _make_observer()

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)
        observer.on_step_end(**_BASE_ATTRS, retry_count=0, success=True)

        cpu.stop.assert_called_once()
        disk.stop.assert_called_once()
        gpu.stop.assert_called_once()

    def test_exports_step_duration(self):
        observer, _, _, _, exporter = _make_observer()

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)
        observer.on_step_end(**_BASE_ATTRS, retry_count=0, success=True)

        calls = [c for c in exporter.record_duration.call_args_list]
        assert len(calls) == 1
        name, seconds, attrs = calls[0].args
        assert name == "step.duration"
        assert seconds >= 0.0
        assert attrs["step_name"] == "train"

    def test_exports_cpu_stats(self):
        cpu = _make_cpu_stats(avg_cpu_percent=55.0, max_cpu_percent=95.0, p95_cpu_percent=88.0)
        observer, _, _, _, exporter = _make_observer(cpu_stats=cpu)

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)
        observer.on_step_end(**_BASE_ATTRS, retry_count=0, success=True)

        gauge_calls = {c.args[0]: c.args[1] for c in exporter.record_gauge.call_args_list}
        assert gauge_calls["step.cpu.avg_percent"] == pytest.approx(55.0)
        assert gauge_calls["step.cpu.max_percent"] == pytest.approx(95.0)
        assert gauge_calls["step.cpu.p95_percent"] == pytest.approx(88.0)

    def test_exports_memory_stats(self):
        cpu = _make_cpu_stats(avg_memory_mb=512.0, max_memory_mb=600.0)
        observer, _, _, _, exporter = _make_observer(cpu_stats=cpu)

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)
        observer.on_step_end(**_BASE_ATTRS, retry_count=0, success=True)

        gauge_calls = {c.args[0]: c.args[1] for c in exporter.record_gauge.call_args_list}
        assert gauge_calls["step.memory.avg_mb"] == pytest.approx(512.0)
        assert gauge_calls["step.memory.max_mb"] == pytest.approx(600.0)

    def test_exports_disk_stats(self):
        disk = _make_disk_stats(
            read_bytes=20_000_000,
            write_bytes=10_000_000,
            avg_read_throughput_mbps=20.0,
            avg_write_throughput_mbps=10.0,
        )
        observer, _, _, _, exporter = _make_observer(disk_stats=disk)

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)
        observer.on_step_end(**_BASE_ATTRS, retry_count=0, success=True)

        counter_calls = {c.args[0]: c.args[1] for c in exporter.record_counter.call_args_list}
        gauge_calls = {c.args[0]: c.args[1] for c in exporter.record_gauge.call_args_list}

        assert counter_calls["step.disk.read_bytes"] == 20_000_000
        assert counter_calls["step.disk.write_bytes"] == 10_000_000
        assert gauge_calls["step.disk.read_throughput_mbps"] == pytest.approx(20.0)
        assert gauge_calls["step.disk.write_throughput_mbps"] == pytest.approx(10.0)

    def test_exports_gpu_stats_per_device(self):
        gpu = _make_gpu_stats(
            devices=[
                GpuDeviceStats(
                    device_index=0,
                    avg_utilization_percent=80.0,
                    max_utilization_percent=95.0,
                    avg_memory_used_mb=3000.0,
                    max_memory_used_mb=3500.0,
                )
            ]
        )
        observer, _, _, _, exporter = _make_observer(gpu_stats=gpu)

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)
        observer.on_step_end(**_BASE_ATTRS, retry_count=0, success=True)

        gauge_calls = exporter.record_gauge.call_args_list
        gpu_gauges = [c for c in gauge_calls if "gpu" in c.args[0]]
        assert len(gpu_gauges) > 0

        by_name = {c.args[0]: (c.args[1], c.args[2]) for c in gpu_gauges}
        assert by_name["step.gpu.avg_utilization_percent"][0] == pytest.approx(80.0)
        assert by_name["step.gpu.avg_utilization_percent"][1]["gpu_index"] == "0"

    def test_skips_gpu_export_when_no_devices(self):
        gpu = _make_gpu_stats(devices=[])
        observer, _, _, _, exporter = _make_observer(gpu_stats=gpu)

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)
        observer.on_step_end(**_BASE_ATTRS, retry_count=0, success=True)

        gpu_calls = [c for c in exporter.record_gauge.call_args_list if "gpu" in c.args[0]]
        assert gpu_calls == []

    def test_increments_failure_counter_on_failure(self):
        observer, _, _, _, exporter = _make_observer()

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)
        observer.on_step_end(**_BASE_ATTRS, retry_count=0, success=False)

        counter_calls = {c.args[0]: c.args[1] for c in exporter.record_counter.call_args_list}
        assert counter_calls["step.failures"] == 1

    def test_does_not_increment_failure_counter_on_success(self):
        observer, _, _, _, exporter = _make_observer()

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)
        observer.on_step_end(**_BASE_ATTRS, retry_count=0, success=True)

        counter_calls = {c.args[0] for c in exporter.record_counter.call_args_list}
        assert "step.failures" not in counter_calls

    def test_exports_retry_count(self):
        observer, _, _, _, exporter = _make_observer()

        observer.on_step_start(**_BASE_ATTRS, retry_count=2)
        observer.on_step_end(**_BASE_ATTRS, retry_count=2, success=True)

        gauge_calls = {c.args[0]: c.args[1] for c in exporter.record_gauge.call_args_list}
        assert gauge_calls["step.retry_count"] == 2

    def test_all_exported_metrics_carry_step_and_flow_attributes(self):
        observer, _, _, _, exporter = _make_observer()

        observer.on_step_start(**_BASE_ATTRS, retry_count=0)
        observer.on_step_end(**_BASE_ATTRS, retry_count=0, success=True)

        all_calls = (
            exporter.record_duration.call_args_list
            + exporter.record_gauge.call_args_list
            + exporter.record_counter.call_args_list
        )
        for c in all_calls:
            attrs = c.args[2]
            assert attrs["step_name"] == "train"
            assert attrs["flow_name"] == "MyFlow"
            assert attrs["run_id"] == "123"
