"""Get ride details and liveboard details for NMBS (Belgian railway).""" from datetime import datetime import logging from typing import Any from pyrail import iRail from pyrail.models import ConnectionDetails, LiveboardDeparture, StationDetails from homeassistant.components.sensor import SensorEntity from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( ATTR_LATITUDE, ATTR_LONGITUDE, CONF_NAME, CONF_SHOW_ON_MAP, UnitOfTime, ) from homeassistant.core import HomeAssistant from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.util import dt as dt_util from .const import ( # noqa: F401 CONF_EXCLUDE_VIAS, CONF_STATION_FROM, CONF_STATION_LIVE, CONF_STATION_TO, DOMAIN, PLATFORMS, find_station, find_station_by_name, ) _LOGGER = logging.getLogger(__name__) DEFAULT_ICON = "mdi:train" DEFAULT_ICON_ALERT = "mdi:alert-octagon" def get_time_until(departure_time: datetime | None = None): """Calculate the time between now and a train's departure time.""" if departure_time is None: return 0 delta = dt_util.as_utc(departure_time) - dt_util.utcnow() return round(delta.total_seconds() / 60) def get_delay_in_minutes(delay=0): """Get the delay in minutes from a delay in seconds.""" return round(int(delay) / 60) def get_ride_duration(departure_time: datetime, arrival_time: datetime, delay=0): """Calculate the total travel time in minutes.""" duration = arrival_time - departure_time duration_time = round(duration.total_seconds() / 60) return duration_time + get_delay_in_minutes(delay) async def async_setup_entry( hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up NMBS sensor entities based on a config entry.""" api_client = iRail(session=async_get_clientsession(hass)) name = config_entry.data.get(CONF_NAME, None) show_on_map = config_entry.data.get(CONF_SHOW_ON_MAP, False) excl_vias = config_entry.data.get(CONF_EXCLUDE_VIAS, False) station_from = find_station(hass, config_entry.data[CONF_STATION_FROM]) station_to = find_station(hass, config_entry.data[CONF_STATION_TO]) # setup the connection from station to station # setup a disabled liveboard for both from and to station async_add_entities( [ NMBSSensor( api_client, name, show_on_map, station_from, station_to, excl_vias ), NMBSLiveBoard( api_client, station_from, station_from, station_to, excl_vias ), NMBSLiveBoard(api_client, station_to, station_from, station_to, excl_vias), ] ) class NMBSLiveBoard(SensorEntity): """Get the next train from a station's liveboard.""" _attr_attribution = "https://api.irail.be/" def __init__( self, api_client: iRail, live_station: StationDetails, station_from: StationDetails, station_to: StationDetails, excl_vias: bool, ) -> None: """Initialize the sensor for getting liveboard data.""" self._station = live_station self._api_client = api_client self._station_from = station_from self._station_to = station_to self._excl_vias = excl_vias self._attrs: LiveboardDeparture | None = None self._state: str | None = None self.entity_registry_enabled_default = False @property def name(self) -> str: """Return the sensor default name.""" return f"Trains in {self._station.standard_name}" @property def unique_id(self) -> str: """Return the unique ID.""" unique_id = f"{self._station.id}_{self._station_from.id}_{self._station_to.id}" vias = "_excl_vias" if self._excl_vias else "" return f"nmbs_live_{unique_id}{vias}" @property def icon(self) -> str: """Return the default icon or an alert icon if delays.""" if self._attrs and int(self._attrs.delay) > 0: return DEFAULT_ICON_ALERT return DEFAULT_ICON @property def native_value(self) -> str | None: """Return sensor state.""" return self._state @property def extra_state_attributes(self) -> dict[str, Any] | None: """Return the sensor attributes if data is available.""" if self._state is None or not self._attrs: return None delay = get_delay_in_minutes(self._attrs.delay) departure = get_time_until(self._attrs.time) attrs = { "departure": f"In {departure} minutes", "departure_minutes": departure, "extra_train": self._attrs.is_extra, "vehicle_id": self._attrs.vehicle, "monitored_station": self._station.standard_name, } if delay > 0: attrs["delay"] = f"{delay} minutes" attrs["delay_minutes"] = delay return attrs async def async_update(self, **kwargs: Any) -> None: """Set the state equal to the next departure.""" liveboard = await self._api_client.get_liveboard(self._station.id) if liveboard is None: _LOGGER.warning("API failed in NMBSLiveBoard") return if not (departures := liveboard.departures): _LOGGER.warning("API returned invalid departures: %r", liveboard) return _LOGGER.debug("API returned departures: %r", departures) if len(departures) == 0: # No trains are scheduled return next_departure = departures[0] self._attrs = next_departure self._state = f"Track {next_departure.platform} - {next_departure.station}" class NMBSSensor(SensorEntity): """Get the total travel time for a given connection.""" _attr_attribution = "https://api.irail.be/" _attr_native_unit_of_measurement = UnitOfTime.MINUTES def __init__( self, api_client: iRail, name: str, show_on_map: bool, station_from: StationDetails, station_to: StationDetails, excl_vias: bool, ) -> None: """Initialize the NMBS connection sensor.""" self._name = name self._show_on_map = show_on_map self._api_client = api_client self._station_from = station_from self._station_to = station_to self._excl_vias = excl_vias self._attrs: ConnectionDetails | None = None self._state = None @property def unique_id(self) -> str: """Return the unique ID.""" unique_id = f"{self._station_from.id}_{self._station_to.id}" vias = "_excl_vias" if self._excl_vias else "" return f"nmbs_connection_{unique_id}{vias}" @property def name(self) -> str: """Return the name of the sensor.""" if self._name is None: return ( f"Train from {self._station_from.standard_name}" f" to {self._station_to.standard_name}" ) return self._name @property def icon(self) -> str: """Return the sensor default icon or an alert icon if any delay.""" if self._attrs: delay = get_delay_in_minutes(self._attrs.departure.delay) if delay > 0: return "mdi:alert-octagon" return "mdi:train" @property def extra_state_attributes(self) -> dict[str, Any] | None: """Return sensor attributes if data is available.""" if self._state is None or not self._attrs: return None delay = get_delay_in_minutes(self._attrs.departure.delay) departure = get_time_until(self._attrs.departure.time) attrs = { "destination": self._attrs.departure.station, "direction": self._attrs.departure.direction.name, "platform_arriving": self._attrs.arrival.platform, "platform_departing": self._attrs.departure.platform, "vehicle_id": self._attrs.departure.vehicle, } attrs["canceled"] = self._attrs.departure.canceled if attrs["canceled"]: attrs["departure"] = None attrs["departure_minutes"] = None else: attrs["departure"] = f"In {departure} minutes" attrs["departure_minutes"] = departure if self._show_on_map and self.station_coordinates: attrs[ATTR_LATITUDE] = self.station_coordinates[0] attrs[ATTR_LONGITUDE] = self.station_coordinates[1] if self.is_via_connection and not self._excl_vias: via = self._attrs.vias[0] attrs["via"] = via.station attrs["via_arrival_platform"] = via.arrival.platform attrs["via_transfer_platform"] = via.departure.platform attrs["via_transfer_time"] = get_delay_in_minutes( via.timebetween ) + get_delay_in_minutes(via.departure.delay) attrs["delay"] = f"{delay} minutes" attrs["delay_minutes"] = delay return attrs @property def native_value(self) -> int | None: """Return the state of the device.""" return self._state @property def station_coordinates(self) -> list[float]: """Get the lat, long coordinates for station.""" if self._state is None or not self._attrs: return [] latitude = float(self._attrs.departure.station_info.latitude) longitude = float(self._attrs.departure.station_info.longitude) return [latitude, longitude] @property def is_via_connection(self) -> bool: """Return whether the connection goes through another station.""" if not self._attrs: return False return self._attrs.vias is not None and len(self._attrs.vias) > 0 async def async_update(self, **kwargs: Any) -> None: """Set the state to the duration of a connection.""" connections = await self._api_client.get_connections( self._station_from.id, self._station_to.id ) if connections is None: _LOGGER.warning("API failed in NMBSSensor") return if not (connection := connections.connections): _LOGGER.warning("API returned invalid connection: %r", connections) return _LOGGER.debug("API returned connection: %r", connection) if connection[0].departure.left: next_connection = connection[1] else: next_connection = connection[0] self._attrs = next_connection if self._excl_vias and self.is_via_connection: _LOGGER.debug( "Skipping update of NMBSSensor because this connection is a via" ) return duration = get_ride_duration( next_connection.departure.time, next_connection.arrival.time, next_connection.departure.delay, ) self._state = duration