core/homeassistant/components/energyid/__init__.py

"""The EnergyID integration."""

from __future__ import annotations

from dataclasses import dataclass
import datetime as dt
from datetime import timedelta
import functools
import logging

from aiohttp import ClientError, ClientResponseError
from energyid_webhooks.client_v2 import WebhookClient

from homeassistant.config_entries import ConfigEntry
from homeassistant.const import STATE_UNAVAILABLE, STATE_UNKNOWN
from homeassistant.core import (
    CALLBACK_TYPE,
    Event,
    EventStateChangedData,
    HomeAssistant,
    callback,
)
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.event import (
    async_track_entity_registry_updated_event,
    async_track_state_change_event,
    async_track_time_interval,
)

from .const import (
    CONF_DEVICE_ID,
    CONF_DEVICE_NAME,
    CONF_ENERGYID_KEY,
    CONF_HA_ENTITY_UUID,
    CONF_PROVISIONING_KEY,
    CONF_PROVISIONING_SECRET,
    DOMAIN,
)

_LOGGER = logging.getLogger(__name__)

type EnergyIDConfigEntry = ConfigEntry[EnergyIDRuntimeData]

DEFAULT_UPLOAD_INTERVAL_SECONDS = 60


@dataclass
class EnergyIDRuntimeData:
    """Runtime data for the EnergyID integration."""

    client: WebhookClient
    mappings: dict[str, str]
    state_listener: CALLBACK_TYPE | None = None
    registry_tracking_listener: CALLBACK_TYPE | None = None
    unavailable_logged: bool = False


async def async_setup_entry(hass: HomeAssistant, entry: EnergyIDConfigEntry) -> bool:
    """Set up EnergyID from a config entry."""
    session = async_get_clientsession(hass)
    client = WebhookClient(
        provisioning_key=entry.data[CONF_PROVISIONING_KEY],
        provisioning_secret=entry.data[CONF_PROVISIONING_SECRET],
        device_id=entry.data[CONF_DEVICE_ID],
        device_name=entry.data[CONF_DEVICE_NAME],
        session=session,
    )

    entry.runtime_data = EnergyIDRuntimeData(
        client=client,
        mappings={},
    )

    is_claimed = None
    try:
        is_claimed = await client.authenticate()
    except TimeoutError as err:
        raise ConfigEntryNotReady(
            f"Timeout authenticating with EnergyID: {err}"
        ) from err
    except ClientResponseError as err:
        # 401/403 = invalid credentials, trigger reauth
        if err.status in (401, 403):
            raise ConfigEntryAuthFailed(f"Invalid credentials: {err}") from err
        # Other HTTP errors are likely temporary
        raise ConfigEntryNotReady(
            f"HTTP error authenticating with EnergyID: {err}"
        ) from err
    except ClientError as err:
        # Network/connection errors are temporary
        raise ConfigEntryNotReady(
            f"Connection error authenticating with EnergyID: {err}"
        ) from err
    except Exception as err:
        # Unknown errors - log and retry (safer than forcing reauth)
        _LOGGER.exception("Unexpected error during EnergyID authentication")
        raise ConfigEntryNotReady(
            f"Unexpected error authenticating with EnergyID: {err}"
        ) from err

    if not is_claimed:
        # Device exists but not claimed = user needs to claim it = auth issue
        raise ConfigEntryAuthFailed("Device is not claimed. Please re-authenticate.")

    _LOGGER.debug("EnergyID device '%s' authenticated successfully", client.device_name)

    async def _async_synchronize_sensors(now: dt.datetime | None = None) -> None:
        """Callback for periodically synchronizing sensor data."""
        try:
            await client.synchronize_sensors()
            if entry.runtime_data.unavailable_logged:
                _LOGGER.debug("Connection to EnergyID re-established")
                entry.runtime_data.unavailable_logged = False
        except (OSError, RuntimeError) as err:
            if not entry.runtime_data.unavailable_logged:
                _LOGGER.debug("EnergyID is unavailable: %s", err)
                entry.runtime_data.unavailable_logged = True

    upload_interval = DEFAULT_UPLOAD_INTERVAL_SECONDS
    if client.webhook_policy:
        upload_interval = client.webhook_policy.get(
            "uploadInterval", DEFAULT_UPLOAD_INTERVAL_SECONDS
        )

    # Schedule the callback and automatically unsubscribe when the entry is unloaded.
    entry.async_on_unload(
        async_track_time_interval(
            hass, _async_synchronize_sensors, timedelta(seconds=upload_interval)
        )
    )
    entry.async_on_unload(entry.add_update_listener(config_entry_update_listener))

    update_listeners(hass, entry)

    _LOGGER.debug(
        "Starting EnergyID background sync for '%s'",
        client.device_name,
    )

    return True


