import os
import statistics
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 CpuMemoryStats:
    avg_cpu_percent: float
    max_cpu_percent: float
    p95_cpu_percent: float
    avg_memory_mb: float
    max_memory_mb: float


@dataclass(frozen=True)
class _Sample:
    cpu_percent: float
    memory_mb: float


class CpuMemorySampler(BaseSampler[CpuMemoryStats]):
    """Samples CPU utilisation and RSS memory of the current process."""

    DEFAULT_INTERVAL = 1.0

    def __init__(
        self,
        interval: float = DEFAULT_INTERVAL,
        process: Optional[psutil.Process] = None,
    ) -> None:
        super().__init__(interval)
        self._process = process or psutil.Process(os.getpid())

    def start(self) -> None:
        # Prime the CPU measurement so the first real sample has a reference delta.
        self._process.cpu_percent()
        super().start()

    def _take_sample(self) -> Optional[_Sample]:
        try:
            return _Sample(
                cpu_percent=self._process.cpu_percent(),
                memory_mb=self._process.memory_info().rss / _BYTES_PER_MB,
            )
        except Exception:
            return None

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

        cpu = [s.cpu_percent for s in self._samples]
        mem = [s.memory_mb for s in self._samples]

        return CpuMemoryStats(
            avg_cpu_percent=statistics.mean(cpu),
            max_cpu_percent=max(cpu),
            p95_cpu_percent=_p95(cpu),
            avg_memory_mb=statistics.mean(mem),
            max_memory_mb=max(mem),
        )


def _p95(values: list[float]) -> float:
    if len(values) < 2:
        return float(values[0])
    return float(statistics.quantiles(values, n=100)[94])
