HomeAssistantVS/custom_components/maintenance_supporter/entity/triggers/counter.py

"""Counter trigger for maintenance tasks."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from homeassistant.core import HomeAssistant

if TYPE_CHECKING:
    from ...sensor import MaintenanceSensor

from .base_trigger import BaseTrigger

_LOGGER = logging.getLogger(__name__)


class CounterTrigger(BaseTrigger):
    """Trigger that activates when a counter reaches a target value.

    Supports:
    - Absolute mode: triggers when value >= target
    - Delta mode: triggers when (value - baseline) >= target
    """

    def __init__(
        self,
        hass: HomeAssistant,
        entity: MaintenanceSensor,
        trigger_config: dict[str, Any],
    ) -> None:
        """Initialize counter trigger."""
        super().__init__(hass, entity, trigger_config)

        self._target_value: float = trigger_config.get("trigger_target_value", 0)
        self._delta_mode: bool = trigger_config.get("trigger_delta_mode", False)
        self._baseline_value: float | None = trigger_config.get("trigger_baseline_value")
        # Set when reset_baseline is called while the source is unavailable, so
        # the re-baseline is applied against the next real value (M4).
        self._reset_pending: bool = False

    async def async_setup(self) -> None:
        """Set up counter trigger with baseline initialization."""
        # Restore persisted baseline BEFORE super().async_setup() which calls
        # evaluate().  Without this, evaluate() sets _baseline_value to the
        # current sensor value, discarding the saved baseline from the Store.
        if self._delta_mode and self._baseline_value is None:
            saved = self.config.get("_trigger_state", {}).get(self.entity_id, {}).get("baseline_value")
            if saved is not None:
                self._baseline_value = saved
                _LOGGER.debug(
                    "Counter baseline restored from store: %s = %s",
                    self.entity_id,
                    self._baseline_value,
                )

        await super().async_setup()

        # Set initial baseline if in delta mode and no baseline exists
        if self._delta_mode and self._baseline_value is None and self._current_value is not None:
            self._baseline_value = self._current_value
            await self._persist_baseline()
            _LOGGER.debug(
                "Counter baseline initialized: %s = %s",
                self.entity_id,
                self._baseline_value,
            )

    def _evaluate_and_update(self, value: float) -> None:
        """Initialize/re-baseline before evaluation, then persist."""
        if self._delta_mode and self._reset_pending:
            # Apply the deferred post-completion re-baseline now that a real
            # value is available (the source was unavailable at completion). (M4)
            self._baseline_value = value
            self._reset_pending = False
            self.hass.async_create_task(
                self._persist_baseline(),
                eager_start=False,
            )
        elif self._delta_mode and self._baseline_value is None:
            self._baseline_value = value
            self.hass.async_create_task(
                self._persist_baseline(),
                eager_start=False,
            )
            _LOGGER.debug(
                "Counter baseline initialized on state change: %s = %s",
                self.entity_id,
                self._baseline_value,
            )
        super()._evaluate_and_update(value)

    def evaluate(self, value: float) -> bool:
        """Evaluate counter condition."""
        if self._delta_mode:
            if self._baseline_value is None:
                # Fallback — should be caught by _evaluate_and_update
                self._baseline_value = value
                return False
            if value < self._baseline_value:
                # Source counter reset / rolled over (device reboot, meter
                # rollover) — re-baseline to the new lower value so we don't
                # miss the next interval waiting for it to climb back past the
                # old baseline. (M3)
                self._baseline_value = value
                self.hass.async_create_task(
                    self._persist_baseline(),
                    eager_start=False,
                )
                return False
            delta = value - self._baseline_value
            return delta >= self._target_value

        # Absolute mode
        return value >= self._target_value

    @property
    def current_delta(self) -> float | None:
        """Return the current delta from baseline."""
        if not self._delta_mode or self._baseline_value is None:
            return None
        if self._current_value is None:
            return None
        return self._current_value - self._baseline_value

    def reset_baseline(self) -> None:
        """Reset the baseline to current value (after maintenance)."""
        if self._current_value is not None:
            self._baseline_value = self._current_value
            self._reset_pending = False
            self.hass.async_create_task(self._persist_baseline())
            _LOGGER.debug(
                "Counter baseline reset: %s = %s",
                self.entity_id,
                self._baseline_value,
            )
        elif self._delta_mode:
            # Source unavailable at completion — defer the re-baseline to the
            # next real value so a stale baseline can't immediately re-trigger
            # the just-completed task when the entity returns. (M4; delta-only,
            # absolute mode has no baseline to defer.)
            self._reset_pending = True
            _LOGGER.debug(
                "Counter baseline reset deferred (source unavailable): %s",
                self.entity_id,
            )

    async def _persist_baseline(self) -> None:
        """Persist baseline value to the Store for survival across restarts."""
        if self._baseline_value is not None:
            await self._coordinator.async_persist_trigger_runtime(
                self._task_id,
                {"baseline_value": self._baseline_value},
                entity_id=self.entity_id,
                immediate=True,
            )

    def reset(self) -> None:
        """Reset trigger and baseline."""
        super().reset()
        self.reset_baseline()