core/homeassistant/components/weatherflow_cloud/sensor.py

"""Sensors for cloud based weatherflow."""

from abc import ABC
from collections.abc import Callable
from dataclasses import dataclass
from datetime import date, datetime
from decimal import Decimal

from weatherflow4py.models.rest.observation import Observation
from weatherflow4py.models.ws.websocket_response import (
    EventDataRapidWind,
    WebsocketObservation,
)

from homeassistant.components.sensor import (
    SensorDeviceClass,
    SensorEntity,
    SensorEntityDescription,
    SensorStateClass,
)
from homeassistant.const import (
    LIGHT_LUX,
    PERCENTAGE,
    UV_INDEX,
    EntityCategory,
    UnitOfIrradiance,
    UnitOfLength,
    UnitOfMass,
    UnitOfPrecipitationDepth,
    UnitOfPressure,
    UnitOfSpeed,
    UnitOfTemperature,
    UnitOfTime,
    UnitOfVolume,
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
from homeassistant.helpers.typing import StateType
from homeassistant.util.dt import UTC

from .coordinator import (
    WeatherFlowCloudConfigEntry,
    WeatherFlowCloudUpdateCoordinatorREST,
    WeatherFlowObservationCoordinator,
    WeatherFlowWindCoordinator,
)
from .entity import WeatherFlowCloudEntity

PRECIPITATION_TYPE = {
    0: "none",
    1: "rain",
    2: "snow",
    3: "sleet",
    4: "storm",
}


@dataclass(frozen=True, kw_only=True)
class WeatherFlowCloudSensorEntityDescription(
    SensorEntityDescription,
):
    """Describes a weatherflow sensor."""

    value_fn: Callable[[Observation], StateType | datetime]


@dataclass(frozen=True, kw_only=True)
class WeatherFlowCloudSensorEntityDescriptionWebsocketWind(
    SensorEntityDescription,
):
    """Describes a weatherflow sensor."""

    value_fn: Callable[[EventDataRapidWind], StateType | datetime]


@dataclass(frozen=True, kw_only=True)
class WeatherFlowCloudSensorEntityDescriptionWebsocketObservation(
    SensorEntityDescription,
):
    """Describes a weatherflow sensor."""

    value_fn: Callable[[WebsocketObservation], StateType | datetime]


WEBSOCKET_WIND_SENSORS: tuple[
    WeatherFlowCloudSensorEntityDescriptionWebsocketWind, ...
] = (
    WeatherFlowCloudSensorEntityDescriptionWebsocketWind(
        key="wind_speed",
        state_class=SensorStateClass.MEASUREMENT,
        device_class=SensorDeviceClass.WIND_SPEED,
        suggested_display_precision=1,
        value_fn=lambda data: data.wind_speed_meters_per_second,
        native_unit_of_measurement=UnitOfSpeed.METERS_PER_SECOND,
    ),
    WeatherFlowCloudSensorEntityDescriptionWebsocketWind(
        key="wind_direction",
        device_class=SensorDeviceClass.WIND_DIRECTION,
        translation_key="wind_direction",
        value_fn=lambda data: data.wind_direction_degrees,
        native_unit_of_measurement="°",
    ),
)

WEBSOCKET_OBSERVATION_SENSORS: tuple[
    WeatherFlowCloudSensorEntityDescriptionWebsocketObservation, ...
] = (
    WeatherFlowCloudSensorEntityDescriptionWebsocketObservation(
        key="wind_lull",
        translation_key="wind_lull",
        state_class=SensorStateClass.MEASUREMENT,
        device_class=SensorDeviceClass.WIND_SPEED,
        suggested_display_precision=1,
        value_fn=lambda data: data.wind_lull,
        native_unit_of_measurement=UnitOfSpeed.METERS_PER_SECOND,
    ),
    WeatherFlowCloudSensorEntityDescriptionWebsocketObservation(
        key="wind_gust",
        translation_key="wind_gust",
        state_class=SensorStateClass.MEASUREMENT,
        device_class=SensorDeviceClass.WIND_SPEED,
        suggested_display_precision=1,
        value_fn=lambda data: data.wind_gust,
        native_unit_of_measurement=UnitOfSpeed.METERS_PER_SECOND,
    ),
    WeatherFlowCloudSensorEntityDescriptionWebsocketObservation(
        key="wind_avg",
        translation_key="wind_avg",
        state_class=SensorStateClass.MEASUREMENT,
        device_class=SensorDeviceClass.WIND_SPEED,
        suggested_display_precision=1,
        value_fn=lambda data: data.wind_avg,
        native_unit_of_measurement=UnitOfSpeed.METERS_PER_SECOND,
    ),
    WeatherFlowCloudSensorEntityDescriptionWebsocketObservation(
        key="wind_sample_interval",
        translation_key="wind_sample_interval",
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        native_unit_of_measurement=UnitOfTime.SECONDS,
        entity_category=EntityCategory.DIAGNOSTIC,
        entity_registry_enabled_default=False,
        value_fn=lambda data: data.wind_sample_interval,
    ),
)


WF_SENSORS: tuple[WeatherFlowCloudSensorEntityDescription, ...] = (
    # Air Sensors
    WeatherFlowCloudSensorEntityDescription(
        key="air_density",
        translation_key="air_density",
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=5,
        value_fn=lambda data: data.air_density,
        native_unit_of_measurement=f"{UnitOfMass.KILOGRAMS}/{UnitOfVolume.CUBIC_METERS}",
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="relative_humidity",
        translation_key="relative_humidity",
        device_class=SensorDeviceClass.HUMIDITY,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=0,
        value_fn=lambda data: data.relative_humidity,
        native_unit_of_measurement=PERCENTAGE,
    ),
    # Light Sensors
    WeatherFlowCloudSensorEntityDescription(
        key="brightness",
        translation_key="illuminance",
        device_class=SensorDeviceClass.ILLUMINANCE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=0,
        value_fn=lambda data: data.brightness,
        native_unit_of_measurement=LIGHT_LUX,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="uv",
        translation_key="uv_index",
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=2,
        value_fn=lambda data: data.uv,
        native_unit_of_measurement=UV_INDEX,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="solar_radiation",
        translation_key="solar_radiation",
        device_class=SensorDeviceClass.IRRADIANCE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=0,
        value_fn=lambda data: data.solar_radiation,
        native_unit_of_measurement=UnitOfIrradiance.WATTS_PER_SQUARE_METER,
    ),
    # Temp Sensors
    WeatherFlowCloudSensorEntityDescription(
        key="air_temperature",
        translation_key="air_temperature",
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.air_temperature,
        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="dew_point",
        translation_key="dew_point",
        value_fn=lambda data: data.dew_point,
        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="feels_like",
        translation_key="feels_like",
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.feels_like,
        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="heat_index",
        translation_key="heat_index",
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.heat_index,
        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="wind_chill",
        translation_key="wind_chill",
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.wind_chill,
        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="wet_bulb_temperature",
        translation_key="wet_bulb_temperature",
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.wet_bulb_temperature,
        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="wet_bulb_globe_temperature",
        translation_key="wet_bulb_globe_temperature",
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.wet_bulb_globe_temperature,
        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
    ),
    # Pressure Sensors
    WeatherFlowCloudSensorEntityDescription(
        key="barometric_pressure",
        translation_key="barometric_pressure",
        value_fn=lambda data: data.barometric_pressure,
        native_unit_of_measurement=UnitOfPressure.MBAR,
        device_class=SensorDeviceClass.ATMOSPHERIC_PRESSURE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=3,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="sea_level_pressure",
        translation_key="sea_level_pressure",
        value_fn=lambda data: data.sea_level_pressure,
        native_unit_of_measurement=UnitOfPressure.MBAR,
        device_class=SensorDeviceClass.ATMOSPHERIC_PRESSURE,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=3,
    ),
    # Rain Sensors
    WeatherFlowCloudSensorEntityDescription(
        key="precip_accum_last_1hr",
        translation_key="precip_accum_last_1hr",
        device_class=SensorDeviceClass.PRECIPITATION,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.precip_accum_last_1hr,
        native_unit_of_measurement=UnitOfPrecipitationDepth.MILLIMETERS,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="precip_accum_local_day",
        translation_key="precip_accum_local_day",
        device_class=SensorDeviceClass.PRECIPITATION,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.precip_accum_local_day,
        native_unit_of_measurement=UnitOfPrecipitationDepth.MILLIMETERS,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="precip_accum_local_day_final",
        translation_key="precip_accum_local_day_final",
        device_class=SensorDeviceClass.PRECIPITATION,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.precip_accum_local_day_final,
        native_unit_of_measurement=UnitOfPrecipitationDepth.MILLIMETERS,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="precip_accum_local_yesterday",
        translation_key="precip_accum_local_yesterday",
        device_class=SensorDeviceClass.PRECIPITATION,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.precip_accum_local_yesterday,
        native_unit_of_measurement=UnitOfPrecipitationDepth.MILLIMETERS,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="precip_accum_local_yesterday_final",
        translation_key="precip_accum_local_yesterday_final",
        device_class=SensorDeviceClass.PRECIPITATION,
        state_class=SensorStateClass.MEASUREMENT,
        suggested_display_precision=1,
        value_fn=lambda data: data.precip_accum_local_yesterday_final,
        native_unit_of_measurement=UnitOfPrecipitationDepth.MILLIMETERS,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="precip_analysis_type_yesterday",
        translation_key="precip_analysis_type_yesterday",
        device_class=SensorDeviceClass.ENUM,
        options=["none", "rain", "snow", "sleet", "storm"],
        suggested_display_precision=1,
        value_fn=lambda data: PRECIPITATION_TYPE.get(
            data.precip_analysis_type_yesterday
        ),
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="precip_minutes_local_day",
        translation_key="precip_minutes_local_day",
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=UnitOfTime.MINUTES,
        suggested_display_precision=1,
        value_fn=lambda data: data.precip_minutes_local_day,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="precip_minutes_local_yesterday",
        translation_key="precip_minutes_local_yesterday",
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=UnitOfTime.MINUTES,
        suggested_display_precision=1,
        value_fn=lambda data: data.precip_minutes_local_yesterday,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="precip_minutes_local_yesterday_final",
        translation_key="precip_minutes_local_yesterday_final",
        state_class=SensorStateClass.MEASUREMENT,
        native_unit_of_measurement=UnitOfTime.MINUTES,
        suggested_display_precision=1,
        value_fn=lambda data: data.precip_minutes_local_yesterday_final,
    ),
    # Lightning Sensors
    WeatherFlowCloudSensorEntityDescription(
        key="lightning_strike_count",
        translation_key="lightning_strike_count",
        state_class=SensorStateClass.TOTAL,
        value_fn=lambda data: data.lightning_strike_count,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="lightning_strike_count_last_1hr",
        translation_key="lightning_strike_count_last_1hr",
        state_class=SensorStateClass.TOTAL,
        value_fn=lambda data: data.lightning_strike_count_last_1hr,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="lightning_strike_count_last_3hr",
        translation_key="lightning_strike_count_last_3hr",
        state_class=SensorStateClass.TOTAL,
        value_fn=lambda data: data.lightning_strike_count_last_3hr,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="lightning_strike_last_distance",
        translation_key="lightning_strike_last_distance",
        state_class=SensorStateClass.MEASUREMENT,
        device_class=SensorDeviceClass.DISTANCE,
        native_unit_of_measurement=UnitOfLength.KILOMETERS,
        value_fn=lambda data: data.lightning_strike_last_distance,
    ),
    WeatherFlowCloudSensorEntityDescription(
        key="lightning_strike_last_epoch",
        translation_key="lightning_strike_last_epoch",
        device_class=SensorDeviceClass.TIMESTAMP,
        value_fn=(
            lambda data: (
                datetime.fromtimestamp(data.lightning_strike_last_epoch, tz=UTC)
                if data.lightning_strike_last_epoch is not None
                else None
            )
        ),
    ),
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: WeatherFlowCloudConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up WeatherFlow sensors based on a config entry."""

    coordinators = entry.runtime_data
    rest_coordinator = coordinators.rest
    wind_coordinator = coordinators.wind
    observation_coordinator = coordinators.observation

    entities: list[SensorEntity] = [
        WeatherFlowCloudSensorREST(rest_coordinator, sensor_description, station_id)
        for station_id in rest_coordinator.data
        for sensor_description in WF_SENSORS
    ]

    entities.extend(
        WeatherFlowWebsocketSensorWind(
            coordinator=wind_coordinator,
            description=sensor_description,
            station_id=station_id,
            device_id=device_id,
        )
        for station_id in wind_coordinator.stations.station_outdoor_device_map
        for device_id in wind_coordinator.stations.station_outdoor_device_map[
            station_id
        ]
        for sensor_description in WEBSOCKET_WIND_SENSORS
    )

    entities.extend(
        WeatherFlowWebsocketSensorObservation(
            coordinator=observation_coordinator,
            description=sensor_description,
            station_id=station_id,
            device_id=device_id,
        )
        for station_id in observation_coordinator.stations.station_outdoor_device_map
        for device_id in observation_coordinator.stations.station_outdoor_device_map[
            station_id
        ]
        for sensor_description in WEBSOCKET_OBSERVATION_SENSORS
    )
    async_add_entities(entities)


class WeatherFlowSensorBase(WeatherFlowCloudEntity, SensorEntity, ABC):
    """Common base class."""

    def __init__(
        self,
        coordinator: (
            WeatherFlowCloudUpdateCoordinatorREST
            | WeatherFlowWindCoordinator
            | WeatherFlowObservationCoordinator
        ),
        description: (
            WeatherFlowCloudSensorEntityDescription
            | WeatherFlowCloudSensorEntityDescriptionWebsocketWind
            | WeatherFlowCloudSensorEntityDescriptionWebsocketObservation
        ),
        station_id: int,
        device_id: int | None = None,
    ) -> None:
        """Initialize a sensor."""
        super().__init__(coordinator, station_id)
        self.station_id = station_id
        self.device_id = device_id
        self.entity_description = description
        self._attr_unique_id = self._generate_unique_id()

    def _generate_unique_id(self) -> str:
        """Generate a unique ID for the sensor."""
        if self.device_id is not None:
            return f"{self.station_id}_{self.device_id}_{self.entity_description.key}"
        return f"{self.station_id}_{self.entity_description.key}"

    @property
    def available(self) -> bool:
        """Get if available."""

        if not super().available:
            return False

        if self.device_id is not None:
            # Websocket sensors - have Device IDs
            return bool(
                self.coordinator.data
                and self.coordinator.data[self.station_id][self.device_id] is not None
            )

        return True


class WeatherFlowWebsocketSensorObservation(WeatherFlowSensorBase):
    """Class for Websocket Observations."""

    entity_description: WeatherFlowCloudSensorEntityDescriptionWebsocketObservation

    @property
    def native_value(self) -> StateType | date | datetime | Decimal:
        """Return the native value."""
        data = self.coordinator.data[self.station_id][self.device_id]
        return self.entity_description.value_fn(data)


class WeatherFlowWebsocketSensorWind(WeatherFlowSensorBase):
    """Class for wind over websockets."""

    entity_description: WeatherFlowCloudSensorEntityDescriptionWebsocketWind

    @property
    def native_value(self) -> StateType | datetime:
        """Return the native value."""

        # This data is often invalid at starutp.
        if self.coordinator.data is not None:
            data = self.coordinator.data[self.station_id][self.device_id]
            return self.entity_description.value_fn(data)
        return None


class WeatherFlowCloudSensorREST(WeatherFlowSensorBase):
    """Class for a REST based sensor."""

    entity_description: WeatherFlowCloudSensorEntityDescription

    coordinator: WeatherFlowCloudUpdateCoordinatorREST

    @property
    def _observation(self) -> Observation | None:
        """Return the current station observation."""
        observations = self.coordinator.data[self.station_id].observation.obs
        if not observations:
            return None
        return observations[0]

    @property
    def available(self) -> bool:
        """Get if available."""
        return super().available and self._observation is not None

    @property
    def native_value(self) -> StateType | datetime:
        """Return the native value."""
        if (observation := self._observation) is None:
            return None
        return self.entity_description.value_fn(observation)