HomeAssistantVS/custom_components/powercalc/sensors/group/config_entry_utils.py

import inspect
import logging
from typing import Any

from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry, ConfigFlow
from homeassistant.const import CONF_NAME
from homeassistant.core import HomeAssistant, callback

from custom_components.powercalc.const import (
    CONF_GROUP,
    CONF_GROUP_MEMBER_SENSORS,
    CONF_GROUP_TYPE,
    CONF_SENSOR_TYPE,
    CONF_SUB_GROUPS,
    DOMAIN,
    ENTRY_GLOBAL_CONFIG_UNIQUE_ID,
    GroupType,
    SensorType,
)

_LOGGER = logging.getLogger(__name__)


def remove_power_sensor_from_associated_groups(
    hass: HomeAssistant,
    config_entry: ConfigEntry,
) -> list[ConfigEntry]:
    """When the user removes a virtual power config entry, update all groups this sensor belongs to."""
    group_entries = get_groups_having_member(hass, config_entry)

    for group_entry in group_entries:
        # Copy the list so we don't mutate the config entry's data in place, which would make
        # async_update_entry's change detection see no change and skip persisting/reloading.
        member_sensors = list(group_entry.data.get(CONF_GROUP_MEMBER_SENSORS) or [])
        member_sensors.remove(config_entry.entry_id)

        hass.config_entries.async_update_entry(
            group_entry,
            data={**group_entry.data, CONF_GROUP_MEMBER_SENSORS: member_sensors},
        )

    return group_entries


async def add_to_associated_groups(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
    """
    When the user has set a group on a virtual power config entry,
    we need to add this config entry to the group members sensors and update the group.
    """
    sensor_type = config_entry.data.get(CONF_SENSOR_TYPE)
    if sensor_type not in [SensorType.VIRTUAL_POWER, SensorType.DAILY_ENERGY]:
        return

    raw_groups = config_entry.data.get(CONF_GROUP)
    if not raw_groups:
        return

    group_ids = raw_groups if isinstance(raw_groups, list) else [raw_groups]
    for group_entry_id in group_ids:
        group_entry = await add_to_associated_group(hass, config_entry, group_entry_id)
        if group_entry:
            _LOGGER.debug(
                "ConfigEntry %s: Added to group %s.",
                config_entry.title,
                group_entry.title,
            )

    # After processed correctly we can want to unset the group, to prevent is being processed again
    new_data = {k: v for k, v in config_entry.data.items() if k != CONF_GROUP}
    hass.config_entries.async_update_entry(config_entry, data=new_data)


async def add_to_associated_group(
    hass: HomeAssistant,
    config_entry: ConfigEntry,
    group_entry_id: str,
) -> ConfigEntry | None:
    """When the user has set a group on a virtual power config entry,
    we need to add this config entry to the group members sensors and update the group.
    """
    group_entry = hass.config_entries.async_get_entry(group_entry_id)

    # When we are not dealing with a uuid, the user has set a group name manually
    # Create a new group entry for this group
    if not group_entry and len(group_entry_id) != 32:
        group_entry = hass.config_entries.async_entry_for_domain_unique_id(DOMAIN, group_entry_id)
        if not group_entry:
            additional_args: dict[str, Any] = {}
            signature = inspect.signature(ConfigEntry.__init__)
            if "discovery_keys" in signature.parameters:
                additional_args["discovery_keys"] = {}

            if "subentries_data" in signature.parameters:
                additional_args["subentries_data"] = None

            group_entry = ConfigEntry(
                version=ConfigFlow.VERSION,
                minor_version=ConfigFlow.MINOR_VERSION,
                domain=DOMAIN,
                source=SOURCE_IMPORT,
                title=group_entry_id,
                data={
                    CONF_SENSOR_TYPE: SensorType.GROUP,
                    CONF_NAME: group_entry_id,
                },
                options={},
                unique_id=group_entry_id,
                **additional_args,
            )
            await hass.config_entries.async_add(group_entry)

    if not group_entry:
        _LOGGER.warning(
            "ConfigEntry %s: Cannot add/remove to group %s. It does not exist.",
            config_entry.title,
            group_entry_id,
        )
        return None

    member_sensors = set(group_entry.data.get(CONF_GROUP_MEMBER_SENSORS) or [])

    # Config entry has already been added to associated group. just skip adding it again
    if config_entry.entry_id in member_sensors:
        return None

    member_sensors.add(config_entry.entry_id)
    hass.config_entries.async_update_entry(
        group_entry,
        data={**group_entry.data, CONF_GROUP_MEMBER_SENSORS: list(member_sensors)},
    )
    return group_entry


def get_entries_having_subgroup(hass: HomeAssistant, subgroup_entry: ConfigEntry) -> list[ConfigEntry]:
    """Get all virtual power entries which have the subgroup in their subgroups list."""
    return [
        entry for entry in get_group_entries(hass) if subgroup_entry.entry_id in (entry.data.get(CONF_SUB_GROUPS) or [])
    ]


def get_groups_having_member(hass: HomeAssistant, member_entry: ConfigEntry) -> list[ConfigEntry]:
    """Get all group entries which have the member sensor in their member list."""
    return [
        entry
        for entry in hass.config_entries.async_entries(DOMAIN)
        if entry.data.get(CONF_SENSOR_TYPE) == SensorType.GROUP
        and member_entry.entry_id in (entry.data.get(CONF_GROUP_MEMBER_SENSORS) or [])
    ]


@callback
def get_group_entries(hass: HomeAssistant, group_type: GroupType | None = None) -> list[ConfigEntry]:
    return [
        entry
        for entry in hass.config_entries.async_entries(DOMAIN)
        if entry.data.get(CONF_SENSOR_TYPE) == SensorType.GROUP
        and (group_type is None or entry.data.get(CONF_GROUP_TYPE, GroupType.CUSTOM) == group_type)
    ]


@callback
def get_entries_excluding_global_config(hass: HomeAssistant) -> list[ConfigEntry]:
    return [
        entry for entry in hass.config_entries.async_entries(DOMAIN) if entry.unique_id != ENTRY_GLOBAL_CONFIG_UNIQUE_ID
    ]