async def config_entry_update_listener(
    hass: HomeAssistant, entry: EnergyIDConfigEntry
) -> None:
    """Handle config entry updates, including subentry changes."""
    _LOGGER.debug("Config entry updated for %s, reloading listeners", entry.entry_id)
    update_listeners(hass, entry)


@callback
def update_listeners(hass: HomeAssistant, entry: EnergyIDConfigEntry) -> None:
    """Set up or update state listeners and queue initial states."""
    runtime_data = entry.runtime_data
    client = runtime_data.client

    # Clean up old state listener
    if runtime_data.state_listener:
        runtime_data.state_listener()
        runtime_data.state_listener = None

    mappings: dict[str, str] = {}
    entities_to_track: list[str] = []
    old_mappings = set(runtime_data.mappings.keys())
    new_mappings = set()
    ent_reg = er.async_get(hass)

    subentries = list(entry.subentries.values())
    _LOGGER.debug(
        "Found %d subentries in entry.subentries: %s",
        len(subentries),
        [s.data for s in subentries],
    )

    # Build current entity mappings
    tracked_entity_ids = []
    for subentry in subentries:
        entity_uuid = subentry.data.get(CONF_HA_ENTITY_UUID)
        energyid_key = subentry.data.get(CONF_ENERGYID_KEY)

        if not (entity_uuid and energyid_key):
            continue

        entity_entry = ent_reg.async_get(entity_uuid)
        if not entity_entry:
            _LOGGER.warning(
                "Entity with UUID %s does not exist, skipping mapping to %s",
                entity_uuid,
                energyid_key,
            )
            continue

        ha_entity_id = entity_entry.entity_id
        tracked_entity_ids.append(ha_entity_id)

        if not hass.states.get(ha_entity_id):
            # Entity exists in registry but is not present in the state machine
            _LOGGER.debug(
                "Entity %s does not exist in state machine yet, will track when available (mapping to %s)",
                ha_entity_id,
                energyid_key,
            )
            # Still add to entities_to_track so we can handle it when state appears
            entities_to_track.append(ha_entity_id)
            continue

        mappings[ha_entity_id] = energyid_key
        entities_to_track.append(ha_entity_id)
        new_mappings.add(ha_entity_id)
        client.get_or_create_sensor(energyid_key)

        if ha_entity_id not in old_mappings:
            _LOGGER.debug(
                "New mapping detected for %s, queuing initial state", ha_entity_id
            )
            if (
                current_state := hass.states.get(ha_entity_id)
            ) and current_state.state not in (
                STATE_UNKNOWN,
                STATE_UNAVAILABLE,
            ):
                try:
                    value = float(current_state.state)
                    timestamp = current_state.last_updated or dt.datetime.now(dt.UTC)
                    client.get_or_create_sensor(energyid_key).update(value, timestamp)
                except ValueError, TypeError:
                    _LOGGER.debug(
                        "Could not convert initial state of %s to float: %s",
                        ha_entity_id,
                        current_state.state,
                    )

    # Clean up old entity registry listener
    if runtime_data.registry_tracking_listener:
        runtime_data.registry_tracking_listener()
        runtime_data.registry_tracking_listener = None

    # Set up listeners for entity registry changes
    if tracked_entity_ids:
        _LOGGER.debug("Setting up entity registry tracking for: %s", tracked_entity_ids)

        def _handle_entity_registry_change(
            event: Event[er.EventEntityRegistryUpdatedData],
        ) -> None:
            """Handle entity registry changes for our tracked entities."""
            _LOGGER.debug("Registry event for tracked entity: %s", event.data)

            if event.data["action"] == "update":
                # Type is now narrowed to _EventEntityRegistryUpdatedData_Update
                if "entity_id" in event.data["changes"]:
                    old_entity_id = event.data["changes"]["entity_id"]
                    new_entity_id = event.data["entity_id"]

                    _LOGGER.debug(
                        "Tracked entity ID changed: %s -> %s",
                        old_entity_id,
                        new_entity_id,
                    )
                    # Entity ID changed, need to reload listeners to track new ID
                    update_listeners(hass, entry)

            elif event.data["action"] == "remove":
                _LOGGER.debug("Tracked entity removed: %s", event.data["entity_id"])
                # reminder: Create repair issue to notify user about removed entity
                update_listeners(hass, entry)

        # Track the specific entity IDs we care about
        unsub_entity_registry = async_track_entity_registry_updated_event(
            hass, tracked_entity_ids, _handle_entity_registry_change
        )
        runtime_data.registry_tracking_listener = unsub_entity_registry

    if removed_mappings := old_mappings - new_mappings:
        _LOGGER.debug("Removed mappings: %s", ", ".join(removed_mappings))

    runtime_data.mappings = mappings

    if not entities_to_track:
        _LOGGER.debug(
            "No valid sensor mappings configured for '%s'", client.device_name
        )
        return

    unsub_state_change = async_track_state_change_event(
        hass,
        entities_to_track,
        functools.partial(_async_handle_state_change, hass, entry.entry_id),
    )
    runtime_data.state_listener = unsub_state_change

    _LOGGER.debug(
        "Now tracking state changes for %d entities for '%s': %s",
        len(entities_to_track),
        client.device_name,
        entities_to_track,
    )


