core/homeassistant/components/nmbs/sensor.py

"""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 = int(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} 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