ha-babytracker/custom_components/babytracker/sensor.py

"""Sensors summarizing a baby's log."""

from __future__ import annotations

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

from homeassistant.components.sensor import (
    SensorDeviceClass,
    SensorEntity,
    SensorEntityDescription,
    SensorStateClass,
)
from homeassistant.const import UnitOfLength, UnitOfMass, UnitOfTemperature, UnitOfTime, UnitOfVolume
from homeassistant.util import dt as dt_util

from .babytracker_api import BOTTLE_KINDS, FEED_KINDS, Baby, Entry, Replica
from .const import RECENT_ENTRIES
from .coordinator import BabyTrackerCoordinator
from .entity import BabyEntity, async_add_baby_entities


def _today(replica: Replica, baby_id: str, kinds: frozenset[str]) -> list[Entry]:
    # An entry counts for the day it started, so a night's sleep counts for one day.
    return replica.since(baby_id, dt_util.start_of_local_day(), kinds)


def _latest_measure(replica: Replica, baby_id: str, kind: str, key: str) -> float | None:
    # Growth entries often hold only some measurements.
    for entry in replica.entries(baby_id, {kind}):
        if entry.details.get(key) is not None:
            return entry.details[key]
    return None


def _entry_attributes(entry: Entry | None) -> dict[str, Any]:
    if entry is None:
        return {}
    return {"type": entry.kind, "summary": entry.summary, "note": entry.note, "end": entry.end, **entry.details}


@dataclass(frozen=True, kw_only=True)
class BabySensorDescription(SensorEntityDescription):
    value_fn: Callable[[Replica, str], datetime | float | None]
    entry_fn: Callable[[Replica, str], Entry | None] | None = None


def _last(kinds: frozenset[str]) -> dict[str, Any]:
    def entry_fn(replica: Replica, baby_id: str) -> Entry | None:
        return replica.latest(baby_id, kinds)

    def value_fn(replica: Replica, baby_id: str) -> datetime | None:
        entry = entry_fn(replica, baby_id)
        return entry.time if entry else None

    return {"device_class": SensorDeviceClass.TIMESTAMP, "value_fn": value_fn, "entry_fn": entry_fn}


def _count_today(kinds: frozenset[str]) -> Callable[[Replica, str], int]:
    return lambda replica, baby_id: len(_today(replica, baby_id, kinds))


DIAPER = frozenset({"Diaper"})
SLEEP = frozenset({"Sleep"})

SENSORS = (
    BabySensorDescription(key="last_diaper", **_last(DIAPER)),
    BabySensorDescription(key="last_feeding", **_last(FEED_KINDS)),
    BabySensorDescription(key="last_sleep", **_last(SLEEP)),
    BabySensorDescription(
        key="diapers_today",
        state_class=SensorStateClass.TOTAL_INCREASING,
        value_fn=_count_today(DIAPER),
    ),
    BabySensorDescription(
        key="feedings_today",
        state_class=SensorStateClass.TOTAL_INCREASING,
        value_fn=_count_today(FEED_KINDS),
    ),
    BabySensorDescription(
        key="sleep_today",
        device_class=SensorDeviceClass.DURATION,
        state_class=SensorStateClass.TOTAL_INCREASING,
        native_unit_of_measurement=UnitOfTime.MINUTES,
        value_fn=lambda replica, baby_id: sum(
            entry.details.get("duration") or 0 for entry in _today(replica, baby_id, SLEEP)
        ),
    ),
    BabySensorDescription(
        key="bottle_volume_today",
        device_class=SensorDeviceClass.VOLUME,
        state_class=SensorStateClass.TOTAL_INCREASING,
        native_unit_of_measurement=UnitOfVolume.MILLILITERS,
        suggested_display_precision=0,
        value_fn=lambda replica, baby_id: sum(
            entry.details.get("volume_ml") or 0 for entry in _today(replica, baby_id, BOTTLE_KINDS)
        ),
    ),
    BabySensorDescription(
        key="temperature",
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
        suggested_display_precision=1,
        value_fn=lambda replica, baby_id: _latest_measure(replica, baby_id, "Temperature", "celsius"),
    ),
    BabySensorDescription(
        key="weight",
        device_class=SensorDeviceClass.WEIGHT,
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=UnitOfMass.KILOGRAMS,
        suggested_display_precision=2,
        value_fn=lambda replica, baby_id: _latest_measure(replica, baby_id, "Growth", "weight_kg"),
    ),
    BabySensorDescription(
        key="length",
        device_class=SensorDeviceClass.DISTANCE,
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=UnitOfLength.CENTIMETERS,
        suggested_display_precision=1,
        value_fn=lambda replica, baby_id: _latest_measure(replica, baby_id, "Growth", "length_cm"),
    ),
    BabySensorDescription(
        key="head_circumference",
        device_class=SensorDeviceClass.DISTANCE,
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=UnitOfLength.CENTIMETERS,
        suggested_display_precision=1,
        value_fn=lambda replica, baby_id: _latest_measure(replica, baby_id, "Growth", "head_cm"),
    ),
)


class BabySensor(BabyEntity, SensorEntity):
    entity_description: BabySensorDescription

    def __init__(self, coordinator: BabyTrackerCoordinator, baby: Baby, description: BabySensorDescription) -> None:
        super().__init__(coordinator, baby, description.key)
        self.entity_description = description

    @property
    def native_value(self) -> datetime | float | None:
        return self.entity_description.value_fn(self.coordinator.data, self.baby_id)

    @property
    def extra_state_attributes(self) -> dict[str, Any] | None:
        entry_fn = self.entity_description.entry_fn
        return _entry_attributes(entry_fn(self.coordinator.data, self.baby_id)) if entry_fn else None


class RecentActivitySensor(BabyEntity, SensorEntity):
    """The latest entries of every type, for showing a log on a dashboard."""

    _attr_device_class = SensorDeviceClass.TIMESTAMP
    # A long list that changes constantly does not belong in the history database.
    _unrecorded_attributes = frozenset({"entries"})

    def __init__(self, coordinator: BabyTrackerCoordinator, baby: Baby) -> None:
        super().__init__(coordinator, baby, "recent_activity")

    @property
    def native_value(self) -> datetime | None:
        entry = self.coordinator.data.latest(self.baby_id)
        return entry.time if entry else None

    @property
    def extra_state_attributes(self) -> dict[str, Any]:
        entries = []
        for entry in self.coordinator.data.entries(self.baby_id):
            if len(entries) == RECENT_ENTRIES:
                break
            entries.append(
                {
                    "id": entry.id,
                    "type": entry.kind,
                    "summary": entry.summary,
                    "time": entry.time.isoformat(),
                    "note": entry.note,
                }
            )
        return {"entries": entries}


async def async_setup_entry(hass, entry, async_add_entities) -> None:
    coordinator: BabyTrackerCoordinator = entry.runtime_data

    def create(baby: Baby) -> list[SensorEntity]:
        sensors: list[SensorEntity] = [BabySensor(coordinator, baby, description) for description in SENSORS]
        sensors.append(RecentActivitySensor(coordinator, baby))
        return sensors

    async_add_baby_entities(entry, coordinator, async_add_entities, create)