zlock_slots/custom_components/zlock_slots/sensor.py

from __future__ import annotations

from homeassistant.components.sensor import SensorEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity

from .const import DOMAIN, CONF_LOCK_ENTITY_IDS, CONF_GROUPS, CONF_GROUP_NAME
from .coordinator import ZLockCoordinator


async def async_setup_entry(
  hass: HomeAssistant,
  entry: ConfigEntry,
  async_add_entities: AddEntitiesCallback,
) -> None:
  coordinator: ZLockCoordinator = hass.data[DOMAIN][entry.entry_id]
  entities = []
  for lock_entity_id in entry.data[CONF_LOCK_ENTITY_IDS]:
    entities.append(ZLockLastActionSensor(coordinator, entry, lock_entity_id))
    entities.append(ZLockLastUserSensor(coordinator, entry, lock_entity_id))
    entities.append(ZLockSlotsSensor(coordinator, entry, lock_entity_id))
  for group in entry.options.get(CONF_GROUPS, []):
    group_name = group[CONF_GROUP_NAME]
    entities.append(ZLockGroupLastActionSensor(coordinator, entry, group_name))
    entities.append(ZLockGroupLastUserSensor(coordinator, entry, group_name))
  async_add_entities(entities)


def _lock_object_id(lock_entity_id: str) -> str:
  return lock_entity_id.split(".")[-1]


class _ZLockBaseSensor(CoordinatorEntity[ZLockCoordinator], SensorEntity):
  _attr_has_entity_name = True

  def __init__(
    self,
    coordinator: ZLockCoordinator,
    entry: ConfigEntry,
    lock_entity_id: str,
  ) -> None:
    super().__init__(coordinator)
    self._lock_entity_id = lock_entity_id
    self._entry = entry

  @property
  def _lock_state(self):
    return self.coordinator.data.get(self._lock_entity_id)


class ZLockLastActionSensor(_ZLockBaseSensor):
  _attr_translation_key = "last_action"

  def __init__(self, coordinator, entry, lock_entity_id):
    super().__init__(coordinator, entry, lock_entity_id)
    self._attr_unique_id = f"{entry.entry_id}_{lock_entity_id}_last_action"

  @property
  def name(self) -> str:
    return f"{_lock_object_id(self._lock_entity_id)} last action"

  @property
  def native_value(self) -> str | None:
    state = self._lock_state
    return state.last_action if state else None


class ZLockLastUserSensor(_ZLockBaseSensor):
  _attr_translation_key = "last_user"

  def __init__(self, coordinator, entry, lock_entity_id):
    super().__init__(coordinator, entry, lock_entity_id)
    self._attr_unique_id = f"{entry.entry_id}_{lock_entity_id}_last_user"

  @property
  def name(self) -> str:
    return f"{_lock_object_id(self._lock_entity_id)} last user"

  @property
  def native_value(self) -> str | None:
    state = self._lock_state
    return state.last_user if state else None


class ZLockSlotsSensor(_ZLockBaseSensor):
  _attr_translation_key = "slots"

  def __init__(self, coordinator, entry, lock_entity_id):
    super().__init__(coordinator, entry, lock_entity_id)
    self._attr_unique_id = f"{entry.entry_id}_{lock_entity_id}_slots"

  @property
  def name(self) -> str:
    return f"{_lock_object_id(self._lock_entity_id)} slots"

  @property
  def native_value(self) -> int:
    state = self._lock_state
    return len(state.slots) if state else 0

  @property
  def extra_state_attributes(self) -> dict:
    state = self._lock_state
    if not state:
      return {"slots": []}
    return {
      "slots": [
        {
          "slot_num": num,
          "status": info.status,
          "name": info.name,
        }
        for num, info in sorted(state.slots.items())
      ]
    }


class _ZLockGroupBaseSensor(CoordinatorEntity[ZLockCoordinator], SensorEntity):
  _attr_has_entity_name = True

  def __init__(
    self,
    coordinator: ZLockCoordinator,
    entry: ConfigEntry,
    group_name: str,
  ) -> None:
    super().__init__(coordinator)
    self._group_name = group_name
    self._entry = entry

  @property
  def _group_state(self):
    return self.coordinator.group_data.get(self._group_name)


class ZLockGroupLastActionSensor(_ZLockGroupBaseSensor):
  _attr_translation_key = "last_action"

  def __init__(self, coordinator, entry, group_name):
    super().__init__(coordinator, entry, group_name)
    self._attr_unique_id = f"{entry.entry_id}_group_{group_name}_last_action"

  @property
  def name(self) -> str:
    return f"{self._group_name} last action"

  @property
  def native_value(self) -> str | None:
    state = self._group_state
    return state.last_action if state else None


class ZLockGroupLastUserSensor(_ZLockGroupBaseSensor):
  _attr_translation_key = "last_user"

  def __init__(self, coordinator, entry, group_name):
    super().__init__(coordinator, entry, group_name)
    self._attr_unique_id = f"{entry.entry_id}_group_{group_name}_last_user"

  @property
  def name(self) -> str:
    return f"{self._group_name} last user"

  @property
  def native_value(self) -> str | None:
    state = self._group_state
    return state.last_user if state else None