core/homeassistant/components/lifx/coordinator.py

"""Coordinator for lifx."""

from __future__ import annotations

import asyncio
from collections.abc import Callable
from datetime import timedelta
from enum import IntEnum
from functools import cached_property, partial
from math import floor, log10
from typing import Any, cast

from aiolifx.aiolifx import (
    Light,
    Message,
    MultiZoneDirection,
    MultiZoneEffectType,
    TileEffectSkyType,
    TileEffectType,
)
from aiolifx.connection import LIFXConnection
from aiolifx_themes.themes import ThemeLibrary, ThemePainter
from awesomeversion import AwesomeVersion

from homeassistant.const import (
    SIGNAL_STRENGTH_DECIBELS,
    SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
    Platform,
)
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.debounce import Debouncer
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator

from .const import (
    _LOGGER,
    ATTR_REMAINING,
    DEFAULT_ATTEMPTS,
    DOMAIN,
    IDENTIFY_WAVEFORM,
    MAX_ATTEMPTS_PER_UPDATE_REQUEST_MESSAGE,
    MAX_UPDATE_TIME,
    MESSAGE_RETRIES,
    MESSAGE_TIMEOUT,
    OVERALL_TIMEOUT,
    TARGET_ANY,
    UNAVAILABLE_GRACE,
)
from .util import (
    async_execute_lifx,
    async_multi_execute_lifx_with_retries,
    get_real_mac_addr,
    infrared_brightness_option_to_value,
    infrared_brightness_value_to_option,
    lifx_features,
)

LIGHT_UPDATE_INTERVAL = 10
REQUEST_REFRESH_DELAY = 0.35
LIFX_IDENTIFY_DELAY = 3.0
ZONES_PER_COLOR_UPDATE_REQUEST = 8

RSSI_DBM_FW = AwesomeVersion("2.77")


class FirmwareEffect(IntEnum):
    """Enumeration of LIFX firmware effects."""

    OFF = 0
    MOVE = 1
    MORPH = 2
    FLAME = 3
    SKY = 5


class SkyType(IntEnum):
    """Enumeration of sky types for SKY firmware effect."""

    SUNRISE = 0
    SUNSET = 1
    CLOUDS = 2


