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

"""The Nvidia Shield integration.

Phase 1 architecture (docs/MVP_PLAN.md §2): the entry owns a single
ShieldConnectionManager stored in ``entry.runtime_data`` — the persistent ADB
socket and all ADB-derived state live there, not on the entities. Entities are
thin callback-views over the manager's StateStore.

ADB is optional: a failed/absent ADB connection never blocks setup (no
ConfigEntryNotReady) — the manager retries in the background, or stays
passive until the user pairs via the Pair ADB button.
"""

import logging
import os

from homeassistant import config_entries
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform
from homeassistant.core import HomeAssistant

from .adb_session import AdbMediaSessionPoller
from .connection import ShieldConnectionManager
from .const import (
    ADB_COMMAND_SERVICE_DESCRIPTION,
    CONF_ADB_ENABLED,
    CONF_ADB_HOST,
    CONF_ADB_PORT,
    CONF_ANDROIDTV_ENTITY,
    CONF_HOST,
    CONF_PROTOCOL_ENABLED,
    CONF_PROTOCOL_HOST,
    DEFAULT_ADB_PORT,
    DEFAULT_NAME,
    DOMAIN,
    ENTRY_VERSION,
    SERVICE_ADB_COMMAND,
)
from .helpers import (
    entry_host,
    protocol_cert_paths,
    resolve_adb_config,
    seed_adb_enabled,
    seed_host,
)

_LOGGER = logging.getLogger(__name__)

# t66: IMAGE is listed unconditionally — the platform's async_setup_entry adds
# nothing unless the user opted in, so an install that never turns screen capture
# on pays nothing for the platform existing (and toggling it is an options reload).
PLATFORMS = [
    Platform.MEDIA_PLAYER,
    Platform.SENSOR,
    Platform.REMOTE,
    Platform.BUTTON,
    Platform.IMAGE,
]


async def async_setup(hass: HomeAssistant, config) -> bool:
    """Domain setup: actions (always-on) + optional YAML seeding (t30 U1).

    YAML is SEED-ONLY (user ruling): every ``nvidia_shield`` block starts an
    interactive SOURCE_IMPORT config flow — nothing is created without the
    user walking the flow, and pairing stays interactive exactly as for UI
    setup. Removing the YAML block creates OR deletes nothing. UI-configured
    entries are never re-seeded (host/MAC duplicate guards).
    """
    domain_config = (config or {}).get(DOMAIN)
    for block in _yaml_seed_blocks(domain_config or []):
        if not seed_host(block):
            _LOGGER.warning(
                "nvidia_shield YAML block without an address ignored: %s",
                block.get("name", "<unnamed>"),
            )
            continue
        hass.async_create_task(
            hass.config_entries.flow.async_init(
                DOMAIN,
                context={"source": config_entries.SOURCE_IMPORT},
                data=block,
            )
        )
    # t29 G8 (kept here across the t30 splice): the action registers ALWAYS.
    # The two APIs the target-bearing registration needs are imported WITH that
    # registration, next to the handler they belong to (t14): importing the package —
    # which every config-flow harness does via
    # ``custom_components.nvidia_shield.config_flow`` — must not drag in the service
    # registration surface for a domain that harness never sets up.
    from homeassistant.core import SupportsResponse
    from homeassistant.helpers.service import async_set_service_schema

    from .media_player import _ADB_COMMAND_SCHEMA, _async_adb_command_service

    # t14: the action must be SELECTABLE for the user's Shield. Measured on HA
    # 2026.10.1: `async_register` has no `target` parameter, so the UI's target
    # picker comes from the service DESCRIPTION — the runtime
    # `async_set_service_schema` (which wins over a services.yaml we do not ship)
    # carries `target: {entity: [{domain: media_player}]}`, so the picker offers the
    # Shield's media player entity. `supports_response=OPTIONAL` is registered FIRST
    # because async_set_service_schema reads `hass.services.supports_response()` to
    # publish the response flag, and the handler returns the command's output under
    # ADB_COMMAND_RESPONSE_KEY.
    hass.services.async_register(
        DOMAIN,
        SERVICE_ADB_COMMAND,
        _async_adb_command_service,
        schema=_ADB_COMMAND_SCHEMA,
        supports_response=SupportsResponse.OPTIONAL,
    )
    async_set_service_schema(
        hass, DOMAIN, SERVICE_ADB_COMMAND, ADB_COMMAND_SERVICE_DESCRIPTION
    )
    return True


def _yaml_seed_blocks(domain_config) -> list[dict]:
    """Normalize the optional YAML domain key into seed blocks (t30 U1)."""
    blocks: list[dict] = []
    raw_list = domain_config if isinstance(domain_config, list) else [domain_config]
    for block in raw_list:
        if not isinstance(block, dict):
            continue
        name = block.get("name") or DEFAULT_NAME
        # t32 F4 (review rider): the seed honors the block's shape — a
        # protocol-only block does NOT force ADB on. t41: the address is ONE
        # field (`host`, or a legacy protocol_host/adb_host).
        seed = {
            "name": name,
            CONF_ADB_ENABLED: seed_adb_enabled(block),
            CONF_HOST: seed_host(block),
            CONF_ADB_PORT: int(block.get(CONF_ADB_PORT) or DEFAULT_ADB_PORT),
        }
        if block.get("unique_id"):
            seed["unique_id"] = str(block["unique_id"])
        blocks.append(seed)
    return blocks


