"""Per-metric reader / normaliser registry for ``SystemMetricsValueStream``.

The stream previously dispatched on a ``self._metric`` string across three
separate ``if / elif`` chains (priming in ``start()``, raw read in
``_read_metric_psutil`` + ``_read_metric_fallback``, normalisation in
``_normalize``). Adding a new metric meant editing every chain.

This module replaces all of that with a single ``METRIC_SPECS`` dict keyed
by metric name. Each :class:`MetricSpec` declares:

* ``read_psutil(stream)`` — the desktop path that uses ``stream._psutil``;
* ``read_fallback(stream)`` — the Android / no-psutil path (returns 0.0
  for desktop-only sensors);
* ``normalize(stream, raw)`` — maps the raw reading to ``[0, 1]``;
* ``prime(stream)`` — optional one-time setup called from ``start()``.

The spec functions operate on the stream's existing attributes
(``_disk_path``, ``_sensor_label``, ``_min_val``, ``_max_val``,
``_max_rate``, ``_gpu_unavailable``, ``_prev_net_bytes``,
``_prev_net_time``). That is intentional: the readers are stateless
strategy callables, but the *stream's* state remains its own. Mutating it
from a reader is documented per function so the contract is explicit.
"""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable, Dict

from ledgrab.utils import get_logger

if TYPE_CHECKING:
    from ledgrab.core.processing.value_stream import SystemMetricsValueStream

logger = get_logger(__name__)


# ---------------------------------------------------------------------------
# Spec dataclass
# ---------------------------------------------------------------------------


ReaderFn = Callable[["SystemMetricsValueStream"], float]
NormalizeFn = Callable[["SystemMetricsValueStream", float], float]
PrimeFn = Callable[["SystemMetricsValueStream"], None]


@dataclass(frozen=True)
class MetricSpec:
    """How to read, normalise, and prime a single system metric."""

    name: str
    read_psutil: ReaderFn
    read_fallback: ReaderFn
    normalize: NormalizeFn
    prime: PrimeFn | None = None


# ---------------------------------------------------------------------------
# Normaliser primitives — shared across metrics
# ---------------------------------------------------------------------------


def _norm_percent(_s, raw: float) -> float:
    """Percent metrics (cpu_load, ram_usage, …) — raw is 0..100."""
    return max(0.0, min(1.0, raw / 100.0))


def _norm_range(s, raw: float) -> float:
    """Temperature / fan-speed metrics normalise against (min_val, max_val)."""
    rng = s._max_val - s._min_val
    if abs(rng) < 1e-9:
        return 0.5
    return max(0.0, min(1.0, (raw - s._min_val) / rng))


def _norm_rate(s, raw: float) -> float:
    """Network rate normalises against ``_max_rate`` (bytes/s)."""
    if s._max_rate <= 0:
        return 0.5
    return max(0.0, min(1.0, raw / s._max_rate))


def _zero(_s) -> float:
    """Desktop-only sensor on a no-psutil platform: report 0.0."""
    return 0.0


# ---------------------------------------------------------------------------
# Read helpers — psutil paths
# ---------------------------------------------------------------------------


def _read_cpu_load_psutil(s) -> float:
    return s._psutil.cpu_percent(interval=None)


def _read_ram_usage_psutil(s) -> float:
    return s._psutil.virtual_memory().percent


def _read_disk_usage_psutil(s) -> float:
    return s._psutil.disk_usage(s._disk_path).percent


def _read_battery_psutil(s) -> float:
    bat = s._psutil.sensors_battery()
    return bat.percent if bat else 0.0


def _read_cpu_temp_psutil(s) -> float:
    psutil = s._psutil
    if psutil is None:
        return 0.0
    temps = psutil.sensors_temperatures()
    if not temps:
        return 0.0
    if s._sensor_label:
        for group_name, entries in temps.items():
            for entry in entries:
                if entry.label == s._sensor_label or group_name == s._sensor_label:
                    return entry.current
    for entries in temps.values():
        if entries:
            return entries[0].current
    return 0.0


def _read_fan_speed_psutil(s) -> float:
    psutil = s._psutil
    if psutil is None:
        return 0.0
    fans = psutil.sensors_fans()
    if not fans:
        return 0.0
    if s._sensor_label:
        for group_name, entries in fans.items():
            for entry in entries:
                if entry.label == s._sensor_label or group_name == s._sensor_label:
                    return entry.current
    for entries in fans.values():
        if entries:
            return entries[0].current
    return 0.0


