import statistics
from dataclasses import dataclass, field
from typing import Any, Optional

from .base import BaseSampler

_BYTES_PER_MB = 1_000_000.0


@dataclass(frozen=True)
class GpuDeviceStats:
    device_index: int
    avg_utilization_percent: float
    max_utilization_percent: float
    avg_memory_used_mb: float
    max_memory_used_mb: float


@dataclass(frozen=True)
class GpuStats:
    devices: list[GpuDeviceStats] = field(default_factory=list)

    @property
    def available(self) -> bool:
        return len(self.devices) > 0


@dataclass(frozen=True)
class _Sample:
    device_index: int
    utilization_percent: float
    memory_used_mb: float


def _load_nvml() -> Optional[Any]:
    try:
        import pynvml

        pynvml.nvmlInit()
        return pynvml
    except Exception:
        return None


class GpuSampler(BaseSampler[GpuStats]):
    """Samples GPU utilisation and memory for all NVIDIA devices.

    Requires the ``pynvml`` package (``pip install metaflow-observability[gpu]``).
    Returns an empty ``GpuStats`` when ``pynvml`` is unavailable or no devices
    are present, so callers never need to branch on GPU availability.
    """

    DEFAULT_INTERVAL = 1.0

    def __init__(
        self,
        interval: float = DEFAULT_INTERVAL,
        nvml: Optional[Any] = _load_nvml,  # callable sentinel — resolved on init
    ) -> None:
        super().__init__(interval)
        # Accept either a pre-built nvml mock (tests) or the sentinel callable.
        self._nvml = nvml() if nvml is _load_nvml else nvml
        self._device_count = self._probe_device_count()

    # ------------------------------------------------------------------
    # BaseSampler interface
    # ------------------------------------------------------------------

    def _collect(self) -> None:
        """Override to fan out across all devices in a single polling cycle."""
        for sample in self._take_all_device_samples():
            self._samples.append(sample)

    def _take_sample(self) -> None:
        # Not used — _collect is overridden directly to fan out across devices.
        pass

    def _compute_stats(self) -> GpuStats:
        if not self._samples:
            return GpuStats()

        by_device: dict[int, list[_Sample]] = {}
        for s in self._samples:
            by_device.setdefault(s.device_index, []).append(s)

        devices = [
            GpuDeviceStats(
                device_index=idx,
                avg_utilization_percent=statistics.mean(s.utilization_percent for s in samples),
                max_utilization_percent=max(s.utilization_percent for s in samples),
                avg_memory_used_mb=statistics.mean(s.memory_used_mb for s in samples),
                max_memory_used_mb=max(s.memory_used_mb for s in samples),
            )
            for idx, samples in sorted(by_device.items())
        ]
        return GpuStats(devices=devices)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _probe_device_count(self) -> int:
        if self._nvml is None:
            return 0
        try:
            return int(self._nvml.nvmlDeviceGetCount())
        except Exception:
            return 0

    def _take_all_device_samples(self) -> list[_Sample]:
        if self._nvml is None or self._device_count == 0:
            return []

        samples = []
        for i in range(self._device_count):
            try:
                handle = self._nvml.nvmlDeviceGetHandleByIndex(i)
                util = self._nvml.nvmlDeviceGetUtilizationRates(handle)
                mem = self._nvml.nvmlDeviceGetMemoryInfo(handle)
                samples.append(
                    _Sample(
                        device_index=i,
                        utilization_percent=float(util.gpu),
                        memory_used_mb=mem.used / _BYTES_PER_MB,
                    )
                )
            except Exception:
                pass
        return samples
