HomeAssistantVS/custom_components/powercalc/sensors/abstract.py

from __future__ import annotations

import logging
from typing import cast

from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
from homeassistant.const import CONF_NAME
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.device_registry import DeviceEntry
from homeassistant.helpers.entity import Entity, async_generate_entity_id
import homeassistant.helpers.entity_registry as er
from homeassistant.helpers.typing import ConfigType

from custom_components.powercalc.common import SourceEntity
from custom_components.powercalc.const import (
    CONF_COST_SENSOR_FRIENDLY_NAMING,
    CONF_COST_SENSOR_NAMING,
    CONF_ENERGY_SENSOR_FRIENDLY_NAMING,
    CONF_ENERGY_SENSOR_NAMING,
    CONF_POWER_SENSOR_FRIENDLY_NAMING,
    CONF_POWER_SENSOR_NAMING,
    DEFAULT_COST_NAME_PATTERN,
    DEFAULT_ENERGY_NAME_PATTERN,
    DEFAULT_POWER_NAME_PATTERN,
    DOMAIN,
)
from custom_components.powercalc.device_binding import bind_entity_to_registry_metadata

ENTITY_ID_FORMAT = SENSOR_DOMAIN + ".{}"

_LOGGER = logging.getLogger(__name__)


class BaseEntity(Entity):
    async def async_added_to_hass(self) -> None:
        """Bind configured registry metadata."""
        await super().async_added_to_hass()

        bind_entity_to_registry_metadata(
            self.hass,
            self.entity_id,
            cast(DeviceEntry | None, getattr(self, "device_entry", None)),
            cast(ConfigType | None, getattr(self, "_sensor_config", None)),
        )


def generate_power_sensor_name(
    sensor_config: ConfigType,
    name: str | None = None,
    source_entity: SourceEntity | None = None,
) -> str:
    """Generates the name to use for a power sensor."""
    return _generate_sensor_name(
        sensor_config,
        CONF_POWER_SENSOR_NAMING,
        CONF_POWER_SENSOR_FRIENDLY_NAMING,
        DEFAULT_POWER_NAME_PATTERN,
        name,
        source_entity,
    )


def generate_energy_sensor_name(
    sensor_config: ConfigType,
    name: str | None = None,
    source_entity: SourceEntity | None = None,
) -> str:
    """Generates the name to use for an energy sensor."""
    return _generate_sensor_name(
        sensor_config,
        CONF_ENERGY_SENSOR_NAMING,
        CONF_ENERGY_SENSOR_FRIENDLY_NAMING,
        DEFAULT_ENERGY_NAME_PATTERN,
        name,
        source_entity,
    )


def generate_cost_sensor_name(
    sensor_config: ConfigType,
    name: str | None = None,
    source_entity: SourceEntity | None = None,
) -> str:
    """Generates the name to use for a cost sensor."""
    return _generate_sensor_name(
        sensor_config,
        CONF_COST_SENSOR_NAMING,
        CONF_COST_SENSOR_FRIENDLY_NAMING,
        DEFAULT_COST_NAME_PATTERN,
        name,
        source_entity,
    )


def _generate_sensor_name(
    sensor_config: ConfigType,
    naming_conf_key: str,
    friendly_naming_conf_key: str,
    default_pattern: str,
    name: str | None = None,
    source_entity: SourceEntity | None = None,
) -> str:
    """Generates the name to use for a sensor."""
    if name is None and source_entity:
        name = source_entity.name

    if friendly_naming_conf_key in sensor_config:
        friendly_name_pattern = str(sensor_config.get(friendly_naming_conf_key))
        return friendly_name_pattern.format(name)

    name_pattern = str(sensor_config.get(naming_conf_key, default_pattern))
    return name_pattern.format(name)


@callback
def generate_power_sensor_entity_id(
    hass: HomeAssistant,
    sensor_config: ConfigType,
    source_entity: SourceEntity | None = None,
    name: str | None = None,
    unique_id: str | None = None,
) -> str:
    """Generates the entity_id to use for a power sensor."""
    return _generate_sensor_entity_id(
        hass,
        sensor_config,
        CONF_POWER_SENSOR_NAMING,
        DEFAULT_POWER_NAME_PATTERN,
        source_entity,
        name,
        unique_id,
    )


@callback
def generate_energy_sensor_entity_id(
    hass: HomeAssistant,
    sensor_config: ConfigType,
    source_entity: SourceEntity | None = None,
    name: str | None = None,
    unique_id: str | None = None,
) -> str:
    """Generates the entity_id to use for an energy sensor."""
    return _generate_sensor_entity_id(
        hass,
        sensor_config,
        CONF_ENERGY_SENSOR_NAMING,
        DEFAULT_ENERGY_NAME_PATTERN,
        source_entity,
        name,
        unique_id,
    )


@callback
def generate_cost_sensor_entity_id(
    hass: HomeAssistant,
    sensor_config: ConfigType,
    source_entity: SourceEntity | None = None,
    name: str | None = None,
    unique_id: str | None = None,
) -> str:
    """Generates the entity_id to use for a cost sensor."""
    return _generate_sensor_entity_id(
        hass,
        sensor_config,
        CONF_COST_SENSOR_NAMING,
        DEFAULT_COST_NAME_PATTERN,
        source_entity,
        name,
        unique_id,
    )


def _generate_sensor_entity_id(
    hass: HomeAssistant,
    sensor_config: ConfigType,
    naming_conf_key: str,
    default_pattern: str,
    source_entity: SourceEntity | None = None,
    name: str | None = None,
    unique_id: str | None = None,
) -> str:
    """Generates the entity_id to use for a sensor."""
    if entity_id := get_entity_id_by_unique_id(hass, unique_id):
        return entity_id
    name_pattern = str(sensor_config.get(naming_conf_key, default_pattern))
    object_id = name or sensor_config.get(CONF_NAME)
    if object_id is None and source_entity:
        object_id = source_entity.object_id
    return async_generate_entity_id(
        ENTITY_ID_FORMAT,
        name_pattern.format(object_id),
        hass=hass,
    )


def get_entity_id_by_unique_id(
    hass: HomeAssistant,
    unique_id: str | None,
) -> str | None:
    if unique_id is None:
        return None
    entity_reg = er.async_get(hass)
    return entity_reg.async_get_entity_id(SENSOR_DOMAIN, DOMAIN, unique_id)