async def _build_poller_factory(hass: HomeAssistant, entry: ConfigEntry):
    """Poller factory when ADB is configured AND usable today, else None.

    Returning None leaves the manager passive: the entry still sets up, the
    entities simply see no ADB state. The Pair ADB button / reconfigure
    triggers a reload that re-evaluates this.
    """
    if not entry.options.get(CONF_ADB_ENABLED, False):
        return None

    adb = resolve_adb_config(hass, entry)
    if adb is None:
        _LOGGER.warning(
            "ADB enabled but no host resolvable for entry %s — continuing "
            "without ADB state.",
            entry.title,
        )
        return None

    # Only auto-connect when a key already exists — an ADB handshake against
    # an unpaired device is pointless noise (pairing happens via the Pair ADB
    # button / config flow, untouched).
    key_path_ready = await hass.async_add_executor_job(
        os.path.exists, adb.key_path
    )
    if not key_path_ready:
        _LOGGER.info(
            "ADB key not found for entry %s — skipping auto-connect. "
            "Use the Pair ADB button or reconfigure the integration.",
            entry.title,
        )
        return None

    # t32 F1 (review repair): the audit-sanctioned FACTORY RE-READ — the
    # host/port resolve from entry.options AT CALL TIME, so the poller built
    # after a zeroconf IP change targets the NEW address without a reload
    # (the setup-time closure pinned the old address). A key-path change
    # still needs the Pair ADB button / reconfigure reload, by design.
    def _factory():
        adb_now = resolve_adb_config(hass, entry)
        if adb_now is None:
            return None
        return AdbMediaSessionPoller(adb_now.host, adb_now.port, adb_now.key_path)

    return _factory


async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Set up Nvidia Shield from a config entry."""
    hass.data.setdefault(DOMAIN, {})
    hass.data[DOMAIN][entry.entry_id] = entry

    manager = ShieldConnectionManager(
        hass,
        entry,
        poller_factory=await _build_poller_factory(hass, entry),
        protocol_settings=await _build_protocol_settings(hass, entry),
        invalid_auth_callback=lambda: entry.async_start_reauth(hass),
    )
    entry.runtime_data = manager
    await manager.async_start()

    # Close the persistent socket on Home Assistant stop as well as unload.
    async def _close_on_hass_stop(_event) -> None:
        await manager.async_stop()

    entry.async_on_unload(
        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _close_on_hass_stop)
    )

    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)

    return True


async def _build_protocol_settings(hass: HomeAssistant, entry: ConfigEntry):
    """Manager protocol_settings when the transport is enabled & resolvable.

    The address is the entry's own single field (t41). Nothing is derived from
    another integration.
    """
    if not entry.options.get(CONF_PROTOCOL_ENABLED, False):
        return None

    host = entry_host(entry)
    if not host:
        _LOGGER.info(
            "Protocol enabled for entry %s but no host resolvable — continuing "
            "without the remote-protocol transport. Pair via the flow.",
            entry.title,
        )
        return None

    certfile, keyfile = await hass.async_add_executor_job(protocol_cert_paths, hass)
    return {
        "host": host,
        "client_name": entry.title or "Home Assistant",
        "certfile": certfile,
        "keyfile": keyfile,
    }


async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
    """Migrate an older config entry to the current version.

    v1 → v2 added the remote-protocol fields (ADB options carried over
    untouched). v2 → v3 (t41) mirrors the two changes in the stored shape:

    * the duplicated address fields collapse into the single CONF_HOST — the
      protocol host wins, else the ADB host, since both name the same Shield;
    * the androidtv-entity key goes. This integration owns the whole stack now:
      the entity it used to configure is neither a state source nor an app-name
      or power tier, so keeping the value would only mislead.

    Downgrades (version > ENTRY_VERSION) are rejected so HA offers a rollback.
    """
    if config_entry.version > ENTRY_VERSION:
        return False
    if config_entry.version >= ENTRY_VERSION:
        return True

    options = {**config_entry.options}
    data = {**config_entry.data}

    if config_entry.version < 2:
        options.setdefault(CONF_PROTOCOL_ENABLED, False)

    if not options.get(CONF_HOST):
        options[CONF_HOST] = (
            options.get(CONF_PROTOCOL_HOST)
            or options.get(CONF_ADB_HOST)
            or data.get(CONF_PROTOCOL_HOST)
            or data.get(CONF_ADB_HOST)
            or ""
        )

    dropped_entity = options.pop(CONF_ANDROIDTV_ENTITY, None)
    dropped_entity = data.pop(CONF_ANDROIDTV_ENTITY, None) or dropped_entity
    for legacy_key in (CONF_PROTOCOL_HOST, CONF_ADB_HOST):
        options.pop(legacy_key, None)
        data.pop(legacy_key, None)

    hass.config_entries.async_update_entry(
        config_entry, data=data, options=options, version=ENTRY_VERSION
    )
    if dropped_entity:
        _LOGGER.info(
            "%s: config entry migrated to version %d — dropped the androidtv "
            "entity %s (this integration reads the Shield directly; the entry "
            "needs no external media_player)",
            config_entry.title,
            ENTRY_VERSION,
            dropped_entity,
        )
    else:
        _LOGGER.info(
            "Migrated %s config entry to version %d",
            config_entry.title,
            ENTRY_VERSION,
        )
    return True


async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Unload a config entry."""
    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

    if unload_ok:
        manager = getattr(entry, "runtime_data", None)
        if manager is not None:
            await manager.async_stop()
            entry.runtime_data = None
        hass.data[DOMAIN].pop(entry.entry_id)

    return unload_ok