ha-reticulum/custom_components/reticulum/sensor.py

"""Sensor platform for Reticulum."""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any

from homeassistant.components.sensor import (
    SensorDeviceClass,
    SensorEntity,
    SensorEntityDescription,
    SensorStateClass,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import (
    PERCENTAGE,
    SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
    EntityCategory,
    UnitOfDataRate,
    UnitOfInformation,
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback

from .const import DOMAIN, SUBENTRY_TYPE_IDENTITY
from .entity import ReticulumIdentityEntity, ReticulumStackEntity
from .reticulum_client import IdentityState, StackState


@dataclass(frozen=True, kw_only=True)
class StackSensorDescription(SensorEntityDescription):
    """Hub sensor."""

    value_fn: Callable[[StackState], Any]
    attr_fn: Callable[[StackState], dict[str, Any]] | None = None


@dataclass(frozen=True, kw_only=True)
class IdentitySensorDescription(SensorEntityDescription):
    """Per-identity sensor."""

    value_fn: Callable[[IdentityState], Any]
    attr_fn: Callable[[IdentityState], dict[str, Any]] | None = None


def _last_message_time(state: IdentityState) -> datetime | None:
    if state.last_message_time is None:
        return None
    return datetime.fromtimestamp(state.last_message_time, tz=timezone.utc)


STACK_SENSORS: tuple[StackSensorDescription, ...] = (
    StackSensorDescription(
        key="peers",
        translation_key="peers",
        icon="mdi:account-group",
        state_class=SensorStateClass.MEASUREMENT,
        value_fn=lambda s: len(s.peers),
        attr_fn=lambda s: {
            "peers": [
                {
                    "address": p.destination_hash,
                    "name": p.display_name,
                    "hops": p.hops,
                    "stamp_cost": p.stamp_cost,
                }
                for p in sorted(
                    s.peers.values(), key=lambda x: x.last_heard, reverse=True
                )
            ]
        },
    ),
    StackSensorDescription(
        key="rx_bytes",
        translation_key="rx_bytes",
        icon="mdi:download-network",
        device_class=SensorDeviceClass.DATA_SIZE,
        native_unit_of_measurement=UnitOfInformation.BYTES,
        suggested_unit_of_measurement=UnitOfInformation.KIBIBYTES,
        suggested_display_precision=1,
        state_class=SensorStateClass.TOTAL_INCREASING,
        entity_category=EntityCategory.DIAGNOSTIC,
        value_fn=lambda s: s.tel_rxb,
    ),
    StackSensorDescription(
        key="tx_bytes",
        translation_key="tx_bytes",
        icon="mdi:upload-network",
        device_class=SensorDeviceClass.DATA_SIZE,
        native_unit_of_measurement=UnitOfInformation.BYTES,
        suggested_unit_of_measurement=UnitOfInformation.KIBIBYTES,
        suggested_display_precision=1,
        state_class=SensorStateClass.TOTAL_INCREASING,
        entity_category=EntityCategory.DIAGNOSTIC,
        value_fn=lambda s: s.tel_txb,
    ),
    StackSensorDescription(
        key="bitrate",
        translation_key="bitrate",
        icon="mdi:speedometer",
        device_class=SensorDeviceClass.DATA_RATE,
        native_unit_of_measurement=UnitOfDataRate.BITS_PER_SECOND,
        state_class=SensorStateClass.MEASUREMENT,
        entity_category=EntityCategory.DIAGNOSTIC,
        value_fn=lambda s: s.tel_bitrate,
    ),
    StackSensorDescription(
        key="rssi",
        translation_key="rssi",
        icon="mdi:signal",
        device_class=SensorDeviceClass.SIGNAL_STRENGTH,
        native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
        state_class=SensorStateClass.MEASUREMENT,
        entity_category=EntityCategory.DIAGNOSTIC,
        entity_registry_enabled_default=False,
        value_fn=lambda s: s.tel_rssi,
    ),
    StackSensorDescription(
        key="snr",
        translation_key="snr",
        icon="mdi:signal-variant",
        native_unit_of_measurement="dB",
        state_class=SensorStateClass.MEASUREMENT,
        entity_category=EntityCategory.DIAGNOSTIC,
        entity_registry_enabled_default=False,
        value_fn=lambda s: s.tel_snr,
    ),
    StackSensorDescription(
        key="link_quality",
        translation_key="link_quality",
        icon="mdi:gauge",
        native_unit_of_measurement=PERCENTAGE,
        state_class=SensorStateClass.MEASUREMENT,
        entity_category=EntityCategory.DIAGNOSTIC,
        entity_registry_enabled_default=False,
        value_fn=lambda s: s.tel_quality,
    ),
)


IDENTITY_SENSORS: tuple[IdentitySensorDescription, ...] = (
    IdentitySensorDescription(
        key="lxmf_address",
        translation_key="lxmf_address",
        icon="mdi:identifier",
        value_fn=lambda s: s.lxmf_address,
    ),
    IdentitySensorDescription(
        key="messages_received",
        translation_key="messages_received",
        icon="mdi:inbox-arrow-down",
        state_class=SensorStateClass.TOTAL_INCREASING,
        value_fn=lambda s: s.messages_received,
    ),
    IdentitySensorDescription(
        key="messages_sent",
        translation_key="messages_sent",
        icon="mdi:inbox-arrow-up",
        state_class=SensorStateClass.TOTAL_INCREASING,
        value_fn=lambda s: s.messages_sent,
    ),
    IdentitySensorDescription(
        key="messages_failed",
        translation_key="messages_failed",
        icon="mdi:message-alert",
        state_class=SensorStateClass.TOTAL_INCREASING,
        value_fn=lambda s: s.messages_failed,
    ),
    IdentitySensorDescription(
        key="last_message",
        translation_key="last_message",
        icon="mdi:message-text",
        value_fn=lambda s: (s.last_message or "")[:255] or None,
        attr_fn=lambda s: {
            "source": s.last_message_source,
            "title": s.last_message_title,
            "full_content": s.last_message,
        },
    ),
    IdentitySensorDescription(
        key="last_message_time",
        translation_key="last_message_time",
        icon="mdi:clock-outline",
        device_class=SensorDeviceClass.TIMESTAMP,
        value_fn=_last_message_time,
    ),
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up hub + per-identity sensors."""
    runtime = hass.data[DOMAIN][entry.entry_id]

    async_add_entities(
        ReticulumStackSensor(runtime.stack, entry, desc) for desc in STACK_SENSORS
    )

    for subentry in entry.subentries.values():
        if subentry.subentry_type != SUBENTRY_TYPE_IDENTITY:
            continue
        manager = runtime.identities.get(subentry.subentry_id)
        if manager is None:
            continue
        async_add_entities(
            (
                ReticulumIdentitySensor(manager, entry, subentry, desc)
                for desc in IDENTITY_SENSORS
            ),
            config_subentry_id=subentry.subentry_id,
        )


class ReticulumStackSensor(ReticulumStackEntity, SensorEntity):
    """Hub sensor."""

    entity_description: StackSensorDescription

    def __init__(self, stack, entry, description: StackSensorDescription) -> None:
        super().__init__(stack, entry)
        self.entity_description = description
        self._attr_unique_id = f"{entry.entry_id}_{description.key}"

    @property
    def native_value(self) -> Any:
        return self.entity_description.value_fn(self._stack.state)

    @property
    def extra_state_attributes(self) -> dict[str, Any] | None:
        if self.entity_description.attr_fn is None:
            return None
        return self.entity_description.attr_fn(self._stack.state)


class ReticulumIdentitySensor(ReticulumIdentityEntity, SensorEntity):
    """Per-identity sensor."""

    entity_description: IdentitySensorDescription

    def __init__(
        self, manager, entry, subentry, description: IdentitySensorDescription
    ) -> None:
        super().__init__(manager, entry, subentry)
        self.entity_description = description
        self._attr_unique_id = f"{subentry.subentry_id}_{description.key}"

    @property
    def native_value(self) -> Any:
        return self.entity_description.value_fn(self._manager.state)

    @property
    def extra_state_attributes(self) -> dict[str, Any] | None:
        if self.entity_description.attr_fn is None:
            return None
        return self.entity_description.attr_fn(self._manager.state)