zlock_slots/custom_components/zlock_slots/coordinator.py

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from homeassistant.core import HomeAssistant, Event, callback
from homeassistant.helpers.event import async_track_state_change_event
from homeassistant.helpers.storage import Store
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator

from .const import (
  DOMAIN,
  STORAGE_KEY,
  STORAGE_VERSION,
  CONF_LOCK_ENTITY_IDS,
  CONF_GROUPS,
  CONF_GROUP_NAME,
  ALARM_MAP,
  UserSource,
)

if TYPE_CHECKING:
  from homeassistant.config_entries import ConfigEntry

_LOGGER = logging.getLogger(__name__)


@dataclass
class SlotInfo:
  status: str  # "enabled" or "disabled"
  name: str | None = None


@dataclass
class GroupState:
  last_action: str | None = None
  last_user: str | None = None


@dataclass
class LockState:
  slots: dict[int, SlotInfo] = field(default_factory=dict)
  last_action: str | None = None
  last_user: str | None = None


def _alarm_entity(lock_entity_id: str, suffix: str) -> str:
  """Derive alarm sensor entity ID from lock entity ID.

  lock.living_room_front_lock -> sensor.living_room_front_lock_alarm_<suffix>
  """
  object_id = lock_entity_id.split(".")[-1]
  return f"sensor.{object_id}_alarm_{suffix}"


