import time
from unittest.mock import MagicMock

import pytest

from metaflow_extensions.observability.plugins.samplers.gpu import (
    GpuSampler,
    GpuStats,
)


def _make_nvml(device_count=1, utilization=60, memory_used=2_000_000_000):
    nvml = MagicMock()
    nvml.nvmlDeviceGetCount.return_value = device_count

    def make_handle(index):
        handle = MagicMock()
        util = MagicMock()
        util.gpu = utilization
        nvml.nvmlDeviceGetUtilizationRates.return_value = util

        mem = MagicMock()
        mem.used = memory_used
        nvml.nvmlDeviceGetMemoryInfo.return_value = mem
        return handle

    nvml.nvmlDeviceGetHandleByIndex.side_effect = make_handle
    return nvml


class TestGpuSampler:
    def test_returns_empty_stats_when_pynvml_unavailable(self):
        sampler = GpuSampler(nvml=None)
        sampler.start()
        stats = sampler.stop()

        assert isinstance(stats, GpuStats)
        assert stats.devices == []
        assert not stats.available

    def test_returns_empty_stats_when_no_devices(self):
        nvml = _make_nvml(device_count=0)
        sampler = GpuSampler(nvml=nvml)
        sampler.start()
        stats = sampler.stop()

        assert stats.devices == []

    def test_samples_single_gpu(self):
        nvml = _make_nvml(device_count=1, utilization=80, memory_used=4_000_000_000)

        sampler = GpuSampler(nvml=nvml, interval=0.01)
        sampler.start()
        time.sleep(0.05)
        stats = sampler.stop()

        assert len(stats.devices) == 1
        device = stats.devices[0]
        assert device.device_index == 0
        assert device.avg_utilization_percent == pytest.approx(80.0, abs=1.0)
        assert device.avg_memory_used_mb == pytest.approx(4000.0, abs=1.0)

    def test_samples_multiple_gpus(self):
        nvml = MagicMock()
        nvml.nvmlDeviceGetCount.return_value = 2

        def make_util_for_device(handle):
            util = MagicMock()
            # We use handle index stored on the mock to vary utilization
            util.gpu = handle._device_index * 10 + 50
            return util

        def make_mem_for_device(handle):
            mem = MagicMock()
            mem.used = 1_000_000_000
            return mem

        def get_handle(index):
            h = MagicMock()
            h._device_index = index
            return h

        nvml.nvmlDeviceGetHandleByIndex.side_effect = get_handle
        nvml.nvmlDeviceGetUtilizationRates.side_effect = make_util_for_device
        nvml.nvmlDeviceGetMemoryInfo.side_effect = make_mem_for_device

        sampler = GpuSampler(nvml=nvml, interval=0.01)
        sampler.start()
        time.sleep(0.05)
        stats = sampler.stop()

        assert len(stats.devices) == 2
        assert stats.devices[0].device_index == 0
        assert stats.devices[1].device_index == 1

    def test_skips_device_on_error(self):
        nvml = MagicMock()
        nvml.nvmlDeviceGetCount.return_value = 2
        nvml.nvmlDeviceGetHandleByIndex.side_effect = [
            MagicMock(),  # device 0 OK
            Exception("device removed"),  # device 1 fails
        ]
        util = MagicMock()
        util.gpu = 70
        nvml.nvmlDeviceGetUtilizationRates.return_value = util
        mem = MagicMock()
        mem.used = 1_000_000_000
        nvml.nvmlDeviceGetMemoryInfo.return_value = mem

        sampler = GpuSampler(nvml=nvml, interval=0.01)
        sampler.start()
        time.sleep(0.05)
        stats = sampler.stop()

        # Should not crash; device 0 stats recorded, device 1 silently skipped
        assert isinstance(stats, GpuStats)

    def test_available_property_reflects_device_presence(self):
        nvml = _make_nvml(device_count=1)
        sampler = GpuSampler(nvml=nvml, interval=0.01)
        sampler.start()
        time.sleep(0.05)
        stats = sampler.stop()

        assert stats.available is True
