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_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_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"})

    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 foreground package from the manager's store (t41).

        StateStore priority chain (P3-3c): protocol current_app push → ADB
        mResumedActivity → ADB recents.
        """
        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.
        """
        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 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