"""Wear counters and degradation models (spec §8.3, FR-8.3.1-5).

Plant parameters (e.g. ``pump_capacity``) age from wear instead of
staying ideal forever, exercising the PLC's compensation and
maintenance logic:

    degradation:
      pump_capacity:                 # parameter name (float 0..1)
        from: 1.0                    # fresh value (also the reset target)
        model: linear                # linear | exponential | staircase
        rate: 0.0002                 # per runtime hour (linear/exp)
        per: runtime_hours           # wear basis
        on: pump_relay               # alias observed for wear (any kind)
        staircase: [0.95, 0.85, 0.7] # staircase model: levels per step
        maintenance_at: 0.75          # parameter threshold -> diagnosis
        diagnosis:                    # injected when threshold crossed
          module: 2                   # per-type DI/DO... module number
          channel: 0
          value: 0x0001

Counters update once per tick from the observed alias truth (runtime
hours, rising/falling edges). The engine reads parameters via
``engine.parameter(name)`` (FR-8.3.3); the maintenance alarm is the
PLC's to raise — the scenario asserts it (FR-8.3.4); the built-in
``maintenance_reset`` action clears diagnosis and restores ``from``
(FR-8.3.5).
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field


@dataclass
class WearCounters:
    """Per-alias wear (FR-8.3.1)."""

    runtime_hours: float = 0.0
    rising_edges: int = 0
    falling_edges: int = 0

    def as_dict(self) -> dict:
        return {
            "runtime_hours": self.runtime_hours,
            "rising_edges": self.rising_edges,
            "falling_edges": self.falling_edges,
        }


@dataclass
class DegradationModel:
    """One aging parameter (FR-8.3.2-5)."""

    name: str
    parameter: str          # parameter key, e.g. 'pump_capacity'
    observes: str           # alias name the wear counts come from
    model: str = "linear"   # linear | exponential | staircase
    start: float = 1.0      # fresh value; maintenance_reset target
    rate: float = 0.0       # per hour (linear/exponential) or per cycle
    basis: str = "runtime_hours"   # runtime_hours | cycles
    staircase: list[float] = field(default_factory=list)
    maintenance_at: float | None = None
    diagnosis: dict | None = None
    # runtime state
    value: float = 1.0
    cycles: int = 0
    diagnosed: bool = False

    def reset(self) -> None:
        """FR-8.3.5: restore to `start`; diagnosis cleared by the engine."""
        self.value = self.start
        self.cycles = 0

    def update(self, dt_hours: float, wear: WearCounters) -> None:
        """Advance the parameter one tick (FR-8.3.2)."""
        if self.basis == "cycles":
            if self.model == "staircase":
                level = min(self.cycles, len(self.staircase) - 1)
                self.value = self.staircase[level]
            elif self.model == "linear":
                self.value = self.start - self.rate * self.cycles
            elif self.model == "exponential":
                self.value = self.start * (self.rate ** self.cycles)
        else:
            if self.model == "staircase":
                # `rate` = levels per runtime hour; rate<=0 = one level/hour
                per_level_hours = (1.0 / self.rate) if self.rate > 0 else 1.0
                level = min(int(wear.runtime_hours / per_level_hours),
                            len(self.staircase) - 1)
                self.value = self.staircase[level]
            elif self.model == "exponential":
                self.value = self.start * (1.0 - self.rate) ** wear.runtime_hours
            else:  # linear
                self.value = self.start - self.rate * wear.runtime_hours
        self.value = max(0.0, min(1.0, self.value))


class DegradationManager:
    """All aging parameters of one continuous run (spec §8.3)."""

    def __init__(self, engine, spec: dict | None) -> None:
        self.engine = engine
        self.models: dict[str, DegradationModel] = {}
        self.wear: dict[str, WearCounters] = {}
        self._last_truth: dict[str, bool] = {}
        for param_name, cfg in (spec or {}).items():
            if not isinstance(cfg, dict):
                raise ValueError(f"degradation {param_name!r} must be a mapping")
            # NB: the key is 'observes', NOT 'on' — YAML 1.1 parses the
            # bare key `on` as boolean True (a classic trap)
            observes = str(cfg.get("observes", ""))
            if observes not in engine._aliases:
                raise ValueError(
                    f"degradation {param_name!r}: unknown alias {observes!r}")
            model = DegradationModel(
                name=param_name,
                parameter=str(cfg.get("parameter", param_name)),
                observes=observes,
                model=str(cfg.get("model", "linear")),
                start=float(cfg.get("from", 1.0)),
                rate=float(cfg.get("rate", 0.0)),
                basis=str(cfg.get("per", "runtime_hours")),
                staircase=[float(v) for v in cfg.get("staircase", [])],
                maintenance_at=(float(cfg["maintenance_at"])
                                 if "maintenance_at" in cfg else None),
                diagnosis=cfg.get("diagnosis"),
            )
            if model.model == "staircase" and not model.staircase:
                raise ValueError(
                    f"degradation {param_name!r}: staircase model needs 'staircase'")
            model.value = model.start
            self.models[param_name] = model
            self.wear.setdefault(observes, WearCounters())

    # -- per-tick update ------------------------------------------------------

    def tick(self, dt_ms: float) -> None:
        """Update wear counters and parameters for one tick (FR-8.2.12 step 4)."""
        dt_hours = dt_ms / 3_600_000.0
        for alias_name, wear in self.wear.items():
            truth = bool(self.engine.alias_value(alias_name))
            before = self._last_truth.get(alias_name, False)
            if truth:
                wear.runtime_hours += dt_hours
            if truth and not before:
                wear.rising_edges += 1
                # count a cycle per completed run: rising edge closes it
                for model in self.models.values():
                    if model.observes == alias_name and model.basis == "cycles":
                        model.cycles += 1
            if (not truth) and before:
                wear.falling_edges += 1
            self._last_truth[alias_name] = truth
        for model in self.models.values():
            model.update(dt_hours, self.wear[model.observes])
            if (model.maintenance_at is not None and not model.diagnosed
                    and model.value <= model.maintenance_at):
                model.diagnosed = True
                self._inject_diagnosis(model)

    def _inject_diagnosis(self, model: DegradationModel) -> None:
        """FR-8.3.4: threshold crossed -> module diagnosis injected."""
        diag = model.diagnosis or {}
        module = int(diag.get("module", 0))
        channel = int(diag.get("channel", 0))
        value = int(diag.get("value", 0x0001))
        registers = getattr(self.engine, "registers", None)
        if registers is not None:
            registers.set_module_diagnosis(module, channel, value)

    # -- engine interface (FR-8.3.3) -------------------------------------------

    def parameter(self, name: str) -> float:
        """Parameter value for plant physics (pump models etc.)."""
        model = self.models.get(name)
        if model is None:
            raise KeyError(f"unknown parameter {name!r}")
        return model.value

    def maintenance_reset(self) -> None:
        """FR-8.3.5 action: restore parameters, clear diagnosis."""
        for model in self.models.values():
            model.reset()
            model.diagnosed = False
        registers = getattr(self.engine, "registers", None)
        if registers is not None:
            registers.set_module_diagnosis(0, 0, 0)

    def snapshot(self) -> dict:
        """Wear + parameter values for report snapshots (FR-8.4.3)."""
        return {
            "parameters": {n: m.value for n, m in self.models.items()},
            "wear": {a: w.as_dict() for a, w in self.wear.items()},
        }