def _read_gpu(metric: str) -> ReaderFn:
    """Build a GPU reader for the given metric name ('gpu_load' or 'gpu_temp')."""

    def _read(s) -> float:
        if s._gpu_unavailable:
            return 0.0
        try:
            from ledgrab.utils.gpu import nvml, nvml_available, nvml_handle

            if not nvml_available or nvml_handle is None:
                s._gpu_unavailable = True
                return 0.0
            if metric == "gpu_load":
                util = nvml.nvmlDeviceGetUtilizationRates(nvml_handle)
                return float(util.gpu)
            # gpu_temp
            return float(nvml.nvmlDeviceGetTemperature(nvml_handle, 0))
        except Exception as e:
            logger.debug("GPU metric read error: %s", e)
            s._gpu_unavailable = True
            return 0.0

    return _read


def _read_network_rate(s) -> float:
    """Bytes/s rate for ``network_rx`` / ``network_tx``.

    Mutates ``s._prev_net_bytes`` and ``s._prev_net_time`` to track the
    delta between calls — the stream owns the cadence state, this reader
    just bumps it forward.
    """
    psutil = s._psutil
    if psutil is None:
        return 0.0
    counters = psutil.net_io_counters()
    if not counters:
        return 0.0
    current_bytes = counters.bytes_recv if s._metric == "network_rx" else counters.bytes_sent
    now = time.monotonic()
    if s._prev_net_bytes is None or s._prev_net_time is None:
        s._prev_net_bytes = current_bytes
        s._prev_net_time = now
        return 0.0
    dt = now - s._prev_net_time
    if dt <= 0:
        return 0.0
    # Cap delta time to avoid spikes after long gaps
    dt = min(dt, s._poll_interval * 2)
    rate = (current_bytes - s._prev_net_bytes) / dt
    s._prev_net_bytes = current_bytes
    s._prev_net_time = now
    return max(0.0, rate)


# ---------------------------------------------------------------------------
# Read helpers — fallback (no-psutil) paths
# ---------------------------------------------------------------------------


def _read_cpu_load_fallback(_s) -> float:
    from ledgrab.utils.metrics import get_metrics_provider

    return get_metrics_provider().cpu_percent()


def _read_ram_usage_fallback(_s) -> float:
    from ledgrab.utils.metrics import get_metrics_provider

    mem = get_metrics_provider().virtual_memory()
    if mem.total_bytes > 0:
        return (mem.used_bytes / mem.total_bytes) * 100.0
    return 0.0


# ---------------------------------------------------------------------------
# Prime helpers — one-time setup from start()
# ---------------------------------------------------------------------------


def _prime_cpu_load(s) -> None:
    # Prime psutil.cpu_percent so the first real call returns meaningful data
    if s._psutil is not None:
        s._psutil.cpu_percent(interval=None)


def _prime_network(s) -> None:
    """Capture initial network counter so the first delta has a baseline."""
    if s._psutil is None:
        return
    counters = s._psutil.net_io_counters()
    if counters:
        s._prev_net_bytes = (
            counters.bytes_recv if s._metric == "network_rx" else counters.bytes_sent
        )
        s._prev_net_time = time.monotonic()


# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------


METRIC_SPECS: Dict[str, MetricSpec] = {
    "cpu_load": MetricSpec(
        "cpu_load", _read_cpu_load_psutil, _read_cpu_load_fallback, _norm_percent, _prime_cpu_load
    ),
    "ram_usage": MetricSpec(
        "ram_usage", _read_ram_usage_psutil, _read_ram_usage_fallback, _norm_percent
    ),
    "disk_usage": MetricSpec("disk_usage", _read_disk_usage_psutil, _zero, _norm_percent),
    "battery_level": MetricSpec("battery_level", _read_battery_psutil, _zero, _norm_percent),
    "cpu_temp": MetricSpec("cpu_temp", _read_cpu_temp_psutil, _zero, _norm_range),
    "fan_speed": MetricSpec("fan_speed", _read_fan_speed_psutil, _zero, _norm_range),
    "gpu_load": MetricSpec("gpu_load", _read_gpu("gpu_load"), _zero, _norm_percent),
    "gpu_temp": MetricSpec("gpu_temp", _read_gpu("gpu_temp"), _zero, _norm_range),
    "network_rx": MetricSpec("network_rx", _read_network_rate, _zero, _norm_rate, _prime_network),
    "network_tx": MetricSpec("network_tx", _read_network_rate, _zero, _norm_rate, _prime_network),
}


def get_spec(metric: str) -> MetricSpec | None:
    """Look up the spec for ``metric``, returning ``None`` for unknown names."""
    return METRIC_SPECS.get(metric)
