HomeAssistantVS/custom_components/spook/__init__.py

"""Spook - Your homie. Custom integration for Home Assistant."""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

from homeassistant.config_entries import ConfigEntryState
from homeassistant.const import (
    EVENT_HOMEASSISTANT_START,
    EVENT_HOMEASSISTANT_STARTED,
    RESTART_EXIT_CODE,
)
from homeassistant.core import (
    CoreState,
    callback,
)
from homeassistant.helpers import issue_registry as ir

from .automation_runs import async_setup_automation_runs
from .const import DOMAIN, LOGGER, PLATFORMS
from .core_compat import async_clear_custom_components_cache
from .dismissals import async_setup_dismissals
from .entity_filtering import async_setup_all_entity_ids_cache_invalidation
from .integration_linking import link_sub_integrations, unlink_sub_integrations
from .listeners import async_listen_once_tracked
from .repairs import SpookRepairManager
from .run_history import async_setup_run_history
from .services import SpookServiceManager
from .setup_helpers import async_forward_setup_entry
from .statistics_sources import async_setup_abandoned_statistics_watching
from .timed_states import async_setup_timed_states

if TYPE_CHECKING:
    from homeassistant.config_entries import ConfigEntry
    from homeassistant.core import Event, HomeAssistant


def _has_entries_waiting(hass: HomeAssistant, domains: set[str]) -> bool:
    """Return if a config entry of these domains is waiting to be set up.

    When Home Assistant could not find the integration at startup, the entry
    was never tried and is left not loaded. Any other state means its code
    was found and ran, so a restart would not change a thing.
    """
    return any(
        entry.state is ConfigEntryState.NOT_LOADED
        for domain in domains
        for entry in hass.config_entries.async_entries(domain, include_disabled=False)
    )


@callback
def _async_restart_for_sub_integrations(hass: HomeAssistant) -> bool:
    """Restart Home Assistant, or ask for it, so sub integrations get set up.

    Returns if Spook can carry on setting up. Not when a restart is on its
    way, as everything set up now would be torn down again right after.
    """

    @callback
    def _restart(_: Event | None = None) -> None:
        """Restart Home Assistant."""
        hass.data["homeassistant_stop"] = asyncio.create_task(
            hass.async_stop(RESTART_EXIT_CODE),
        )

    # Should be OK to restart. Better to do it before anything else started.
    if hass.state == CoreState.starting:
        LOGGER.info("Restarting Home Assistant to set up Spook's helpers")
        _restart()
        return False

    # If all other fails, but we are not running yet... wait for it.
    if hass.state == CoreState.not_running:
        LOGGER.info("Restarting Home Assistant to set up Spook's helpers")
        # Listen to both... just in case.
        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_START, _restart)
        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STARTED, _restart)
        return False

    LOGGER.info(
        "Home Assistant needs to be restarted for Spook to complete setting up",
    )
    ir.async_create_issue(
        hass=hass,
        domain=DOMAIN,
        issue_id="restart_required",
        is_fixable=True,
        severity=ir.IssueSeverity.WARNING,
        translation_key="restart_required",
    )
    return True


async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Set up from a config entry."""
    # Symlink all sub integrations from Spook to the parent integrations folder.
    # This is a workaround for the fact that Home Assistant doesn't support
    # sub integrations. The loader only scans custom_components once, so a
    # freshly linked one needs it to look again, no restart required.
    if linked := await hass.async_add_executor_job(link_sub_integrations, hass):
        LOGGER.debug("Newly symlinked sub integrations, rescanning custom components")
        async_clear_custom_components_cache(hass)

        # A helper made with a sub integration before its link went missing
        # had nothing to set it up with when Home Assistant started, and that
        # is not tried again. Only a restart picks it up.
        waiting = _has_entries_waiting(hass, linked)
        if waiting and not _async_restart_for_sub_integrations(hass):
            return False

    # Forward async_setup_entry to ectoplasms
    await async_forward_setup_entry(hass, entry)

    # Before the services, because the actions that hold an automation in a
    # state reach for this the moment they are called, and registering them
    # first leaves a window where calling one finds nothing there.
    #
    # This also picks up automations held before the last restart. An
    # automation keeps whatever state it had, so without this a snooze that
    # spanned a restart would be a disable nobody remembers making.
    entry.async_on_unload(await async_setup_timed_states(hass))

    # Set up services. Unloading is registered first: setting up starts
    # listening before it is done, and Home Assistant runs these callbacks
    # when a setup fails or is cancelled too. Registered afterwards, a failed
    # setup would leave the manager listening and registering actions for a
    # Spook that never loaded.
    services = SpookServiceManager(hass)
    entry.async_on_unload(services.async_on_unload)
    await services.async_setup()

    # Watching before anything can start a wait it has to observe. A repair
    # can put a statistic on the clock the first time it looks, and on the
    # reload path that happens below rather than at some later event, so a
    # watcher registered afterwards would already have missed the beginning
    # of what it is meant to be watching.
    entry.async_on_unload(async_setup_abandoned_statistics_watching(hass))

    # Loaded before any repair looks, because the first thing a repair does
    # with what it finds is check whether somebody already said to leave it.
    entry.async_on_unload(await async_setup_dismissals(hass))

    # Who you gonna call? SpookRepairManager!
    repairs = SpookRepairManager(hass)

    # Starting the repairs takes a moment, and on a normal start that moment
    # comes after this setup is long done. Spook can be disabled or reloaded
    # in it, and the unload that runs then has not been told about repairs
    # that are not there yet. They finished starting afterwards with nothing
    # left to stop them.
    unloaded = False

    @callback
    def _note_the_unload() -> None:
        """Remember Spook went, for repairs that were still on their way."""
        nonlocal unloaded
        unloaded = True

    entry.async_on_unload(_note_the_unload)

    async def _ghost_busters(_: Event | None = None) -> None:
        """Send them in, time for some ghost chasing."""
        await repairs.async_setup()

        if unloaded:
            await repairs.async_on_unload()
            return

        entry.async_on_unload(repairs.async_on_unload)

    if hass.state == CoreState.running:
        # Home Assistant is already up, so the started event has been and gone.
        # This is the reload path, and setting up for the first time on a
        # running instance. Waiting for an event that cannot fire again would
        # leave every repair check dead until a core restart, silently.
        await _ghost_busters()
    else:
        # Otherwise wait for a settled instance, so repairs do not inspect a
        # half-loaded Home Assistant and report things that are still coming.
        entry.async_on_unload(
            async_listen_once_tracked(hass, EVENT_HOMEASSISTANT_STARTED, _ghost_busters)
        )

    # Set up platforms
    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)

    # Set up the all entity ids cache invalidation
    entry.async_on_unload(async_setup_all_entity_ids_cache_invalidation(hass))

    # Start noting which automation runs under which context. It has to begin
    # here rather than when a condition first asks: a condition inside an
    # action sequence is only built once that sequence runs, by which point
    # the automation it wants to know about has already been and gone.
    entry.async_on_unload(async_setup_automation_runs(hass))

    # And when they ran, for the conditions that count runs rather than
    # contexts. Same reason it cannot wait to be asked.
    entry.async_on_unload(async_setup_run_history(hass))

    # Yay, we didn't got spooked!
    return True


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


async def async_remove_entry(hass: HomeAssistant, _: ConfigEntry) -> None:
    """Remove a config entry."""
    await hass.async_add_executor_job(unlink_sub_integrations, hass)