@callback
def _async_handle_state_change(
    hass: HomeAssistant, entry_id: str, event: Event[EventStateChangedData]
) -> None:
    """Handle state changes for tracked entities."""
    entity_id = event.data["entity_id"]
    new_state = event.data["new_state"]

    _LOGGER.debug(
        "State change detected for entity: %s, new value: %s",
        entity_id,
        new_state.state if new_state else "None",
    )

    if not new_state or new_state.state in (STATE_UNKNOWN, STATE_UNAVAILABLE):
        return

    entry = hass.config_entries.async_get_entry(entry_id)
    if not entry or not hasattr(entry, "runtime_data"):
        # Entry is being unloaded or not yet fully initialized
        return

    runtime_data = entry.runtime_data
    client = runtime_data.client

    # Check if entity is already mapped
    if energyid_key := runtime_data.mappings.get(entity_id):
        # Entity already mapped, just update value
        _LOGGER.debug(
            "Updating EnergyID sensor %s with value %s", energyid_key, new_state.state
        )
    else:
        # Entity not mapped yet - check if it should be (handles late-appearing entities)
        ent_reg = er.async_get(hass)
        for subentry in entry.subentries.values():
            entity_uuid = subentry.data.get(CONF_HA_ENTITY_UUID)
            energyid_key_candidate = subentry.data.get(CONF_ENERGYID_KEY)

            if not (entity_uuid and energyid_key_candidate):
                continue

            entity_entry = ent_reg.async_get(entity_uuid)
            if entity_entry and entity_entry.entity_id == entity_id:
                # Found it! Add to mappings and send initial value
                energyid_key = energyid_key_candidate
                runtime_data.mappings[entity_id] = energyid_key
                client.get_or_create_sensor(energyid_key)
                _LOGGER.debug(
                    "Entity %s now available in state machine, adding to mappings (key: %s)",
                    entity_id,
                    energyid_key,
                )
                break
        else:
            # Not a tracked entity, ignore
            return

    try:
        value = float(new_state.state)
    except ValueError, TypeError:
        return

    client.get_or_create_sensor(energyid_key).update(value, new_state.last_updated)


async def async_unload_entry(hass: HomeAssistant, entry: EnergyIDConfigEntry) -> bool:
    """Unload a config entry."""
    _LOGGER.debug("Unloading EnergyID entry for %s", entry.title)

    try:
        # Unload subentries if present (guarded for test and reload scenarios)
        if hasattr(hass.config_entries, "async_entries") and hasattr(entry, "entry_id"):
            subentries = [
                e.entry_id
                for e in hass.config_entries.async_entries(DOMAIN)
                if getattr(e, "parent_entry", None) == entry.entry_id
            ]
            for subentry_id in subentries:
                await hass.config_entries.async_unload(subentry_id)

        # Only clean up listeners and client if runtime_data is present
        if hasattr(entry, "runtime_data"):
            runtime_data = entry.runtime_data

            # Remove state listener
            if runtime_data.state_listener:
                runtime_data.state_listener()

            # Remove registry tracking listener
            if runtime_data.registry_tracking_listener:
                runtime_data.registry_tracking_listener()

            try:
                await runtime_data.client.close()
            except Exception:
                _LOGGER.exception("Error closing EnergyID client for %s", entry.title)
            del entry.runtime_data
    except Exception:
        _LOGGER.exception("Error during async_unload_entry for %s", entry.title)
        return False
    return True