adaptive-lighting/custom_components/adaptive_lighting/adaptation_utils.py

"""Utility functions for adaptation commands."""

import logging
from collections.abc import AsyncGenerator
from dataclasses import dataclass
from enum import IntFlag, auto
from typing import Any

from homeassistant.components.light import (
    ATTR_BRIGHTNESS,
    ATTR_BRIGHTNESS_PCT,
    ATTR_BRIGHTNESS_STEP,
    ATTR_BRIGHTNESS_STEP_PCT,
    ATTR_COLOR_NAME,
    ATTR_COLOR_TEMP_KELVIN,
    ATTR_EFFECT,
    ATTR_FLASH,
    ATTR_HS_COLOR,
    ATTR_RGB_COLOR,
    ATTR_RGBW_COLOR,
    ATTR_RGBWW_COLOR,
    ATTR_TRANSITION,
    ATTR_XY_COLOR,
)
from homeassistant.const import ATTR_ENTITY_ID
from homeassistant.core import Context, HomeAssistant, State

_LOGGER = logging.getLogger(__name__)

COLOR_ATTRS = {  # Should ATTR_PROFILE be in here?
    ATTR_COLOR_NAME,
    ATTR_COLOR_TEMP_KELVIN,
    ATTR_HS_COLOR,
    ATTR_RGB_COLOR,
    ATTR_XY_COLOR,
    ATTR_RGBW_COLOR,
    ATTR_RGBWW_COLOR,
}


BRIGHTNESS_ATTRS = {
    ATTR_BRIGHTNESS,
    ATTR_BRIGHTNESS_PCT,
    ATTR_BRIGHTNESS_STEP,
    ATTR_BRIGHTNESS_STEP_PCT,
}

# Worst-case rounding error when Home Assistant's 0-255 brightness scale
# round-trips through a device with coarser resolution (e.g., the 0-99 Z-Wave
# Multilevel Switch scale). A light cannot report back a value more precise than
# its own scale, so exact equality would never hold for such targets and
# 'skip_redundant_commands' would keep sending them forever. The tolerance sits
# far below the manual-control-detection threshold (BRIGHTNESS_CHANGE = 25), so
# it cannot mask a genuine user change.
BRIGHTNESS_TOLERANCE = 2

ServiceData = dict[str, Any]


class LightControlAttributes(IntFlag):
    """Attributes of lights that the adaptation engine can control."""

    NONE = 0
    BRIGHTNESS = auto()
    COLOR = auto()

    ALL = BRIGHTNESS | COLOR

    def __str__(self) -> str:
        """Return a string representation of the attributes."""
        if self == LightControlAttributes.NONE:
            return "NONE"

        return "|".join(
            member.name
            for member in type(self)
            if member is not LightControlAttributes.NONE
            and member in self
            and member.name is not None
        )

    def has_any(self) -> bool:
        """Determine whether any attribute is selected."""
        return self != LightControlAttributes.NONE

    def has_none(self) -> bool:
        """Determine whether no attribute is selected."""
        return self == LightControlAttributes.NONE

    def has_all(self) -> bool:
        """Determine whether all attributes are selected."""
        return (self & LightControlAttributes.ALL) == LightControlAttributes.ALL


def _split_service_call_data(service_data: ServiceData) -> list[ServiceData]:
    """Splits the service data by the adapted attributes.

    i.e., into separate data items for brightness and color.
    """
    common_attrs = {ATTR_ENTITY_ID}
    common_data = {k: service_data[k] for k in common_attrs if k in service_data}

    attributes_split_sequence = [BRIGHTNESS_ATTRS, COLOR_ATTRS]
    service_datas: list[dict[str, Any]] = []

    for attributes in attributes_split_sequence:
        split_data = {
            attribute: service_data[attribute]
            for attribute in attributes
            if service_data.get(attribute)
        }
        if split_data:
            service_datas.append(common_data | split_data)

    # Distribute the transition duration across all service calls
    if service_datas and (transition := service_data.get(ATTR_TRANSITION)) is not None:
        transition /= len(service_datas)

        for _service_data in service_datas:
            _service_data[ATTR_TRANSITION] = transition

    return service_datas


