ha-xiaoxiang-bms/custom_components/xiaoxiang_bms/sensor.py

"""Sensor entities for the Xiaoxiang Smart BMS integration."""
from __future__ import annotations

from homeassistant.components.sensor import (
    SensorDeviceClass,
    SensorEntity,
    SensorEntityDescription,
    SensorStateClass,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import (
    PERCENTAGE,
    UnitOfElectricCurrent,
    UnitOfElectricPotential,
    UnitOfPower,
    UnitOfTemperature,
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity

from .const import CONF_ADDRESS, DOMAIN
from .coordinator import BmsCoordinator

# ---------------------------------------------------------------------------
# Static sensor definitions — one entity per key in coordinator.data
# ---------------------------------------------------------------------------

STATIC_SENSORS: tuple[SensorEntityDescription, ...] = (
    SensorEntityDescription(
        key="voltage",
        name="Voltage",
        native_unit_of_measurement=UnitOfElectricPotential.VOLT,
        device_class=SensorDeviceClass.VOLTAGE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=2,
    ),
    SensorEntityDescription(
        key="current",
        name="Current",
        native_unit_of_measurement=UnitOfElectricCurrent.AMPERE,
        device_class=SensorDeviceClass.CURRENT,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=2,
    ),
    SensorEntityDescription(
        key="power",
        name="Power",
        native_unit_of_measurement=UnitOfPower.WATT,
        device_class=SensorDeviceClass.POWER,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
    ),
    SensorEntityDescription(
        key="state_of_charge",
        name="State of Charge",
        native_unit_of_measurement=PERCENTAGE,
        device_class=SensorDeviceClass.BATTERY,
        state_class=SensorStateClass.MEASUREMENT,
    ),
    SensorEntityDescription(
        key="residual_capacity",
        name="Remaining Capacity",
        native_unit_of_measurement="Ah",
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=2,
        icon="mdi:battery-charging",
    ),
    SensorEntityDescription(
        key="nominal_capacity",
        name="Nominal Capacity",
        native_unit_of_measurement="Ah",
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=2,
        icon="mdi:battery",
    ),
    SensorEntityDescription(
        key="cycle_count",
        name="Cycle Count",
        state_class=SensorStateClass.TOTAL_INCREASING,
        icon="mdi:battery-sync",
    ),
    SensorEntityDescription(
        key="cell_delta",
        name="Cell Voltage Delta",
        native_unit_of_measurement="mV",
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        icon="mdi:battery-alert-variant",
    ),
)


# ---------------------------------------------------------------------------
# Platform setup
# ---------------------------------------------------------------------------

async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Create all BMS sensor entities after the first successful poll."""
    coordinator: BmsCoordinator = hass.data[DOMAIN][entry.entry_id]
    address: str = entry.data[CONF_ADDRESS]

    entities: list[SensorEntity] = [
        BmsStaticSensor(coordinator, description, address)
        for description in STATIC_SENSORS
    ]

    # Temperature sensors — count determined from first poll
    for i in range(len(coordinator.data.get("temperatures", []))):
        entities.append(
            BmsDynamicSensor(
                coordinator=coordinator,
                address=address,
                unique_key=f"temperature_{i + 1}",
                friendly_name=f"Temperature {i + 1}",
                data_key="temperatures",
                index=i,
                unit=UnitOfTemperature.CELSIUS,
                device_class=SensorDeviceClass.TEMPERATURE,
                icon=None,
            )
        )

    # Cell voltage sensors — count determined from first poll
    for i in range(len(coordinator.data.get("cell_voltages", []))):
        entities.append(
            BmsDynamicSensor(
                coordinator=coordinator,
                address=address,
                unique_key=f"cell_voltage_{i + 1}",
                friendly_name=f"Cell {i + 1} Voltage",
                data_key="cell_voltages",
                index=i,
                unit=UnitOfElectricPotential.VOLT,
                device_class=SensorDeviceClass.VOLTAGE,
                icon=None,
            )
        )

    async_add_entities(entities)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _device_info(address: str) -> DeviceInfo:
    return DeviceInfo(
        identifiers={(DOMAIN, address)},
        name="Xiaoxiang Smart BMS",
        manufacturer="Xiaoxiang",
        model="Smart BMS",
    )


# ---------------------------------------------------------------------------
# Entity classes
# ---------------------------------------------------------------------------

class BmsStaticSensor(CoordinatorEntity[BmsCoordinator], SensorEntity):
    """A sensor whose key maps directly to a scalar value in coordinator.data."""

    _attr_has_entity_name = True

    def __init__(
        self,
        coordinator: BmsCoordinator,
        description: SensorEntityDescription,
        address: str,
    ) -> None:
        super().__init__(coordinator)
        self.entity_description = description
        self._attr_unique_id = f"{address}_{description.key}"
        self._attr_device_info = _device_info(address)

    @property
    def native_value(self):
        return self.coordinator.data.get(self.entity_description.key)


class BmsDynamicSensor(CoordinatorEntity[BmsCoordinator], SensorEntity):
    """A sensor that reads an element from a list inside coordinator.data.

    Used for per-cell voltages and per-probe temperatures, whose count is
    only known after the first successful BMS poll.
    """

    _attr_has_entity_name = True
    _attr_state_class = SensorStateClass.MEASUREMENT

    def __init__(
        self,
        coordinator: BmsCoordinator,
        address: str,
        unique_key: str,
        friendly_name: str,
        data_key: str,
        index: int,
        unit: str,
        device_class: SensorDeviceClass | None,
        icon: str | None,
    ) -> None:
        super().__init__(coordinator)
        self._data_key = data_key
        self._index = index
        self._attr_unique_id = f"{address}_{unique_key}"
        self._attr_name = friendly_name
        self._attr_native_unit_of_measurement = unit
        self._attr_device_class = device_class
        self._attr_device_info = _device_info(address)
        if icon:
            self._attr_icon = icon
        if device_class == SensorDeviceClass.VOLTAGE:
            self._attr_suggested_display_precision = 3
        elif device_class == SensorDeviceClass.TEMPERATURE:
            self._attr_suggested_display_precision = 1

    @property
    def native_value(self):
        values: list = self.coordinator.data.get(self._data_key, [])
        if self._index < len(values):
            return values[self._index]
        return None