heating-mpc-ha/custom_components/heating_mpc/__init__.py

"""Heating MPC custom integration, initial release shadow-only."""
from __future__ import annotations

from types import MappingProxyType

from homeassistant.const import Platform
from homeassistant.config_entries import ConfigEntry, ConfigSubentry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.storage import Store
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.event import async_track_state_change_event

from .const import CONF_HEATER, CONF_ROOMS, CONF_WEIGHTS, DOMAIN, CONF_OUTSIDE, WEATHER_SOURCES
from .weather_sources import sources_for
from .coordinator import HeatingMPCCoordinator, mark_model_removed, storage_lock
from .zones import SUBENTRY_ZONE, ZONE_KEY, ZONE_NAME, legacy_zone_data

PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.BINARY_SENSOR]


async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    coordinator = HeatingMPCCoordinator(hass, entry)
    await coordinator.async_restore_model()
    # Read-only: timestamps relay changes between polls for the radiator-lag estimate.
    unsub_relay = async_track_state_change_event(
        hass, [coordinator.config[CONF_HEATER]], coordinator.async_relay_changed)
    try:
        await coordinator.async_config_entry_first_refresh()
        hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
        await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
    except BaseException:
        # HA skips unload callbacks for unexpected setup errors, so clean up here,
        # including a refresh scheduled by entities that were already added.
        unsub_relay()
        hass.data.get(DOMAIN, {}).pop(entry.entry_id, None)
        await coordinator.async_shutdown()
        raise
    entry.async_on_unload(unsub_relay)
    # Remove only obsolete zone forecasts belonging to this entry/platform. Zones
    # deleted in the UI are also removed by HA through their subentry link.
    registry = er.async_get(hass)
    expected = {f"{entry.entry_id}_{zone}_predicted_60min" for zone in coordinator.rooms}
    for entity in er.async_entries_for_config_entry(registry, entry.entry_id):
        if (entity.platform == DOMAIN and
                entity.unique_id.startswith(f"{entry.entry_id}_") and
                entity.unique_id.endswith("_predicted_60min") and
                entity.unique_id not in expected):
            registry.async_remove(entity.entity_id)
    entry.async_on_unload(entry.add_update_listener(async_reload_entry))
    return True


async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
    await hass.config_entries.async_reload(entry.entry_id)


async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    coordinator = hass.data.get(DOMAIN, {}).get(entry.entry_id)
    if coordinator is not None:
        # Before anything is unloaded: an update still running must not publish or save.
        coordinator.async_begin_shutdown()
    ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
    if ok:
        hass.data[DOMAIN].pop(entry.entry_id, None)
    elif coordinator is not None:
        coordinator.async_cancel_shutdown()
    return ok


async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
    """Remove the learned model when this integration instance is deleted."""
    store = Store(hass, 1, f"{DOMAIN}.{entry.entry_id}")
    # Wait for a save still in progress, which would otherwise recreate the file;
    # a save still waiting for the lock sees the mark and does not write.
    async with storage_lock(hass, entry.entry_id):
        mark_model_removed(hass, entry.entry_id)
        await store.async_remove()
    hass.data.get(f"{DOMAIN}_storage_locks", {}).pop(entry.entry_id, None)


async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Migrate older entries without changing the user's selections.

    1 -> 2: single outdoor sensor becomes a prioritized source list.
    2 -> 3: each room sensor becomes a one-sensor zone subentry with its comfort
    weight. Its entity ID stays the zone key, so learned models, forecast entities
    and accuracy statistics carry over. Older versions cannot load version 3.
    """
    if entry.version > 3:
        return False
    if entry.version == 1:
        def migrate(config):
            result = dict(config)
            if CONF_OUTSIDE in result:
                result[WEATHER_SOURCES["temperature"]] = sources_for(result, "temperature")
                result.pop(CONF_OUTSIDE)
            return result
        hass.config_entries.async_update_entry(
            entry, data=migrate(entry.data), options=migrate(entry.options),
            version=2, minor_version=1,
        )
    if entry.version == 2:
        config = {**entry.data, **entry.options}
        rooms = [room for room in config.get(CONF_ROOMS, []) if isinstance(room, str)]
        registry = er.async_get(hass)
        names = {}
        for room in rooms:
            state, registered = hass.states.get(room), registry.async_get(room)
            names[room] = ((state.name if state else None)
                           or (registered.name or registered.original_name if registered else None)
                           or room)
        existing = {sub.unique_id for sub in entry.subentries.values()}
        for data in legacy_zone_data(rooms, config.get(CONF_WEIGHTS), names):
            if data[ZONE_KEY] not in existing:
                hass.config_entries.async_add_subentry(entry, ConfigSubentry(
                    data=MappingProxyType(data), subentry_type=SUBENTRY_ZONE,
                    title=data[ZONE_NAME], unique_id=data[ZONE_KEY]))

        def strip(config):
            return {k: v for k, v in config.items() if k not in (CONF_ROOMS, CONF_WEIGHTS)}
        hass.config_entries.async_update_entry(
            entry, data=strip(entry.data), options=strip(entry.options), version=3, minor_version=1)
    return True