def _is_attribute_satisfied(key: str, value: Any, attributes: dict[str, Any]) -> bool:
    """Whether the light's current state already satisfies this target value."""
    if key not in attributes:
        return False
    current = attributes[key]
    if not isinstance(current, (int, float)) or not isinstance(value, (int, float)):
        return value == current
    if key == ATTR_BRIGHTNESS:
        return abs(value - current) <= BRIGHTNESS_TOLERANCE
    if key == ATTR_COLOR_TEMP_KELVIN and value > 0 and current > 0:
        # Compare in mired space: most integrations quantize color temperature
        # to whole mireds, and the kelvin error of that quantization grows
        # quadratically with kelvin (~21 K at 6500 K, ~50 K at 10000 K), so no
        # fixed kelvin tolerance fits the whole range. The tolerance of one
        # mired absorbs the difference between conversion schemes: HA core's
        # helpers floor (e.g. 5500 K -> 181 mired -> 5524 K) while some
        # integrations round (5500 K -> 182 mired -> 5495 K), and no exact
        # equality converges for both. One mired is far below the ~5.5 mired
        # just-noticeable difference for color temperature.
        return abs(round(1_000_000 / value) - round(1_000_000 / current)) <= 1
    return value == current


def _remove_redundant_attributes(
    service_data: ServiceData,
    state: State,
) -> ServiceData:
    """Filter service data by removing attributes already satisfied by the state.

    Removes all attributes from service call data whose values are already present
    in the target entity's state. Quantized attributes (brightness, color temp) are
    compared with a small tolerance: a light whose resolution is coarser than Home
    Assistant's cannot report back the exact value it was given, so exact equality
    would never hold and the attribute would never be filtered.
    """
    attributes: dict[str, Any] = dict(state.attributes)
    return {
        k: v
        for k, v in service_data.items()
        if not _is_attribute_satisfied(k, v, attributes)
    }


def _has_relevant_service_data_attributes(service_data: ServiceData) -> bool:
    """Determines whether the service data justifies an adaptation service call.

    A service call is not justified for data which does not contain any entries that
    change relevant attributes of an adapting entity, e.g., brightness or color.
    """
    common_attrs = {ATTR_ENTITY_ID, ATTR_TRANSITION}

    return any(attr not in common_attrs for attr in service_data)


async def _create_service_call_data_iterator(
    hass: HomeAssistant,
    service_datas: list[ServiceData],
    filter_by_state: bool,
) -> AsyncGenerator[ServiceData]:
    """Enumerates and filters a list of service datas on the fly.

    If filtering is enabled, every service data is filtered by the current state of
    the related entity and only returned if it contains relevant data that justifies
    a service call.
    The main advantage of this generator over a list is that it applies the filter
    at the time when the service data is read instead of up front. This gives greater
    flexibility because entity states can change while the items are iterated.
    """
    for service_data in service_datas:
        if filter_by_state and (entity_id := service_data.get(ATTR_ENTITY_ID)):
            current_entity_state = hass.states.get(entity_id)

            # Filter data to remove attributes that equal the current state
            if current_entity_state is not None:
                service_data = _remove_redundant_attributes(  # noqa: PLW2901
                    service_data,
                    state=current_entity_state,
                )

            # Emit service data if it still contains relevant attributes (else try next)
            if _has_relevant_service_data_attributes(service_data):
                yield service_data
        else:
            yield service_data


@dataclass
class AdaptationData:
    """Holds all data required to execute an adaptation."""

    entity_id: str
    context: Context
    sleep_time: float
    service_call_datas: AsyncGenerator[ServiceData]
    force: bool
    max_length: int
    attributes: LightControlAttributes
    initial_sleep: bool = False

    async def next_service_call_data(self) -> ServiceData | None:
        """Return data for the next service call, or none if no more data exists."""
        return await anext(self.service_call_datas, None)

    def __str__(self) -> str:
        """Return a string representation of the data."""
        return (
            f"{self.__class__.__name__}("
            f"entity_id={self.entity_id}, "
            f"context_id={self.context.id}, "
            f"sleep_time={self.sleep_time}, "
            f"force={self.force}, "
            f"max_length={self.max_length}, "
            f"attributes={self.attributes}, "
            f"initial_sleep={self.initial_sleep}"
            ")"
        )


class NoColorOrBrightnessInServiceDataError(Exception):
    """Exception raised when no color or brightness attributes are found in service data."""


