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 (
    CONF_ADB_ENABLED,
    CONF_ADB_HOST,
    CONF_ADB_PORT,
    CONF_ANDROIDTV_ENTITY,
    CONF_PROTOCOL_ENABLED,
    CONF_PROTOCOL_HOST,
    DEFAULT_ADB_PORT,
    DEFAULT_NAME,
    DOMAIN,
    SERVICE_ADB_COMMAND,
)
from .helpers import discover_adb_host, protocol_cert_paths, resolve_adb_config

_LOGGER = logging.getLogger(__name__)

PLATFORMS = [Platform.MEDIA_PLAYER, Platform.SENSOR, Platform.REMOTE, Platform.BUTTON]


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 block.get(CONF_ADB_HOST) and not block.get(CONF_PROTOCOL_HOST):
            _LOGGER.warning(
                "nvidia_shield YAML block without protocol/adb host 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
    from .media_player import _ADB_COMMAND_SCHEMA, _async_adb_command_service

    hass.services.async_register(
        DOMAIN,
        SERVICE_ADB_COMMAND,
        _async_adb_command_service,
        schema=_ADB_COMMAND_SCHEMA,
    )
    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; ADB comes only when the
        # block carries an adb host (interactive cosmetics stay in the flow).
        seed = {
            "name": name,
            CONF_ADB_ENABLED: bool(block.get(CONF_ADB_HOST)),
            CONF_ADB_HOST: str(block.get(CONF_ADB_HOST) or ""),
            CONF_ADB_PORT: int(block.get(CONF_ADB_PORT) or DEFAULT_ADB_PORT),
            # protocol fields are prefill hints ONLY — pairing stays interactive
            CONF_PROTOCOL_HOST: str(block.get(CONF_PROTOCOL_HOST) or ""),
        }
        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.

    Falls back to the androidtv-remote host discovery (the ADB host fallback)
    until the user pairs, after which the explicit protocol host persists.
    """
    if not entry.options.get(CONF_PROTOCOL_ENABLED, False):
        return None

    host = entry.options.get(CONF_PROTOCOL_HOST) or discover_adb_host(
        hass, entry.data.get(CONF_ANDROIDTV_ENTITY)
    )
    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 v1 entries to v2 (adds the remote-protocol fields).

    ADB options carry over untouched; the ADB-host discovery fallback keeps
    serving as the protocol host until pairing populates it. Downgrades
    (version > 2) are rejected so HA offers a rollback.
    """
    if config_entry.version > 2:
        return False
    if config_entry.version < 2:
        options = {**config_entry.options}
        options.setdefault(CONF_PROTOCOL_ENABLED, False)
        options.setdefault(CONF_PROTOCOL_HOST, "")
        hass.config_entries.async_update_entry(
            config_entry, options=options, version=2
        )
        _LOGGER.info(
            "Migrated %s config entry to version 2", config_entry.title
        )
    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