import contextlib
import os
from dataclasses import dataclass
from typing import Optional

import psutil

from .base import BaseSampler

_BYTES_PER_MB = 1_000_000.0


@dataclass(frozen=True)
class DiskStats:
    read_bytes: int
    write_bytes: int
    avg_read_throughput_mbps: float
    avg_write_throughput_mbps: float


@dataclass(frozen=True)
class _Sample:
    read_bytes_delta: int
    write_bytes_delta: int


class DiskSampler(BaseSampler[DiskStats]):
    """Samples disk I/O for the current process and optionally its children.

    Falls back to zero stats on platforms where ``io_counters`` is unavailable
    (e.g. macOS) or when access is denied.
    """

    DEFAULT_INTERVAL = 1.0

    def __init__(
        self,
        interval: float = DEFAULT_INTERVAL,
        include_children: bool = True,
        process: Optional[psutil.Process] = None,
    ) -> None:
        super().__init__(interval)
        self._include_children = include_children
        self._process = process or psutil.Process(os.getpid())
        self._available = self._probe_availability()
        self._baseline: Optional[tuple[int, int]] = None

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

    def start(self) -> None:
        if self._available:
            self._baseline = self._read_io_counters()
        super().start()

    def _take_sample(self) -> Optional[_Sample]:
        if not self._available or self._baseline is None:
            return None
        try:
            current = self._read_io_counters()
            delta = _Sample(
                read_bytes_delta=current[0] - self._baseline[0],
                write_bytes_delta=current[1] - self._baseline[1],
            )
            self._baseline = current
            return delta
        except Exception:
            return None

    def _compute_stats(self) -> DiskStats:
        if not self._samples:
            return DiskStats(0, 0, 0.0, 0.0)

        total_read = sum(s.read_bytes_delta for s in self._samples)
        total_write = sum(s.write_bytes_delta for s in self._samples)
        n = len(self._samples)

        return DiskStats(
            read_bytes=total_read,
            write_bytes=total_write,
            avg_read_throughput_mbps=(total_read / n) / _BYTES_PER_MB / self._interval,
            avg_write_throughput_mbps=(total_write / n) / _BYTES_PER_MB / self._interval,
        )

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

    def _probe_availability(self) -> bool:
        try:
            self._process.io_counters()
            return True
        except (AttributeError, psutil.AccessDenied, NotImplementedError):
            return False

    def _read_io_counters(self) -> tuple[int, int]:
        procs: list[psutil.Process] = [self._process]
        if self._include_children:
            with contextlib.suppress(psutil.NoSuchProcess):
                procs += self._process.children(recursive=True)

        read = write = 0
        for proc in procs:
            try:
                counters = proc.io_counters()
                read += counters.read_bytes
                write += counters.write_bytes
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
        return read, write