def _identify_light_control_attributes(
    service_data: ServiceData,
) -> LightControlAttributes:
    """Extract the 'which' attribute from the service data."""
    has_brightness = ATTR_BRIGHTNESS in service_data
    has_color = any(attr in service_data for attr in COLOR_ATTRS)

    parameters = LightControlAttributes.NONE

    if has_brightness:
        parameters |= LightControlAttributes.BRIGHTNESS
    if has_color:
        parameters |= LightControlAttributes.COLOR

    if parameters == LightControlAttributes.NONE:
        msg = f"Invalid service_data, no brightness or color attributes found: {service_data=}"
        raise NoColorOrBrightnessInServiceDataError(msg)

    return parameters


def prepare_adaptation_data(
    hass: HomeAssistant,
    entity_id: str,
    context: Context,
    transition: float | None,
    split_delay: float,
    service_data: ServiceData,
    split: bool,
    filter_by_state: bool,
    force: bool,
    already_applied: LightControlAttributes = LightControlAttributes.NONE,
) -> AdaptationData:
    """Prepares a data object carrying all data required to execute an adaptation."""
    _LOGGER.debug(
        "Preparing adaptation data for %s with service data %s",
        entity_id,
        service_data,
    )
    service_datas = _split_service_call_data(service_data) if split else [service_data]

    service_datas_length = len(service_datas)

    if transition is not None:
        transition_duration_per_data = transition / max(1, service_datas_length)
        sleep_time = transition_duration_per_data + split_delay
    else:
        sleep_time = split_delay

    # Keep the original split timing, but omit attributes carried by an
    # intercepted turn-on shared with other lights. Do this before state
    # filtering: members can have different brightness/color already satisfied.
    applied_attrs = (
        BRIGHTNESS_ATTRS
        if LightControlAttributes.BRIGHTNESS in already_applied
        else set()
    ) | (COLOR_ATTRS if LightControlAttributes.COLOR in already_applied else set())
    if applied_attrs:
        service_datas = [
            {key: value for key, value in data.items() if key not in applied_attrs}
            for data in service_datas
        ]
        service_datas = [
            data
            for data in service_datas
            if _has_relevant_service_data_attributes(data)
        ]

    service_data_iterator = _create_service_call_data_iterator(
        hass,
        service_datas,
        filter_by_state,
    )

    attributes = _identify_light_control_attributes(service_data)

    return AdaptationData(
        entity_id=entity_id,
        context=context,
        sleep_time=sleep_time,
        service_call_datas=service_data_iterator,
        force=force,
        max_length=len(service_datas),
        attributes=attributes & ~already_applied,
    )


def manual_control_event_attribute_to_flags(
    manual_control_attribute: bool | str,
) -> LightControlAttributes:
    """Convert manual control event data to light control attributes."""
    if isinstance(manual_control_attribute, bool) and manual_control_attribute:
        return LightControlAttributes.ALL
    if manual_control_attribute == "brightness":
        return LightControlAttributes.BRIGHTNESS
    if manual_control_attribute == "color":
        return LightControlAttributes.COLOR
    return LightControlAttributes.NONE


def has_brightness_attribute(
    service_data: ServiceData,
) -> bool:
    """Determine whether the service data contains brightness attributes."""
    return any(attr in BRIGHTNESS_ATTRS for attr in service_data)


def has_color_attribute(
    service_data: ServiceData,
) -> bool:
    """Determine whether the service data contains color attributes."""
    return any(attr in COLOR_ATTRS for attr in service_data)


def has_effect_attribute(
    service_data: ServiceData,
) -> bool:
    """Determine whether the service data contains effect attributes."""
    return ATTR_FLASH in service_data or ATTR_EFFECT in service_data


def get_light_control_attributes(
    service_data: ServiceData,
) -> LightControlAttributes:
    """Get the light control attributes affected by the service call data."""
    parameters = LightControlAttributes.NONE

    if has_brightness_attribute(service_data):
        parameters |= LightControlAttributes.BRIGHTNESS

    if has_color_attribute(service_data):
        parameters |= LightControlAttributes.COLOR

    if has_effect_attribute(service_data):
        parameters |= LightControlAttributes.BRIGHTNESS
        parameters |= LightControlAttributes.COLOR

    return parameters