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_PLAYING, STATE_PAUSED
from homeassistant.core import HomeAssistant, callback, Event
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.event import async_track_state_change_event

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."""

    _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[CONF_ANDROIDTV_ENTITY]
        self._attr_unique_id = f"{DOMAIN}_{entry.entry_id}_current_app"

    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
            )
        )

    @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._androidtv_entity)
        if not state:
            return None
        return state.attributes.get("app_name") or state.attributes.get("app_id")

    @property
    def extra_state_attributes(self) -> dict:
        state = self.hass.states.get(self._androidtv_entity)
        if not state:
            return {}
        return {
            "app_id": state.attributes.get("app_id"),
            "app_name": state.attributes.get("app_name"),
        }


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