ha-nvidia-shield/custom_components/nvidia_shield/sensor.py

"""Sensor entities for Nvidia Shield."""

import logging

from homeassistant.components.sensor import SensorEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import STATE_PAUSED, STATE_PLAYING
from homeassistant.core import Event, HomeAssistant, callback
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.event import async_track_state_change_event

from .connection import ShieldConnectionManager
from .const import (
    CONF_ADB_ENABLED,
    CONF_APP_NAME_MAP,
    CONF_JELLYFIN_ENTITY,
    DOMAIN,
)
from .media_player import resolve_app_name

_LOGGER = logging.getLogger(__name__)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up the Nvidia Shield sensor entities."""
    entities = [NvidiaShieldCurrentAppSensor(hass, entry)]
    if entry.options.get(CONF_ADB_ENABLED, True):
        # discovery is an ADB feature; with the transport off there is nothing
        # to enumerate and the entity would only ever report "unavailable"
        entities.append(NvidiaShieldAppsSensor(hass, entry))

    if entry.options.get(CONF_JELLYFIN_ENTITY):
        entities.append(NvidiaShieldNowPlayingSensor(hass, entry))

    async_add_entities(entities)


class NvidiaShieldSensorBase(SensorEntity):
    """Base class for Shield sensors."""

    _attr_has_entity_name = True
    _attr_should_poll = False
    # G12 (t30): poll-driven attributes never belong in the database
    _unrecorded_attributes = frozenset(
        {"app_id", "app_name", "playback_state", "installed_apps"}
    )

    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
        self.hass = hass
        self._entry = entry
        self._unsub_listeners: list = []

    @property
    def device_info(self) -> DeviceInfo:
        return DeviceInfo(
            identifiers={(DOMAIN, self._entry.entry_id)},
            name=self._entry.title,
            manufacturer="NVIDIA",
            model="Shield TV",
        )

    async def async_will_remove_from_hass(self) -> None:
        for unsub in self._unsub_listeners:
            unsub()
        self._unsub_listeners.clear()


class NvidiaShieldCurrentAppSensor(NvidiaShieldSensorBase):
    """Current running app on the Shield.

    Source: the ShieldConnectionManager's StateStore — the protocol current-app
    push, or the most recent visible ADB recent-task app when that has not
    reported (t41: our own transports only).
    """

    _attr_translation_key = "current_app"
    _attr_icon = "mdi:application"

    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
        super().__init__(hass, entry)
        self._attr_unique_id = f"{DOMAIN}_{entry.entry_id}_current_app"

    def _manager(self) -> ShieldConnectionManager | None:
        manager = getattr(self._entry, "runtime_data", None)
        return manager if isinstance(manager, ShieldConnectionManager) else None

    async def async_added_to_hass(self) -> None:
        manager = self._manager()
        if manager is not None:
            # Refresh when the connection manager updates its store.
            self._unsub_listeners.append(manager.register_listener(self._handle_store_update))

    @callback
    def _handle_store_update(self) -> None:
        self.async_write_ha_state()

    def _package(self) -> str | None:
        """The raw package of the app that is on screen (t41, t55).

        StateStore priority chain (P3-3c): protocol current_app push → ADB
        mResumedActivity (both via ``manager.foreground_app``) → the most
        recent visible ADB recent task. The recents tier is a last resort for
        a moment when neither live signal has reported yet; it keeps the
        sensor from going None while ADB is still warming up.
        """
        manager = self._manager()
        if manager is None:
            return None
        foreground = manager.foreground_app
        if foreground:
            return foreground
        running = manager.running_apps
        if running:
            return running[0].get("package")
        return None

    @property
    def native_value(self) -> str | None:
        """The resolved friendly name (t40), not the raw package id.

        Exactly the ``media_player.app_name`` chain via the shared
        ``resolve_app_name``: apps map → cached manifest label → the package
        text itself. A bare or unknown package shows its package text rather
        than None, like the media player's app_name.

        This is why ``_package`` must be the CURRENT app and the label fetch
        must be driven off it (t55): the device-manifest label is resolved on
        demand for exactly the package passed here, so a live run showing a
        package id means the label tier never answered — see
        ``ShieldConnectionManager.app_label`` for the caching rule that makes
        a miss retry instead of being stamped as an eternal "no label".
        """
        manager = self._manager()
        return resolve_app_name(
            self._package(),
            dict(self._entry.options.get(CONF_APP_NAME_MAP) or {}),
            manager.app_label if manager is not None else None,
        )

    @property
    def extra_state_attributes(self) -> dict:
        """The raw package stays visible here (t40) next to the shown name."""
        package = self._package()
        attrs: dict = {}
        if package:
            attrs["app_id"] = package
        name = self.native_value
        if name:
            attrs["app_name"] = name
        return attrs


class NvidiaShieldAppsSensor(NvidiaShieldSensorBase):
    """The apps installed on the Shield (t57) — a dashboard-facing surface.

    The state is the app COUNT and the ``installed_apps`` attribute is the
    table: ``[{package, name, icon}, ...]``, exactly what the Configure page's
    Apps sub-screen edits. It is the answer to "it doesn't find all the apps":
    the list comes from the device (``pm list packages -3`` plus the active
    launcher), labelled through the cached ``pm dump`` path.

    Only ADB-enabled entries get one. Discovery is a handful of shells, so it
    runs once when the entity is added and then only when the table has gone
    stale: ``async_update`` goes through the manager's TTL cache, so a poll on a
    fresh table issues NO shell. On top of that the user can ask for a refresh
    with the standard ``homeassistant.update_entity`` action, and the Configure
    page runs the same sweep from its discovery row. With ADB off, or
    unreachable, the attributes say why and the state is 0 rather than raising.
    """

    _attr_translation_key = "installed_apps"
    _attr_icon = "mdi:apps"
    # Polling is what makes the standard `homeassistant.update_entity` action
    # work for it — the user-facing "refresh the app list" trigger. The poll is
    # free while the table is fresh: the manager returns the cache and issues no
    # shell until APP_LIST_TTL expires.
    _attr_should_poll = True

    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
        super().__init__(hass, entry)
        self._attr_unique_id = f"{DOMAIN}_{entry.entry_id}_installed_apps"

    def _manager(self) -> ShieldConnectionManager | None:
        manager = getattr(self._entry, "runtime_data", None)
        return manager if isinstance(manager, ShieldConnectionManager) else None

    async def async_added_to_hass(self) -> None:
        manager = self._manager()
        if manager is not None:
            self._unsub_listeners.append(manager.register_listener(self._handle_store_update))
        # one discovery per entity lifetime; a load is already a wake-up moment,
        # and nothing here runs on a state poll
        self.hass.async_create_task(self.async_update())

    @callback
    def _handle_store_update(self) -> None:
        if self.hass is not None:
            self.async_write_ha_state()

    async def async_update(self) -> None:
        """Refresh the table on demand (explicitly requested only)."""
        manager = self._manager()
        if manager is None:
            return
        await manager.async_refresh_installed_apps()

    @property
    def native_value(self) -> int:
        manager = self._manager()
        return len(manager.installed_apps()) if manager is not None else 0

    @property
    def extra_state_attributes(self) -> dict:
        """The table itself, plus why it may be empty (never a crash)."""
        manager = self._manager()
        if manager is None:
            return {"installed_apps": [], "discovery": "unavailable"}
        apps = manager.installed_apps()
        attrs = {"installed_apps": apps}
        if not apps:
            if not self._entry.options.get(CONF_ADB_ENABLED, True):
                attrs["discovery"] = "adb_disabled"
            elif not manager.connected:
                attrs["discovery"] = "not_connected"
            else:
                attrs["discovery"] = "pending"
        else:
            attrs["discovery"] = "ok"
        return attrs


class NvidiaShieldNowPlayingSensor(NvidiaShieldSensorBase):
    """What's currently playing via Jellyfin."""

    _attr_translation_key = "now_playing"
    _attr_icon = "mdi:play-circle"

    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
        super().__init__(hass, entry)
        self._jellyfin_entity = entry.options[CONF_JELLYFIN_ENTITY]
        self._attr_unique_id = f"{DOMAIN}_{entry.entry_id}_now_playing"

    async def async_added_to_hass(self) -> None:
        self._unsub_listeners.append(
            async_track_state_change_event(
                self.hass, [self._jellyfin_entity], self._handle_state_change
            )
        )

    @callback
    def _handle_state_change(self, event: Event) -> None:
        self.async_write_ha_state()

    @property
    def native_value(self) -> str | None:
        state = self.hass.states.get(self._jellyfin_entity)
        if not state or state.state not in (STATE_PLAYING, STATE_PAUSED):
            return "Idle"

        attrs = state.attributes
        title = attrs.get("media_title", "")
        series = attrs.get("media_series_title")

        if series:
            season = attrs.get("media_season", "")
            episode = attrs.get("media_episode", "")
            return f"{series} S{season}E{episode} - {title}"
        return title or "Playing"

    @property
    def extra_state_attributes(self) -> dict:
        state = self.hass.states.get(self._jellyfin_entity)
        if not state or state.state not in (STATE_PLAYING, STATE_PAUSED):
            return {"playback_state": "idle"}

        attrs = state.attributes
        result = {
            "playback_state": state.state,
            "title": attrs.get("media_title"),
            "media_type": attrs.get("media_content_type"),
        }

        if attrs.get("media_series_title"):
            result["series"] = attrs["media_series_title"]
        if attrs.get("media_season") is not None:
            result["season"] = attrs["media_season"]
        if attrs.get("media_episode") is not None:
            result["episode"] = attrs["media_episode"]
        if attrs.get("media_duration") is not None:
            result["duration"] = attrs["media_duration"]
        if attrs.get("media_position") is not None:
            result["position"] = attrs["media_position"]
        if attrs.get("entity_picture"):
            result["thumbnail"] = attrs["entity_picture"]

        return result