core/homeassistant/components/wled/select.py

"""Support for LED selects."""

from functools import partial
from typing import override

from wled import LiveDataOverride

from homeassistant.components.select import SelectEntity
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback

from .coordinator import WLEDConfigEntry, WLEDDataUpdateCoordinator
from .entity import WLEDEntity
from .helpers import wled_exception_handler

PARALLEL_UPDATES = 1


async def async_setup_entry(
    hass: HomeAssistant,
    entry: WLEDConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up WLED select based on a config entry."""
    coordinator = entry.runtime_data

    async_add_entities(
        [
            WLEDLiveOverrideSelect(coordinator),
            WLEDPlaylistSelect(coordinator),
            WLEDPresetSelect(coordinator),
        ]
    )

    update_segments = partial(
        async_update_segments,
        coordinator,
        set(),
        async_add_entities,
    )
    coordinator.async_add_listener(update_segments)
    update_segments()


class WLEDLiveOverrideSelect(WLEDEntity, SelectEntity):
    """Defined a WLED Live Override select."""

    _attr_entity_category = EntityCategory.CONFIG
    _attr_translation_key = "live_override"

    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
        """Initialize WLED ."""
        super().__init__(coordinator=coordinator)

        self._attr_unique_id = f"{coordinator.data.info.mac_address}_live_override"
        self._attr_options = [str(live.value) for live in LiveDataOverride]

    @property
    @override
    def current_option(self) -> str:
        """Return the current selected live override."""
        return str(self.coordinator.data.state.live_data_override.value)

    @wled_exception_handler
    @override
    async def async_select_option(self, option: str) -> None:
        """Set WLED state to the selected live override state."""
        await self.coordinator.wled.live(live=LiveDataOverride(int(option)))


class WLEDPresetSelect(WLEDEntity, SelectEntity):
    """Defined a WLED Preset select."""

    _attr_translation_key = "preset"

    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
        """Initialize WLED ."""
        super().__init__(coordinator=coordinator)

        self._attr_unique_id = f"{coordinator.data.info.mac_address}_preset"

    @property
    @override
    def available(self) -> bool:
        """Return True if entity is available."""
        return len(self.coordinator.data.presets) > 0 and super().available

    @property
    @override
    def current_option(self) -> str | None:
        """Return the current selected preset."""
        if not self.coordinator.data.state.preset_id:
            return None
        if preset := self.coordinator.data.presets.get(
            self.coordinator.data.state.preset_id
        ):
            return preset.name
        return None

    @property
    @override
    def options(self) -> list[str]:
        """Return a list of selectable options."""
        sorted_values = sorted(
            self.coordinator.data.presets.values(), key=lambda preset: preset.name
        )
        return [preset.name for preset in sorted_values]

    @wled_exception_handler
    @override
    async def async_select_option(self, option: str) -> None:
        """Set WLED segment to the selected preset."""
        await self.coordinator.wled.preset(preset=option)


class WLEDPlaylistSelect(WLEDEntity, SelectEntity):
    """Define a WLED Playlist select."""

    _attr_translation_key = "playlist"

    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
        """Initialize WLED playlist."""
        super().__init__(coordinator=coordinator)

        self._attr_unique_id = f"{coordinator.data.info.mac_address}_playlist"

    @property
    @override
    def available(self) -> bool:
        """Return True if entity is available."""
        return len(self.coordinator.data.playlists) > 0 and super().available

    @property
    @override
    def current_option(self) -> str | None:
        """Return the currently selected playlist."""
        if not self.coordinator.data.state.playlist_id:
            return None
        if playlist := self.coordinator.data.playlists.get(
            self.coordinator.data.state.playlist_id
        ):
            return playlist.name
        return None

    @property
    @override
    def options(self) -> list[str]:
        """Return a list of selectable options."""
        sorted_values = sorted(
            self.coordinator.data.playlists.values(), key=lambda playlist: playlist.name
        )
        return [playlist.name for playlist in sorted_values]

    @wled_exception_handler
    @override
    async def async_select_option(self, option: str) -> None:
        """Set WLED segment to the selected playlist."""
        await self.coordinator.wled.playlist(playlist=option)


class WLEDPaletteSelect(WLEDEntity, SelectEntity):
    """Defines a WLED Palette select."""

    _attr_entity_category = EntityCategory.CONFIG
    _attr_translation_key = "color_palette"
    _segment: int

    def __init__(self, coordinator: WLEDDataUpdateCoordinator, segment: int) -> None:
        """Initialize WLED ."""
        super().__init__(coordinator=coordinator)

        # Segment 0 uses a simpler name, which is more natural for when using
        # a single segment / using WLED with one big LED strip.
        if segment != 0:
            self._attr_translation_key = "segment_color_palette"
            self._attr_translation_placeholders = {"segment": str(segment)}

        self._attr_unique_id = f"{coordinator.data.info.mac_address}_palette_{segment}"
        self._segment = segment

    @property
    @override
    def available(self) -> bool:
        """Return True if entity is available."""
        return (
            super().available and self._segment in self.coordinator.data.state.segments
        )

    @property
    @override
    def current_option(self) -> str | None:
        """Return the current selected color palette."""
        if not self.coordinator.data.palettes:
            return None
        if (segment := self.coordinator.data.state.segments.get(self._segment)) is None:
            return None
        palette_id = int(segment.palette_id)
        if (palette := self.coordinator.data.palettes.get(palette_id)) is None:
            return None
        return palette.name

    @property
    @override
    def options(self) -> list[str]:
        """Return a list of selectable options."""
        sorted_values = sorted(
            self.coordinator.data.palettes.values(), key=lambda palette: palette.name
        )
        return [palette.name for palette in sorted_values]

    @wled_exception_handler
    @override
    async def async_select_option(self, option: str) -> None:
        """Set WLED segment to the selected color palette."""
        await self.coordinator.wled.segment(segment_id=self._segment, palette=option)


@callback
def async_update_segments(
    coordinator: WLEDDataUpdateCoordinator,
    current_ids: set[int],
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Update segments."""
    new_entities: list[WLEDPaletteSelect] = []

    # Process new segments, add them to Home Assistant
    for segment_id in coordinator.segment_ids - current_ids:
        current_ids.add(segment_id)
        new_entities.append(WLEDPaletteSelect(coordinator, segment_id))

    async_add_entities(new_entities)