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_ANDROIDTV_ENTITY,
    CONF_JELLYFIN_ENTITY,
    DOMAIN,
)

_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

    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.

    Primary source: the androidtv entity's own state. When that entity is
    unavailable, falls back to the ShieldConnectionManager's StateStore
    (most recent visible ADB recent-task app).
    """

    _attr_name = "Current App"
    _attr_icon = "mdi:application"

    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
        super().__init__(hass, entry)
        self._androidtv_entity = entry.data.get(CONF_ANDROIDTV_ENTITY, "")
        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:
        self._unsub_listeners.append(
            async_track_state_change_event(
                self.hass, [self._androidtv_entity], self._handle_state_change
            )
        )
        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_state_change(self, event: Event) -> None:
        self.async_write_ha_state()

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

    @property
    def native_value(self) -> str | None:
        state = self.hass.states.get(self._androidtv_entity)
        if state:
            return state.attributes.get("app_name") or state.attributes.get("app_id")

        # androidtv entity unavailable — StateStore fallback (ADB recents)
        manager = self._manager()
        if manager is not None:
            running = manager.running_apps
            if running:
                return running[0].get("package")
        return None

    @property
    def extra_state_attributes(self) -> dict:
        state = self.hass.states.get(self._androidtv_entity)
        if state:
            return {
                "app_id": state.attributes.get("app_id"),
                "app_name": state.attributes.get("app_name"),
            }
        manager = self._manager()
        if manager is not None:
            running = manager.running_apps
            if running:
                return {"app_id": running[0].get("package")}
        return {}


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

    _attr_name = "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