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

"""Media player entity for Nvidia Shield.

State architecture:
  - Power/availability: androidtv_remote (local push, fast)
  - Playback state: Jellyfin (polling, slow 10-30s)
  - Metadata: Jellyfin attributes
  - ALL commands: androidtv_remote (Jellyfin is read-only)

Intent tracking bridges the gap between instant commands and slow Jellyfin polling:
  1. User presses pause → intent recorded → UI shows paused immediately
  2. Command sent to androidtv_remote → Shield pauses on screen
  3. 10-30s later Jellyfin confirms paused → intent cleared
  4. If Jellyfin doesn't confirm in 30s → retry once → if still no → surrender to truth
"""

import logging
import time

from .ha_adb_common import PlaybackIntent, running_apps_from_tasks

from homeassistant.components.media_player import (
    MediaPlayerDeviceClass,
    MediaPlayerEntity,
    MediaPlayerEntityFeature,
    MediaPlayerState,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import STATE_PLAYING, STATE_PAUSED, STATE_IDLE, STATE_OFF
from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback, Event
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.event import async_call_later, async_track_state_change_event

from .const import (
    ADB_POLL_INTERVAL,
    ADB_RETRY_INTERVAL,
    ATTR_MEDIA_DAG_JELLYFIN_ENTITY,
    ATTR_MEDIA_DAG_NOW_PLAYING,
    ATTR_MEDIA_DAG_RECEIVER_ENTITY,
    ATTR_MEDIA_DAG_RECEIVER_INPUT,
    ATTR_MEDIA_DAG_ROLE,
    ATTR_MEDIA_DAG_RUNNING_APPS,
    CONF_ADB_ENABLED,
    CONF_ADB_HOST,
    CONF_ADB_PORT,
    CONF_ADB_STORE_APP_INFO,
    CONF_ANDROIDTV_ENTITY,
    CONF_JELLYFIN_ENTITY,
    CONF_POWER_ENTITY,
    CONF_RECEIVER_ENTITY,
    CONF_RECEIVER_INPUT,
    DEFAULT_ADB_PORT,
    DOMAIN,
)
from .adb_session import AdbMediaSessionPoller, ensure_adb_key
from .jellyfin import fetch_session_media_info
from .power import discover_power_siblings

_LOGGER = logging.getLogger(__name__)

INTENT_TIMEOUT = 10  # seconds before self-heal fires
MAX_RETRIES = 1

# Map service names to the MediaPlayerState we expect after the command
SERVICE_TO_STATE = {
    "media_pause": MediaPlayerState.PAUSED,
    "media_play": MediaPlayerState.PLAYING,
    "media_stop": MediaPlayerState.IDLE,
}

_BASE_FEATURES = (
    MediaPlayerEntityFeature.TURN_ON
    | MediaPlayerEntityFeature.TURN_OFF
    | MediaPlayerEntityFeature.VOLUME_STEP
    | MediaPlayerEntityFeature.VOLUME_MUTE
)

_PLAYBACK_FEATURES = (
    MediaPlayerEntityFeature.PLAY
    | MediaPlayerEntityFeature.PAUSE
    | MediaPlayerEntityFeature.STOP
    | MediaPlayerEntityFeature.NEXT_TRACK
    | MediaPlayerEntityFeature.PREVIOUS_TRACK
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    async_add_entities([NvidiaShieldMediaPlayer(hass, entry)])


class NvidiaShieldMediaPlayer(MediaPlayerEntity):

    _attr_device_class = MediaPlayerDeviceClass.TV
    _attr_has_entity_name = True
    _attr_name = None
    _attr_should_poll = False

    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
        self.hass = hass
        self._entry = entry
        self._androidtv_entity = entry.data[CONF_ANDROIDTV_ENTITY]
        self._attr_unique_id = f"{DOMAIN}_{entry.entry_id}_media_player"
        self._unsub_listeners: list = []

        # Intent tracking
        self._intent: PlaybackIntent | None = None
        self._intent_timer_unsub: CALLBACK_TYPE | None = None

        # Power monitoring — discovered from power switch entity
        self._power_sensors: dict[str, str] = {}

        # Jellyfin media stream info (resolution, codec, etc.)
        self._media_info: dict | None = None

        # ADB media session, wake state, and running apps
        self._adb_poller: AdbMediaSessionPoller | None = None
        self._adb_poll_unsub: CALLBACK_TYPE | None = None
        self._adb_awake: bool | None = None
        self._running_apps: list[dict] = []

    @property
    def _jellyfin_entity(self) -> str | None:
        return self._entry.options.get(CONF_JELLYFIN_ENTITY) or None

    @property
    def _receiver_entity(self) -> str | None:
        return self._entry.options.get(CONF_RECEIVER_ENTITY) or None

    @property
    def _receiver_input(self) -> str | None:
        return self._entry.options.get(CONF_RECEIVER_INPUT) or None

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

    @property
    def _power_entity(self) -> str | None:
        return self._entry.options.get(CONF_POWER_ENTITY) or None

    @property
    def _adb_enabled(self) -> bool:
        return self._entry.options.get(CONF_ADB_ENABLED, False)

    @property
    def _adb_store_info(self) -> bool:
        return self._entry.options.get(CONF_ADB_STORE_APP_INFO, True)

    async def async_added_to_hass(self) -> None:
        # Discover power sensor siblings from the power switch
        if self._power_entity:
            self._power_sensors = discover_power_siblings(self.hass, self._power_entity)

        # Set up ADB media session poller
        if self._adb_enabled:
            await self._setup_adb()

        entities = [self._androidtv_entity]
        if self._jellyfin_entity:
            entities.append(self._jellyfin_entity)
        if self._power_entity:
            entities.append(self._power_entity)
        self._unsub_listeners.append(
            async_track_state_change_event(
                self.hass, entities, self._on_source_change
            )
        )

    async def async_will_remove_from_hass(self) -> None:
        for unsub in self._unsub_listeners:
            unsub()
        self._unsub_listeners.clear()
        self._clear_intent()
        if self._adb_poll_unsub:
            self._adb_poll_unsub()
            self._adb_poll_unsub = None
        if self._adb_poller:
            await self._adb_poller.disconnect()

    # ── Intent management ──

    def _set_intent(self, service: str, target: MediaPlayerState) -> None:
        """Record a new playback intent. Replaces any existing intent."""
        # Cancel existing timer
        if self._intent_timer_unsub:
            self._intent_timer_unsub()

        # Capture current Jellyfin state before the intent
        jf = self._read(self._jellyfin_entity)
        prior = jf.state if jf else None

        self._intent = PlaybackIntent.create(
            target_state=target.value,
            command=service,
            prior_state=prior,
        )

        # Schedule self-heal timeout
        self._intent_timer_unsub = async_call_later(
            self.hass, INTENT_TIMEOUT, self._on_intent_timeout
        )

        _LOGGER.debug("Intent set: %s → %s (prior JF: %s)", service, target, prior)

    def _clear_intent(self) -> None:
        """Clear the pending intent and its timer."""
        if self._intent_timer_unsub:
            self._intent_timer_unsub()
            self._intent_timer_unsub = None
        self._intent = None

    @callback
    def _on_intent_timeout(self, _now) -> None:
        """Self-heal: intent wasn't confirmed by Jellyfin in time."""
        intent = self._intent
        if intent is None:
            return

        jf = self._read(self._jellyfin_entity)
        jf_state = jf.state if jf else None

        # Check if Jellyfin caught up during the timeout window
        if jf_state == intent.target_state:
            _LOGGER.debug("Intent %s confirmed at timeout check", intent.command)
            self._clear_intent()
            self.async_write_ha_state()
            return

        # Retry once
        if intent.retry_count < MAX_RETRIES:
            intent.retry_count += 1
            intent.created_at = time.monotonic()
            _LOGGER.warning(
                "Intent %s not confirmed (JF=%s, want=%s), retrying (attempt %d)",
                intent.command, jf_state, intent.target_state, intent.retry_count,
            )
            self.hass.async_create_task(
                self._cmd(intent.command)
            )
            self._intent_timer_unsub = async_call_later(
                self.hass, INTENT_TIMEOUT, self._on_intent_timeout
            )
            return

        # Exhausted retries — surrender to Jellyfin truth
        _LOGGER.error(
            "Intent %s failed after %d retries (JF=%s), falling back to truth",
            intent.command, MAX_RETRIES, jf_state,
        )
        self._clear_intent()
        self.async_write_ha_state()

    # ── State change listener ──

    @callback
    def _on_source_change(self, event: Event) -> None:
        """React to state changes in androidtv or Jellyfin."""
        entity_id = event.data.get("entity_id")
        new_state_obj = event.data.get("new_state")

        # Power off clears any intent — device is off, intent is moot
        if entity_id in (self._androidtv_entity, self._power_entity) and new_state_obj:
            if not self._is_power_on():
                if self._intent:
                    _LOGGER.debug("Intent cleared: device powered off")
                    self._clear_intent()

        # Jellyfin state change — check for intent confirmation and fetch media info
        if entity_id == self._jellyfin_entity and new_state_obj:
            new_jf = new_state_obj.state

            # Intent confirmation
            if self._intent:
                intent = self._intent
                if new_jf == intent.target_state:
                    _LOGGER.debug("Intent %s confirmed by Jellyfin", intent.command)
                    self._clear_intent()
                elif new_jf not in (
                    intent.prior_state, "unavailable", "unknown", None
                ):
                    _LOGGER.debug(
                        "Intent %s overridden by external action (JF=%s)",
                        intent.command, new_jf,
                    )
                    self._clear_intent()

            # Fetch media stream info when playing
            if new_jf in (STATE_PLAYING, STATE_PAUSED):
                entity_picture = new_state_obj.attributes.get("entity_picture", "")
                self.hass.async_create_task(
                    self._fetch_media_info(entity_picture)
                )
            else:
                self._media_info = None


        # Always re-evaluate state
        self.async_write_ha_state()

    # ── Helpers ──

    def _read(self, entity_id: str | None):
        if not entity_id:
            return None
        return self.hass.states.get(entity_id)

    def _discover_adb_host(self) -> str | None:
        """Auto-detect ADB host from the androidtv_remote config entry for our entity."""
        from homeassistant.helpers import entity_registry as er

        registry = er.async_get(self.hass)
        entity_entry = registry.async_get(self._androidtv_entity)
        if not entity_entry or not entity_entry.config_entry_id:
            return None

        config_entry = self.hass.config_entries.async_get_entry(entity_entry.config_entry_id)
        if config_entry:
            return config_entry.data.get("host")
        return None

    async def _setup_adb(self) -> None:
        """Initialize ADB connection and start polling."""
        try:
            # Only auto-connect if ADB was previously paired (key exists)
            import os
            key_dir = os.path.join(self.hass.config.config_dir, ".storage", DOMAIN)
            key_file = os.path.join(key_dir, "adbkey")
            if not await self.hass.async_add_executor_job(os.path.exists, key_file):
                _LOGGER.info(
                    "ADB key not found — skipping auto-connect. "
                    "Use the Pair ADB button or reconfigure the integration."
                )
                return

            # Get device IP: explicit config > auto-detect from androidtv_remote
            host = self._entry.options.get(CONF_ADB_HOST)
            port = self._entry.options.get(CONF_ADB_PORT, DEFAULT_ADB_PORT)

            if not host:
                host = self._discover_adb_host()

            if not host:
                _LOGGER.warning("Could not determine ADB host. Set it in integration options.")
                return

            key_path = key_file
            self._adb_poller = AdbMediaSessionPoller(host, port, key_path)

            if await self._adb_poller.connect():
                # Start periodic polling
                self._adb_poll_unsub = async_call_later(
                    self.hass, ADB_POLL_INTERVAL, self._adb_poll_callback
                )
                _LOGGER.info("ADB media session polling started for %s", host)
            else:
                _LOGGER.warning(
                    "ADB connection to %s failed. If this is the first connection, "
                    "approve the ADB debugging prompt on the device. "
                    "Retrying in %d seconds.",
                    host, ADB_RETRY_INTERVAL,
                )
                self._adb_poll_unsub = async_call_later(
                    self.hass, ADB_RETRY_INTERVAL, self._adb_poll_callback
                )
        except Exception:
            _LOGGER.exception("Failed to set up ADB")

    @callback
    def _adb_poll_callback(self, _now) -> None:
        """Schedule ADB poll as a task."""
        self.hass.async_create_task(self._adb_poll())

    async def _adb_poll(self) -> None:
        """Poll ADB for wake state and media session, then schedule next poll."""
        try:
            if not self._adb_poller:
                return

            # Always check wake state
            prev_awake = self._adb_awake
            self._adb_awake = await self._adb_poller.is_awake()

            # Only poll media session when awake and Jellyfin is NOT active
            jf = self._read(self._jellyfin_entity)
            jf_active = jf and jf.state in (STATE_PLAYING, STATE_PAUSED)

            if self._adb_awake:
                # Always poll ADB when awake — pass the foreground app
                # so the parser prefers its session over stale ones
                await self._adb_poller.poll(foreground_package=self.app_id)
                self._running_apps = await self._adb_poller.poll_running_apps()
            elif self._adb_poller.last_info:
                # Clear media info when asleep
                self._adb_poller.clear_info()
                self._running_apps = []

            # Always update state (wake state may have changed)
            if self._adb_awake != prev_awake or self._adb_poller.last_info:
                self.async_write_ha_state()
        except Exception:
            _LOGGER.debug("ADB poll error", exc_info=True)
        finally:
            # Use slow retry if ADB is disconnected, normal interval if connected
            if self._adb_enabled:
                interval = ADB_POLL_INTERVAL if self._adb_poller.connected else ADB_RETRY_INTERVAL
                self._adb_poll_unsub = async_call_later(
                    self.hass, interval, self._adb_poll_callback
                )

    async def _fetch_media_info(self, entity_picture: str) -> None:
        """Read stream info from the HA Jellyfin integration's coordinator."""
        info = await fetch_session_media_info(
            self.hass, self._jellyfin_entity, self._entry.title
        )
        if info != self._media_info:
            self._media_info = info
            self.async_write_ha_state()

    def _jf_active(self) -> bool:
        """Is Jellyfin reporting active playback?"""
        jf = self._read(self._jellyfin_entity)
        if jf and jf.state in (STATE_PLAYING, STATE_PAUSED):
            return True
        # If we have an active intent for play/pause, media is conceptually active
        if self._intent and self._intent.target_state in (
            STATE_PLAYING, STATE_PAUSED
        ):
            return True
        return False

    # ── State ──

    def _is_power_on(self) -> bool | None:
        """Determine if the Shield is powered on.

        Priority: ADB wake state > power entity plug off > unknown.
        ADB is the authority — it knows if the device is awake or asleep.
        The plug being "on" just means the outlet has power (standby).
        """
        # 1. ADB wake state — most accurate
        if self._adb_awake is not None:
            return self._adb_awake

        # 2. Power entity — only trust "off" (plug pulled = definitely off)
        if self._power_entity:
            power_state = self._read(self._power_entity)
            if power_state and power_state.state == "off":
                return False

        return None

    @property
    def state(self) -> MediaPlayerState:
        # Check if plug is off first — truly off
        if self._power_entity:
            power_state = self._read(self._power_entity)
            if power_state and power_state.state == "off":
                return MediaPlayerState.OFF

        power = self._is_power_on()

        if power is False:
            return MediaPlayerState.IDLE  # ADB says asleep — standby (yellow dot)

        if power is None:
            return MediaPlayerState.IDLE

        # Tier 1: Active intent overrides Jellyfin
        if self._intent is not None:
            return MediaPlayerState(self._intent.target_state)

        # Tier 2: Jellyfin is authoritative
        jf = self._read(self._jellyfin_entity)
        if jf:
            if jf.state == STATE_PLAYING:
                return MediaPlayerState.PLAYING
            if jf.state == STATE_PAUSED:
                return MediaPlayerState.PAUSED

        # Tier 3: ADB media session (non-Jellyfin apps)
        if self._adb_poller and self._adb_poller.last_info:
            adb_info = self._adb_poller.last_info
            if adb_info.playing:
                return MediaPlayerState.PLAYING
            if adb_info.paused:
                return MediaPlayerState.PAUSED

        # Tier 4: On but no media
        return MediaPlayerState.ON

    @property
    def supported_features(self) -> MediaPlayerEntityFeature:
        if self._jf_active():
            return _BASE_FEATURES | _PLAYBACK_FEATURES
        # ADB media session also supports playback controls
        if self._adb_poller and self._adb_poller.last_info:
            adb_info = self._adb_poller.last_info
            if adb_info.playing or adb_info.paused:
                return _BASE_FEATURES | _PLAYBACK_FEATURES
        return _BASE_FEATURES

    # ── Metadata from Jellyfin ──

    def _jf_attr(self, key: str):
        jf = self._read(self._jellyfin_entity)
        return jf.attributes.get(key) if jf else None

    @property
    def app_id(self):
        atv = self._read(self._androidtv_entity)
        return atv.attributes.get("app_id") if atv else None

    @property
    def app_name(self):
        atv = self._read(self._androidtv_entity)
        return atv.attributes.get("app_name") if atv else None

    @property
    def media_content_type(self):
        return self._jf_attr("media_content_type") if self._jf_active() else None

    @property
    def media_title(self):
        return self._jf_attr("media_title") if self._jf_active() else None

    @property
    def media_series_title(self):
        return self._jf_attr("media_series_title") if self._jf_active() else None

    @property
    def media_season(self):
        return self._jf_attr("media_season") if self._jf_active() else None

    @property
    def media_episode(self):
        return self._jf_attr("media_episode") if self._jf_active() else None

    @property
    def media_duration(self):
        return self._jf_attr("media_duration") if self._jf_active() else None

    @property
    def media_position(self):
        return self._jf_attr("media_position") if self._jf_active() else None

    @property
    def media_position_updated_at(self):
        return self._jf_attr("media_position_updated_at") if self._jf_active() else None

    @property
    def media_image_url(self):
        return self._jf_attr("entity_picture") if self._jf_active() else None

    def _get_now_playing(self) -> dict | None:
        if not self._jf_active():
            return None
        jf = self._read(self._jellyfin_entity)
        if not jf:
            return None
        attrs = jf.attributes
        np = {"state": self.state.value, "title": attrs.get("media_title"),
              "media_type": attrs.get("media_content_type")}
        for src, dst in [("media_series_title", "series"), ("media_season", "season"),
                         ("media_episode", "episode"), ("media_duration", "duration"),
                         ("media_position", "position")]:
            v = attrs.get(src)
            if v is not None:
                np[dst] = v
        if attrs.get("entity_picture"):
            np["thumbnail"] = attrs["entity_picture"]
        if self._media_info:
            np["media_info"] = self._media_info
        return np

    def _get_adb_now_playing(self) -> dict | None:
        """Build now-playing from ADB media session (non-Jellyfin apps)."""
        if not self._adb_poller or not self._adb_poller.last_info:
            return None
        if not self._adb_store_info:
            # Privacy mode: show that something is playing but not what
            info = self._adb_poller.last_info
            if info.playing or info.paused:
                return {
                    "state": "playing" if info.playing else "paused",
                    "title": self.app_name or self.app_id or "Playing",
                    "source": "media_session",
                }
            return None

        info = self._adb_poller.last_info
        if not info.title:
            return None
        np = {
            "state": "playing" if info.playing else ("paused" if info.paused else "idle"),
            "title": info.title,
            "source": "media_session",
        }
        if info.artist:
            np["artist"] = info.artist
        if info.album:
            np["album"] = info.album
        return np

    def _get_power_data(self) -> dict | None:
        """Read current values from discovered power sensors."""
        if not self._power_sensors:
            return None
        data = {}
        for key, entity_id in self._power_sensors.items():
            if key == "switch":
                continue
            state = self._read(entity_id)
            if state and state.state not in ("unavailable", "unknown"):
                try:
                    data[key] = float(state.state)
                except (ValueError, TypeError):
                    data[key] = state.state
        return data if data else None

    def _get_running_apps(self) -> list[dict] | None:
        """Build a list of running apps from ADB recent tasks + media sessions."""
        if not self._adb_poller or not self._running_apps:
            return None
        if not self._adb_store_info:
            return None
        return running_apps_from_tasks(
            self._running_apps,
            self._adb_poller.all_sessions,
            foreground_package=self.app_id,
        )

    @property
    def extra_state_attributes(self) -> dict:
        # Now playing: prefer Jellyfin, fall back to ADB media session
        now_playing = self._get_now_playing()
        if not now_playing:
            now_playing = self._get_adb_now_playing()

        attrs = {"app_id": self.app_id, "app_name": self.app_name,
                 ATTR_MEDIA_DAG_ROLE: "source",
                 ATTR_MEDIA_DAG_NOW_PLAYING: now_playing}
        if self._receiver_entity:
            attrs[ATTR_MEDIA_DAG_RECEIVER_ENTITY] = self._receiver_entity
        if self._receiver_input:
            attrs[ATTR_MEDIA_DAG_RECEIVER_INPUT] = self._receiver_input
        if self._jellyfin_entity:
            attrs[ATTR_MEDIA_DAG_JELLYFIN_ENTITY] = self._jellyfin_entity

        # Running apps from ADB media sessions
        running_apps = self._get_running_apps()
        if running_apps:
            attrs[ATTR_MEDIA_DAG_RUNNING_APPS] = running_apps

        # Power monitoring
        if self._power_entity:
            attrs["power_switch"] = self._power_entity
        power_data = self._get_power_data()
        if power_data:
            attrs["power_monitoring"] = power_data

        # Intent state for debugging
        if self._intent:
            attrs["_intent"] = self._intent.to_debug_dict()
        return attrs

    # ── Commands — ALL go to androidtv_remote ──

    async def _cmd(self, service: str, **extra) -> None:
        data = {"entity_id": self._androidtv_entity, **extra}
        await self.hass.services.async_call("media_player", service, data)

    async def async_turn_on(self):
        await self._cmd("turn_on")

    async def async_turn_off(self):
        await self._cmd("turn_off")

    async def async_volume_up(self):
        await self._cmd("volume_up")

    async def async_volume_down(self):
        await self._cmd("volume_down")

    async def async_mute_volume(self, mute):
        await self._cmd("volume_mute", is_volume_muted=mute)

    async def _media_cmd(self, ha_service: str, adb_command: str, intent_state: MediaPlayerState | None = None) -> None:
        """Send a media command. Uses ADB dispatch for non-Jellyfin apps when available."""
        if intent_state:
            self._set_intent(ha_service, intent_state)
            self.async_write_ha_state()

        # Use ADB media dispatch when ADB is active and Jellyfin isn't playing
        jf = self._read(self._jellyfin_entity)
        jf_active = jf and jf.state in (STATE_PLAYING, STATE_PAUSED)

        if self._adb_poller and self._adb_poller.connected and not jf_active:
            success = await self._adb_poller.send_media_command(adb_command)
            if success:
                return

        # Fall back to Android TV Remote
        await self._cmd(ha_service)

    async def async_media_play(self):
        await self._media_cmd("media_play", "play", MediaPlayerState.PLAYING)

    async def async_media_pause(self):
        await self._media_cmd("media_pause", "pause", MediaPlayerState.PAUSED)

    async def async_media_stop(self):
        await self._media_cmd("media_stop", "stop", MediaPlayerState.IDLE)

    async def async_media_next_track(self):
        await self._media_cmd("media_next_track", "next")

    async def async_media_previous_track(self):
        await self._media_cmd("media_previous_track", "previous")

    async def async_bring_to_foreground(self, package: str) -> bool:
        """Bring a specific app to the foreground via ADB."""
        if not self._adb_poller:
            _LOGGER.warning("ADB not available for bring_to_foreground")
            return False
        return await self._adb_poller.bring_to_foreground(package)