class ZLockCoordinator(DataUpdateCoordinator[dict[str, LockState]]):
  """Manages slot data, name storage, and event listening for all configured locks."""

  def __init__(self, hass: HomeAssistant, entry: "ConfigEntry") -> None:
    super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=None)
    self._entry = entry
    self._lock_entity_ids: list[str] = entry.data[CONF_LOCK_ENTITY_IDS]
    self._store: Store = Store(hass, STORAGE_VERSION, STORAGE_KEY)
    self._names: dict[str, str] = {}  # "{entity_id}:{slot_num}" → name
    self._alarm_type_to_lock: dict[str, str] = {
      _alarm_entity(eid, "type"): eid for eid in self._lock_entity_ids
    }
    self._alarm_level_to_lock: dict[str, str] = {
      _alarm_entity(eid, "level"): eid for eid in self._lock_entity_ids
    }
    # State machine for pairing alarm_type and alarm_level, which arrive separately.
    # _pending_slot_alarm: alarm_type fired with SLOT source; waiting for alarm_level.
    # _recent_level: alarm_level arrived first; alarm_type may follow within the window.
    self._pending_slot_alarm: dict[str, AlarmMapping] = {}
    self._recent_level: dict[str, tuple[int, float]] = {}  # lock → (slot, monotonic_time)
    self._unsub_event: callback | None = None
    self.data: dict[str, LockState] = {eid: LockState() for eid in self._lock_entity_ids}

    groups: list[dict] = entry.options.get(CONF_GROUPS, [])
    self.group_data: dict[str, GroupState] = {
      g[CONF_GROUP_NAME]: GroupState() for g in groups
    }
    # lock_entity_id → list of group names it belongs to
    self._lock_to_groups: dict[str, list[str]] = {}
    for g in groups:
      for lid in g.get(CONF_LOCK_ENTITY_IDS, []):
        self._lock_to_groups.setdefault(lid, []).append(g[CONF_GROUP_NAME])

  async def async_setup(self) -> None:
    stored = await self._store.async_load()
    if stored:
      self._names = stored.get("names", {})
      last_events = stored.get("last_events", {})
      for eid, ev in last_events.get("locks", {}).items():
        if eid in self.data:
          self.data[eid].last_action = ev.get("last_action")
          self.data[eid].last_user = ev.get("last_user")
      for gname, ev in last_events.get("groups", {}).items():
        if gname in self.group_data:
          self.group_data[gname].last_action = ev.get("last_action")
          self.group_data[gname].last_user = ev.get("last_user")

    await self._async_fetch_all_slots()

    self._unsub_event = async_track_state_change_event(
      self.hass,
      list(self._alarm_type_to_lock) + list(self._alarm_level_to_lock),
      self._handle_alarm_state_change,
    )

  def async_teardown(self) -> None:
    if self._unsub_event:
      self._unsub_event()
      self._unsub_event = None

  async def _async_fetch_all_slots(self) -> None:
    for lock_entity_id in self._lock_entity_ids:
      slots = await self._fetch_slots(lock_entity_id)
      self.data[lock_entity_id].slots = slots
    self.async_update_listeners()

  async def _fetch_slots(self, lock_entity_id: str) -> dict[int, SlotInfo]:
    try:
      from homeassistant.components.zwave_js.helpers import async_get_node_from_entity_id
      node = async_get_node_from_entity_id(self.hass, lock_entity_id)
      users = await node.access_control.get_users_cached()
    except Exception:
      _LOGGER.exception("Failed to fetch code slots for %s", lock_entity_id)
      return {}

    slots: dict[int, SlotInfo] = {}
    for user in users:
      slot_num = user.user_id
      if slot_num == 0:
        continue
      slots[slot_num] = SlotInfo(
        status="enabled" if user.active else "disabled",
        name=self.get_name(lock_entity_id, slot_num),
      )
    return slots

  # alarm_type and alarm_level are separate Z-Wave parameters and can arrive in
  # either order. We subscribe to both and pair them with a short time window.
  _LEVEL_WINDOW_S: float = 5.0

  @callback
  def _handle_alarm_state_change(self, event: Event) -> None:
    entity_id: str = event.data["entity_id"]
    new_state = event.data.get("new_state")
    if new_state is None:
      return

    if lock_entity_id := self._alarm_type_to_lock.get(entity_id):
      self._on_alarm_type(lock_entity_id, new_state.state)
    elif lock_entity_id := self._alarm_level_to_lock.get(entity_id):
      self._on_alarm_level(lock_entity_id, new_state.state)

  @callback
  def _on_alarm_type(self, lock_entity_id: str, state: str) -> None:
    # Clear any pending SLOT alarm — a new alarm_type supersedes it.
    self._pending_slot_alarm.pop(lock_entity_id, None)

    try:
      alarm_type = int(state)
    except (ValueError, TypeError):
      return

    mapping = ALARM_MAP.get(alarm_type)
    if not mapping:
      return

    if mapping.user_source == UserSource.MANUAL:
      self._emit_alarm(lock_entity_id, mapping.action, "<Manual>")
    elif mapping.user_source == UserSource.RADIO:
      self._emit_alarm(lock_entity_id, mapping.action, "<Radio>")
    else:
      # SLOT: check if alarm_level arrived recently (level-first ordering).
      recent = self._recent_level.get(lock_entity_id)
      if recent and (time.monotonic() - recent[1]) < self._LEVEL_WINDOW_S:
        slot_num = recent[0]
        self._emit_alarm(lock_entity_id, mapping.action,
                         self.get_name(lock_entity_id, slot_num) or str(slot_num))
      else:
        # Level not yet received; store pending and resolve when it arrives.
        self._pending_slot_alarm[lock_entity_id] = mapping

  @callback
  def _on_alarm_level(self, lock_entity_id: str, state: str) -> None:
    try:
      slot_num = int(state)
    except (ValueError, TypeError):
      return

    self._recent_level[lock_entity_id] = (slot_num, time.monotonic())

    # Resolve a pending SLOT alarm that arrived before the level.
    pending = self._pending_slot_alarm.pop(lock_entity_id, None)
    if pending:
      self._emit_alarm(lock_entity_id, pending.action,
                       self.get_name(lock_entity_id, slot_num) or str(slot_num))

  @callback
  def _emit_alarm(self, lock_entity_id: str, action: str, user_str: str) -> None:
    self.data[lock_entity_id].last_action = action
    self.data[lock_entity_id].last_user = user_str
    for group_name in self._lock_to_groups.get(lock_entity_id, []):
      self.group_data[group_name].last_action = action
      self.group_data[group_name].last_user = user_str
    self.hass.async_create_task(self._async_save())
    self.async_update_listeners()

  def get_name(self, lock_entity_id: str, slot_num: int) -> str | None:
    return self._names.get(f"{lock_entity_id}:{slot_num}")

  async def _async_save(self) -> None:
    await self._store.async_save({
      "names": self._names,
      "last_events": {
        "locks": {
          eid: {"last_action": s.last_action, "last_user": s.last_user}
          for eid, s in self.data.items()
          if s.last_action is not None or s.last_user is not None
        },
        "groups": {
          gname: {"last_action": g.last_action, "last_user": g.last_user}
          for gname, g in self.group_data.items()
          if g.last_action is not None or g.last_user is not None
        },
      },
    })

  async def async_set_name(self, lock_entity_id: str, slot_num: int, name: str) -> None:
    key = f"{lock_entity_id}:{slot_num}"
    if name:
      self._names[key] = name
    else:
      self._names.pop(key, None)
    await self._async_save()

    state = self.data.get(lock_entity_id)
    if state and slot_num in state.slots:
      state.slots[slot_num].name = name or None
    self.async_update_listeners()

  async def async_refresh_lock(self, lock_entity_id: str) -> None:
    slots = await self._fetch_slots(lock_entity_id)
    self.data[lock_entity_id].slots = slots
    self.async_update_listeners()

  async def _async_update_data(self) -> dict[str, LockState]:
    await self._async_fetch_all_slots()
    return self.data