"""Support for Launch Library sensors.""" from __future__ import annotations from collections.abc import Callable from dataclasses import dataclass from datetime import datetime from typing import Any from pylaunches.types import Event, Launch from homeassistant.components.sensor import ( SensorDeviceClass, SensorEntity, SensorEntityDescription, ) from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_NAME, PERCENTAGE from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import ( CoordinatorEntity, DataUpdateCoordinator, ) from homeassistant.util.dt import parse_datetime from . import LaunchLibraryData from .const import DOMAIN DEFAULT_NEXT_LAUNCH_NAME = "Next launch" @dataclass(frozen=True, kw_only=True) class LaunchLibrarySensorEntityDescription(SensorEntityDescription): """Describes a Next Launch sensor entity.""" value_fn: Callable[[Launch | Event], datetime | int | str | None] attributes_fn: Callable[[Launch | Event], dict[str, Any] | None] SENSOR_DESCRIPTIONS: tuple[LaunchLibrarySensorEntityDescription, ...] = ( LaunchLibrarySensorEntityDescription( key="next_launch", icon="mdi:rocket-launch", translation_key="next_launch", value_fn=lambda nl: nl["name"], attributes_fn=lambda nl: { "provider": nl["launch_service_provider"]["name"], "pad": nl["pad"]["name"], "facility": nl["pad"]["location"]["name"], "provider_country_code": nl["pad"]["location"]["country_code"], }, ), LaunchLibrarySensorEntityDescription( key="launch_time", icon="mdi:clock-outline", translation_key="launch_time", device_class=SensorDeviceClass.TIMESTAMP, value_fn=lambda nl: parse_datetime(nl["net"]), attributes_fn=lambda nl: { "window_start": nl["window_start"], "window_end": nl["window_end"], "stream_live": nl["window_start"], }, ), LaunchLibrarySensorEntityDescription( key="launch_probability", icon="mdi:dice-multiple", translation_key="launch_probability", native_unit_of_measurement=PERCENTAGE, value_fn=lambda nl: None if nl["probability"] == -1 else nl["probability"], attributes_fn=lambda nl: None, ), LaunchLibrarySensorEntityDescription( key="launch_status", icon="mdi:rocket-launch", translation_key="launch_status", value_fn=lambda nl: nl["status"]["name"], attributes_fn=lambda nl: {"reason": nl.get("holdreason")}, ), LaunchLibrarySensorEntityDescription( key="launch_mission", icon="mdi:orbit", translation_key="launch_mission", value_fn=lambda nl: nl["mission"]["name"], attributes_fn=lambda nl: { "mission_type": nl["mission"]["type"], "target_orbit": nl["mission"]["orbit"]["name"], "description": nl["mission"]["description"], }, ), LaunchLibrarySensorEntityDescription( key="starship_launch", icon="mdi:rocket", translation_key="starship_launch", device_class=SensorDeviceClass.TIMESTAMP, value_fn=lambda sl: parse_datetime(sl["net"]), attributes_fn=lambda sl: { "title": sl["mission"]["name"], "status": sl["status"]["name"], "target_orbit": sl["mission"]["orbit"]["name"], "description": sl["mission"]["description"], }, ), LaunchLibrarySensorEntityDescription( key="starship_event", icon="mdi:calendar", translation_key="starship_event", device_class=SensorDeviceClass.TIMESTAMP, value_fn=lambda se: parse_datetime(se["date"]), attributes_fn=lambda se: { "title": se["name"], "location": se["location"], "stream": se["video_url"], "description": se["description"], }, ), ) async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the sensor platform.""" name = entry.data.get(CONF_NAME, DEFAULT_NEXT_LAUNCH_NAME) coordinator: DataUpdateCoordinator[LaunchLibraryData] = hass.data[DOMAIN] async_add_entities( LaunchLibrarySensor( coordinator=coordinator, entry_id=entry.entry_id, description=description, name=name, ) for description in SENSOR_DESCRIPTIONS ) class LaunchLibrarySensor( CoordinatorEntity[DataUpdateCoordinator[LaunchLibraryData]], SensorEntity ): """Representation of the next launch sensors.""" _attr_attribution = "Data provided by Launch Library." _attr_has_entity_name = True _next_event: Launch | Event | None = None entity_description: LaunchLibrarySensorEntityDescription def __init__( self, coordinator: DataUpdateCoordinator[LaunchLibraryData], entry_id: str, description: LaunchLibrarySensorEntityDescription, name: str, ) -> None: """Initialize a Launch Library sensor.""" super().__init__(coordinator) self._attr_unique_id = f"{entry_id}_{description.key}" self.entity_description = description self._attr_device_info = DeviceInfo( identifiers={(DOMAIN, entry_id)}, entry_type=DeviceEntryType.SERVICE, name=name, ) @property def native_value(self) -> datetime | str | int | None: """Return the state of the sensor.""" if self._next_event is None: return None return self.entity_description.value_fn(self._next_event) @property def extra_state_attributes(self) -> dict[str, Any] | None: """Return the attributes of the sensor.""" if self._next_event is None: return None return self.entity_description.attributes_fn(self._next_event) @property def available(self) -> bool: """Return if the sensor is available.""" return super().available and self._next_event is not None @callback def _handle_coordinator_update(self) -> None: """Handle updated data from the coordinator.""" if self.entity_description.key == "starship_launch": events = self.coordinator.data["starship_events"]["upcoming"]["launches"] elif self.entity_description.key == "starship_event": events = self.coordinator.data["starship_events"]["upcoming"]["events"] else: events = self.coordinator.data["upcoming_launches"] self._next_event = next((event for event in (events)), None) super()._handle_coordinator_update() async def async_added_to_hass(self) -> None: """When entity is added to hass.""" await super().async_added_to_hass() self._handle_coordinator_update()