class LIFXUpdateCoordinator(DataUpdateCoordinator[None]):  # noqa: PLR0904
    """DataUpdateCoordinator to gather data for a specific lifx device."""

    def __init__(
        self,
        hass: HomeAssistant,
        connection: LIFXConnection,
        title: str,
    ) -> None:
        """Initialize DataUpdateCoordinator."""
        assert connection.device is not None
        self.connection = connection
        self.device: Light = connection.device
        self.lock = asyncio.Lock()
        self.active_effect = FirmwareEffect.OFF
        self._update_rssi: bool = False
        self._rssi: int = 0
        self.last_used_theme: str = ""

        super().__init__(
            hass,
            _LOGGER,
            name=f"{title} ({self.device.ip_addr})",
            update_interval=timedelta(seconds=LIGHT_UPDATE_INTERVAL),
            # We don't want an immediate refresh since the device
            # takes a moment to reflect the state change
            request_refresh_debouncer=Debouncer(
                hass, _LOGGER, cooldown=REQUEST_REFRESH_DELAY, immediate=False
            ),
        )

    @callback
    def async_setup(self) -> None:
        """Change timeouts."""
        self.device.timeout = MESSAGE_TIMEOUT
        self.device.retry_count = MESSAGE_RETRIES
        self.device.unregister_timeout = UNAVAILABLE_GRACE

    @property
    def rssi(self) -> int:
        """Return stored RSSI value."""
        return self._rssi

    @property
    def rssi_uom(self) -> str:
        """Return the RSSI unit of measurement."""
        if AwesomeVersion(self.device.host_firmware_version) <= RSSI_DBM_FW:
            return SIGNAL_STRENGTH_DECIBELS

        return SIGNAL_STRENGTH_DECIBELS_MILLIWATT

    @property
    def current_infrared_brightness(self) -> str | None:
        """Return the current infrared brightness as a string."""
        return infrared_brightness_value_to_option(self.device.infrared_brightness)

    @cached_property
    def serial_number(self) -> str:
        """Return the internal mac address."""
        return cast(
            str, self.device.mac_addr
        )  # device.mac_addr is not the mac_address, its the serial number

    @cached_property
    def mac_address(self) -> str:
        """Return the physical mac address."""
        return get_real_mac_addr(
            # device.mac_addr is not the mac_address, its the serial number
            self.device.mac_addr,
            self.device.host_firmware_version,
        )

    @property
    def label(self) -> str:
        """Return the label of the bulb."""
        return cast(str, self.device.label)

    @cached_property
    def is_extended_multizone(self) -> bool:
        """Return true if this is a multizone device."""
        return bool(lifx_features(self.device)["extended_multizone"])

    @cached_property
    def is_legacy_multizone(self) -> bool:
        """Return true if this is a legacy multizone device."""
        return bool(
            lifx_features(self.device)["multizone"] and not self.is_extended_multizone
        )

    @cached_property
    def is_matrix(self) -> bool:
        """Return true if this is a matrix device."""
        return bool(lifx_features(self.device)["matrix"])

    async def diagnostics(self) -> dict[str, Any]:
        """Return diagnostic information about the device."""
        features = lifx_features(self.device)
        device_data = {
            "firmware": self.device.host_firmware_version,
            "vendor": self.device.vendor,
            "product_id": self.device.product,
            "features": features,
            "hue": self.device.color[0],
            "saturation": self.device.color[1],
            "brightness": self.device.color[2],
            "kelvin": self.device.color[3],
            "power": self.device.power_level,
        }

        if features["multizone"] is True:
            zones = {"count": self.device.zones_count, "state": {}}
            for index, zone_color in enumerate(self.device.color_zones):
                zones["state"][index] = {
                    "hue": zone_color[0],
                    "saturation": zone_color[1],
                    "brightness": zone_color[2],
                    "kelvin": zone_color[3],
                }
            device_data["zones"] = zones

        if features["hev"] is True:
            device_data["hev"] = {
                "hev_cycle": self.device.hev_cycle,
                "hev_config": self.device.hev_cycle_configuration,
                "last_result": self.device.last_hev_cycle_result,
            }

        if features["infrared"] is True:
            device_data["infrared"] = {"brightness": self.device.infrared_brightness}

        return device_data

    def async_get_entity_id(self, platform: Platform, key: str) -> str | None:
        """Return the entity_id from the platform and key provided."""
        ent_reg = er.async_get(self.hass)
        return ent_reg.async_get_entity_id(
            platform, DOMAIN, f"{self.serial_number}_{key}"
        )

    async def _async_populate_device_info(self) -> None:
        """Populate device info."""
        methods: list[Callable] = []
        device = self.device
        if self.device.host_firmware_version is None:
            methods.append(device.get_hostfirmware)
        if self.device.product is None:
            methods.append(device.get_version)
        if self.device.group is None:
            methods.append(device.get_group)
        assert methods, "Device info already populated"
        await async_multi_execute_lifx_with_retries(
            methods, DEFAULT_ATTEMPTS, OVERALL_TIMEOUT
        )

    def get_number_of_zones(self) -> int:
        """Return the number of zones.

        If the number of zones is not yet populated, return 1 since
        the device will have a least one zone.
        """
        return len(self.device.color_zones) if self.device.color_zones else 1

    @callback
    def _async_build_color_zones_update_requests(self) -> list[Callable]:
        """Build a color zones update request."""
        device = self.device
        calls: list[Callable] = []
        for zone in range(
            0, self.get_number_of_zones(), ZONES_PER_COLOR_UPDATE_REQUEST
        ):

            def _wrap_get_color_zones(
                callb: Callable[[Message, dict[str, Any] | None], None],
                get_color_zones_args: dict[str, Any],
            ) -> None:
                """Capture the callback and make sure resp_set_multizonemultizone is called before."""

                def _wrapped_callback(
                    bulb: Light,
                    response: Message,
                    **kwargs: Any,
                ) -> None:
                    # We need to call resp_set_multizonemultizone to populate
                    # the color_zones attribute before calling the callback
                    device.resp_set_multizonemultizone(response)
                    # Now call the original callback
                    callb(bulb, response, **kwargs)

                device.get_color_zones(**get_color_zones_args, callb=_wrapped_callback)

            calls.append(
                partial(
                    _wrap_get_color_zones,
                    get_color_zones_args={
                        "start_index": zone,
                        "end_index": zone + ZONES_PER_COLOR_UPDATE_REQUEST - 1,
                    },
                )
            )

        return calls

    async def _async_update_data(self) -> None:
        """Fetch all device data from the api."""
        device = self.device
        if (
            device.host_firmware_version is None
            or device.product is None
            or device.group is None
        ):
            await self._async_populate_device_info()

        num_zones = self.get_number_of_zones()
        features = lifx_features(self.device)
        update_rssi = self._update_rssi
        methods: list[Callable] = [self.device.get_color]
        if update_rssi:
            methods.append(self.device.get_wifiinfo)
        if self.is_matrix:
            methods.extend(
                [
                    self.device.get_tile_effect,
                    self.device.get_device_chain,
                    self.device.get64,
                ]
            )
        if self.is_extended_multizone:
            methods.append(self.device.get_extended_color_zones)
        elif self.is_legacy_multizone:
            methods.extend(self._async_build_color_zones_update_requests())
        if self.is_extended_multizone or self.is_legacy_multizone:
            methods.append(self.device.get_multizone_effect)
        if features["hev"]:
            methods.append(self.device.get_hev_cycle)
        if features["infrared"]:
            methods.append(self.device.get_infrared)

        responses = await async_multi_execute_lifx_with_retries(
            methods, MAX_ATTEMPTS_PER_UPDATE_REQUEST_MESSAGE, MAX_UPDATE_TIME
        )
        # device.mac_addr is not the mac_address, its the serial number
        if device.mac_addr == TARGET_ANY:
            device.mac_addr = responses[0].target_addr

        if update_rssi:
            # We always send the rssi request second
            self._rssi = int(floor(10 * log10(responses[1].signal) + 0.5))

        if self.is_matrix or self.is_extended_multizone or self.is_legacy_multizone:
            self.active_effect = FirmwareEffect[self.device.effect.get("effect", "OFF")]
        if self.is_legacy_multizone and num_zones != self.get_number_of_zones():
            # The number of zones has changed so we need
            # to update the zones again. This happens rarely.
            await self.async_get_color_zones()

    async def async_get_color_zones(self) -> None:
        """Get updated color information for each zone."""
        await async_multi_execute_lifx_with_retries(
            self._async_build_color_zones_update_requests(),
            DEFAULT_ATTEMPTS,
            OVERALL_TIMEOUT,
        )

    async def async_get_extended_color_zones(self) -> None:
        """Get updated color information for all zones."""
        try:
            await async_execute_lifx(self.device.get_extended_color_zones)
        except TimeoutError as ex:
            raise HomeAssistantError(
                f"Timeout getting color zones from {self.name}"
            ) from ex

    async def async_set_waveform_optional(
        self, value: dict[str, Any], rapid: bool = False
    ) -> None:
        """Send a set_waveform_optional message to the device."""
        await async_execute_lifx(
            partial(self.device.set_waveform_optional, value=value, rapid=rapid)
        )

    async def async_get_color(self) -> None:
        """Send a get color message to the device."""
        await async_execute_lifx(self.device.get_color)

    async def async_set_power(self, state: bool, duration: int | None) -> None:
        """Send a set power message to the device."""
        await async_execute_lifx(
            partial(self.device.set_power, state, duration=duration)
        )

    async def async_set_color(
        self, hsbk: list[float | int | None], duration: int | None
    ) -> None:
        """Send a set color message to the device."""
        await async_execute_lifx(
            partial(self.device.set_color, hsbk, duration=duration)
        )

    async def async_set_color_zones(
        self,
        start_index: int,
        end_index: int,
        hsbk: list[float | int | None],
        duration: int | None,
        apply: int,
    ) -> None:
        """Send a set color zones message to the device."""
        await async_execute_lifx(
            partial(
                self.device.set_color_zones,
                start_index=start_index,
                end_index=end_index,
                color=hsbk,
                duration=duration,
                apply=apply,
            )
        )

    async def async_set_extended_color_zones(
        self,
        colors: list[tuple[int | float, int | float, int | float, int | float]],
        colors_count: int | None = None,
        duration: int = 0,
        apply: int = 1,
    ) -> None:
        """Send a single set extended color zones message to the device."""

        if colors_count is None:
            colors_count = len(colors)

        # pad the color list with blanks if necessary
        if len(colors) < 82:
            colors.extend([(0, 0, 0, 0) for _ in range(82 - len(colors))])

        await async_execute_lifx(
            partial(
                self.device.set_extended_color_zones,
                colors=colors,
                colors_count=colors_count,
                duration=duration,
                apply=apply,
            )
        )

    async def async_set_multizone_effect(
        self,
        effect: str,
        speed: float = 3.0,
        direction: str = "RIGHT",
        theme_name: str | None = None,
        power_on: bool = True,
    ) -> None:
        """Control the firmware-based Move effect on a multizone device."""
        if self.is_extended_multizone or self.is_legacy_multizone:
            if power_on and self.device.power_level == 0:
                await self.async_set_power(True, 0)

            if theme_name is not None:
                theme = ThemeLibrary().get_theme(theme_name)
                await ThemePainter(self.hass.loop).paint(
                    theme, [self.device], round(speed)
                )

            await async_execute_lifx(
                partial(
                    self.device.set_multizone_effect,
                    effect=MultiZoneEffectType[effect.upper()].value,
                    speed=speed,
                    direction=MultiZoneDirection[direction.upper()].value,
                )
            )
            self.active_effect = FirmwareEffect[effect.upper()]

    async def async_set_matrix_effect(  # noqa: PLR0917
        self,
        effect: str,
        palette: list[tuple[int, int, int, int]] | None = None,
        speed: float | None = None,
        power_on: bool = True,
        sky_type: str | None = None,
        cloud_saturation_min: int | None = None,
        cloud_saturation_max: int | None = None,
    ) -> None:
        """Control the firmware-based effects on a matrix device."""
        if self.is_matrix:
            if power_on and self.device.power_level == 0:
                await self.async_set_power(True, 0)

            if palette is None:
                palette = []

            if sky_type is not None:
                sky_type = TileEffectSkyType[sky_type.upper()].value

            await async_execute_lifx(
                partial(
                    self.device.set_tile_effect,
                    effect=TileEffectType[effect.upper()].value,
                    speed=speed,
                    palette=palette,
                    sky_type=sky_type,
                    cloud_saturation_min=cloud_saturation_min,
                    cloud_saturation_max=cloud_saturation_max,
                )
            )
            self.active_effect = FirmwareEffect[effect.upper()]

    def async_get_active_effect(self) -> int:
        """Return the enum value of the currently active firmware effect."""
        return self.active_effect.value

    async def async_set_infrared_brightness(self, option: str) -> None:
        """Set infrared brightness."""
        infrared_brightness = infrared_brightness_option_to_value(option)
        await async_execute_lifx(partial(self.device.set_infrared, infrared_brightness))

    async def async_identify_bulb(self) -> None:
        """Identify the device by flashing it three times."""
        bulb: Light = self.device
        if bulb.power_level:
            # just flash the bulb for three seconds
            await self.async_set_waveform_optional(value=IDENTIFY_WAVEFORM)
            return
        # Turn the bulb on first, flash for 3 seconds, then turn off
        await self.async_set_power(state=True, duration=1)
        await self.async_set_waveform_optional(value=IDENTIFY_WAVEFORM)
        await asyncio.sleep(LIFX_IDENTIFY_DELAY)
        await self.async_set_power(state=False, duration=1)

    def async_enable_rssi_updates(self) -> Callable[[], None]:
        """Enable RSSI signal strength updates."""

        @callback
        def _async_disable_rssi_updates() -> None:
            """Disable RSSI updates when sensor removed."""
            self._update_rssi = False

        self._update_rssi = True
        return _async_disable_rssi_updates

    def async_get_hev_cycle_state(self) -> bool | None:
        """Return the current HEV cycle state."""
        if self.device.hev_cycle is None:
            return None
        return bool(self.device.hev_cycle.get(ATTR_REMAINING, 0) > 0)

    async def async_set_hev_cycle_state(self, enable: bool, duration: int = 0) -> None:
        """Start or stop an HEV cycle on a LIFX Clean bulb."""
        if lifx_features(self.device)["hev"]:
            await async_execute_lifx(
                partial(self.device.set_hev_cycle, enable=enable, duration=duration)
            )

    async def async_apply_theme(self, theme_name: str) -> None:
        """Apply the selected theme to the device."""
        self.last_used_theme = theme_name
        theme = ThemeLibrary().get_theme(theme_name)
        await ThemePainter(self.hass.loop).paint(theme, [self.device])