core/homeassistant/components/unifiprotect/select.py

"""Component providing select entities for UniFi Protect."""

from collections.abc import Callable, Sequence
from dataclasses import dataclass
from enum import Enum
import logging
from typing import Any

from uiprotect.api import ProtectApiClient
from uiprotect.data import (
    NVR,
    Camera,
    ChimeType,
    DoorbellMessageType,
    Doorlock,
    IRLEDMode,
    Light,
    LightModeEnableType,
    LightModeType,
    ModelType,
    MountType,
    ProtectAdoptableDeviceModel,
    PTZPatrol,
    PublicHdrMode,
    RecordingMode,
    Sensor,
    Viewer,
)
from uiprotect.exceptions import GlobalAlarmManagerError

from homeassistant.components.select import SelectEntity, SelectEntityDescription
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity import EntityDescription
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback

from .const import DOMAIN, TYPE_EMPTY_VALUE
from .data import ProtectData, ProtectDeviceType, UFPConfigEntry
from .entity import (
    PermRequired,
    ProtectDeviceEntity,
    ProtectEntityDescription,
    ProtectNVREntity,
    ProtectSettableKeysMixin,
    T,
    async_all_device_entities,
)
from .utils import async_get_light_motion_current, async_ufp_instance_command

_LOGGER = logging.getLogger(__name__)
_KEY_LIGHT_MOTION = "light_motion"
PARALLEL_UPDATES = 0

HDR_MODES = [
    {"id": "always", "name": "always"},
    {"id": "off", "name": "off"},
    {"id": "auto", "name": "auto"},
]

INFRARED_MODES = [
    {"id": IRLEDMode.AUTO.value, "name": "auto"},
    {"id": IRLEDMode.ON.value, "name": "on"},
    {"id": IRLEDMode.AUTO_NO_LED.value, "name": "auto_filter_only"},
    {"id": IRLEDMode.CUSTOM.value, "name": "custom"},
    {"id": IRLEDMode.OFF.value, "name": "off"},
]

CHIME_TYPES = [
    {"id": ChimeType.NONE.value, "name": "none"},
    {"id": ChimeType.MECHANICAL.value, "name": "mechanical"},
    {"id": ChimeType.DIGITAL.value, "name": "digital"},
]

MOUNT_TYPES = [
    {"id": MountType.NONE.value, "name": MountType.NONE.value},
    {"id": MountType.DOOR.value, "name": MountType.DOOR.value},
    {"id": MountType.WINDOW.value, "name": MountType.WINDOW.value},
    {"id": MountType.GARAGE.value, "name": MountType.GARAGE.value},
    {"id": MountType.LEAK.value, "name": MountType.LEAK.value},
]

LIGHT_MODE_MOTION = "motion"
LIGHT_MODE_MOTION_DARK = "motion_dark"
LIGHT_MODE_DARK = "when_dark"
LIGHT_MODE_OFF = "manual"
LIGHT_MODES = [LIGHT_MODE_MOTION, LIGHT_MODE_DARK, LIGHT_MODE_OFF]

LIGHT_MODE_TO_SETTINGS = {
    LIGHT_MODE_MOTION: (LightModeType.MOTION.value, LightModeEnableType.ALWAYS.value),
    LIGHT_MODE_MOTION_DARK: (
        LightModeType.MOTION.value,
        LightModeEnableType.DARK.value,
    ),
    LIGHT_MODE_DARK: (LightModeType.WHEN_DARK.value, LightModeEnableType.DARK.value),
    LIGHT_MODE_OFF: (LightModeType.MANUAL.value, None),
}

MOTION_MODE_TO_LIGHT_MODE = [
    {"id": LightModeType.MOTION.value, "name": LIGHT_MODE_MOTION},
    {"id": f"{LightModeType.MOTION.value}_dark", "name": LIGHT_MODE_MOTION_DARK},
    {"id": LightModeType.WHEN_DARK.value, "name": LIGHT_MODE_DARK},
    {"id": LightModeType.MANUAL.value, "name": LIGHT_MODE_OFF},
]

PTZ_PATROL_STOP = "stop"
_KEY_PTZ_PATROL = "ptz_patrol"

DEVICE_RECORDING_MODES = [
    {"id": mode.value, "name": mode.value} for mode in list(RecordingMode)
]


@dataclass(frozen=True, kw_only=True)
class ProtectSelectEntityDescription(
    ProtectSettableKeysMixin[T], SelectEntityDescription
):
    """Describes UniFi Protect Select entity."""

    ufp_options: list[dict[str, Any]] | None = None
    ufp_options_fn: Callable[[ProtectApiClient], list[dict[str, Any]]] | None = None
    ufp_enum_type: type[Enum] | None = None


def _get_viewer_options(api: ProtectApiClient) -> list[dict[str, Any]]:
    return [
        {"id": item.id, "name": item.name} for item in api.bootstrap.liveviews.values()
    ]


def _get_doorbell_options(api: ProtectApiClient) -> list[dict[str, Any]]:
    default_message = api.bootstrap.nvr.doorbell_settings.default_message_text
    messages = api.bootstrap.nvr.doorbell_settings.all_messages
    built_messages: list[dict[str, str]] = []

    for item in messages:
        msg_type = item.type.value
        if item.type is DoorbellMessageType.CUSTOM_MESSAGE:
            msg_type = f"{DoorbellMessageType.CUSTOM_MESSAGE.value}:{item.text}"

        built_messages.append({"id": msg_type, "name": item.text})

    return [
        {"id": "", "name": f"Default Message ({default_message})"},
        *built_messages,
    ]


def _get_paired_camera_options(api: ProtectApiClient) -> list[dict[str, Any]]:
    options = [{"id": TYPE_EMPTY_VALUE, "name": "Not Paired"}]
    options.extend(
        {"id": camera.id, "name": camera.display_name or camera.type}
        for camera in api.bootstrap.cameras.values()
    )

    return options


def _get_viewer_current(obj: Viewer) -> str:
    return obj.liveview_id


def _get_doorbell_current(obj: Camera) -> str | None:
    if obj.lcd_message is None:
        return None
    return obj.lcd_message.text


async def _set_light_mode(obj: Light, mode: str) -> None:
    lightmode, timing = LIGHT_MODE_TO_SETTINGS[mode]
    await obj.set_light_settings(
        LightModeType(lightmode),
        enable_at=None if timing is None else LightModeEnableType(timing),
    )


async def _set_paired_camera(obj: Light | Sensor | Doorlock, camera_id: str) -> None:
    if camera_id == TYPE_EMPTY_VALUE:
        camera: Camera | None = None
    else:
        camera = obj.api.bootstrap.cameras.get(camera_id)
    await obj.set_paired_camera(camera)


async def _set_doorbell_message(obj: Camera, message: str) -> None:
    if message.startswith(DoorbellMessageType.CUSTOM_MESSAGE.value):
        message = message.rsplit(":", maxsplit=1)[-1]
        await obj.set_lcd_message_public(
            DoorbellMessageType.CUSTOM_MESSAGE, text=message
        )
    elif message == TYPE_EMPTY_VALUE:
        # Public API has no endpoint to clear the LCD message; fall back to
        # the non-deprecated legacy helper.
        await obj.set_lcd_text(None)
    else:
        await obj.set_lcd_message_public(DoorbellMessageType(message))


async def _set_liveview(obj: Viewer, liveview_id: str) -> None:
    """Set the liveview for a viewer."""
    liveview = obj.api.bootstrap.liveviews[liveview_id]
    await obj.set_liveview(liveview)


async def _set_ptz_patrol(obj: Camera, patrol_slot: str) -> None:
    """Start or stop PTZ patrol."""
    if patrol_slot == PTZ_PATROL_STOP:
        await obj.ptz_patrol_stop_public()
    else:
        slot = int(patrol_slot)
        await obj.ptz_patrol_start_public(slot=slot)


_HDR_MODE_MAP = {
    "auto": PublicHdrMode.AUTO,
    "always": PublicHdrMode.ON,
    "off": PublicHdrMode.OFF,
}


async def _set_hdr_mode(obj: Camera, mode: str) -> None:
    """Set HDR mode via the public API."""
    await obj.set_hdr_mode_public(_HDR_MODE_MAP[mode])


PTZ_PATROL_DESCRIPTION = ProtectSelectEntityDescription[Camera](
    key=_KEY_PTZ_PATROL,
    translation_key="ptz_patrol",
    entity_category=EntityCategory.CONFIG,
    ufp_required_field="feature_flags.is_ptz",
    ufp_set_method_fn=_set_ptz_patrol,
    ufp_perm=PermRequired.WRITE,
)

CAMERA_SELECTS: tuple[ProtectSelectEntityDescription, ...] = (
    ProtectSelectEntityDescription(
        key="recording_mode",
        translation_key="recording_mode",
        entity_category=EntityCategory.CONFIG,
        ufp_options=DEVICE_RECORDING_MODES,
        ufp_enum_type=RecordingMode,
        ufp_value="recording_settings.mode",
        ufp_set_method="set_recording_mode",
        ufp_perm=PermRequired.WRITE,
    ),
    ProtectSelectEntityDescription(
        key="infrared",
        translation_key="infrared_mode",
        entity_category=EntityCategory.CONFIG,
        ufp_required_field="feature_flags.has_led_ir",
        ufp_options=INFRARED_MODES,
        ufp_enum_type=IRLEDMode,
        ufp_value="isp_settings.ir_led_mode",
        ufp_set_method="set_ir_led_model",
        ufp_perm=PermRequired.WRITE,
    ),
    ProtectSelectEntityDescription[Camera](
        key="doorbell_text",
        translation_key="doorbell_text",
        entity_category=EntityCategory.CONFIG,
        ufp_required_field="feature_flags.has_lcd_screen",
        ufp_value_fn=_get_doorbell_current,
        ufp_options_fn=_get_doorbell_options,
        ufp_set_method_fn=_set_doorbell_message,
        ufp_perm=PermRequired.WRITE,
    ),
    ProtectSelectEntityDescription(
        key="chime_type",
        translation_key="chime_type",
        entity_category=EntityCategory.CONFIG,
        ufp_required_field="feature_flags.has_chime",
        ufp_options=CHIME_TYPES,
        ufp_enum_type=ChimeType,
        ufp_value="chime_type",
        ufp_set_method="set_chime_type",
        ufp_perm=PermRequired.WRITE,
    ),
    ProtectSelectEntityDescription[Camera](
        key="hdr_mode",
        translation_key="hdr_mode",
        entity_category=EntityCategory.CONFIG,
        ufp_required_field="feature_flags.has_hdr",
        ufp_options=HDR_MODES,
        ufp_value="hdr_mode_display",
        ufp_set_method_fn=_set_hdr_mode,
        ufp_perm=PermRequired.WRITE,
    ),
)

LIGHT_SELECTS: tuple[ProtectSelectEntityDescription, ...] = (
    ProtectSelectEntityDescription[Light](
        key=_KEY_LIGHT_MOTION,
        translation_key="light_mode",
        entity_category=EntityCategory.CONFIG,
        ufp_options=MOTION_MODE_TO_LIGHT_MODE,
        ufp_value_fn=async_get_light_motion_current,
        ufp_set_method_fn=_set_light_mode,
        ufp_perm=PermRequired.WRITE,
    ),
    ProtectSelectEntityDescription[Light](
        key="paired_camera",
        translation_key="paired_camera",
        entity_category=EntityCategory.CONFIG,
        ufp_value="camera_id",
        ufp_options_fn=_get_paired_camera_options,
        ufp_set_method_fn=_set_paired_camera,
        ufp_perm=PermRequired.WRITE,
    ),
)

SENSE_SELECTS: tuple[ProtectSelectEntityDescription, ...] = (
    ProtectSelectEntityDescription(
        key="mount_type",
        translation_key="mount_type",
        entity_category=EntityCategory.CONFIG,
        ufp_options=MOUNT_TYPES,
        ufp_enum_type=MountType,
        ufp_value="mount_type",
        ufp_set_method="set_mount_type",
        ufp_perm=PermRequired.WRITE,
    ),
    ProtectSelectEntityDescription[Sensor](
        key="paired_camera",
        translation_key="paired_camera",
        entity_category=EntityCategory.CONFIG,
        ufp_value="camera_id",
        ufp_options_fn=_get_paired_camera_options,
        ufp_set_method_fn=_set_paired_camera,
        ufp_perm=PermRequired.WRITE,
    ),
)

DOORLOCK_SELECTS: tuple[ProtectSelectEntityDescription, ...] = (
    ProtectSelectEntityDescription[Doorlock](
        key="paired_camera",
        translation_key="paired_camera",
        entity_category=EntityCategory.CONFIG,
        ufp_value="camera_id",
        ufp_options_fn=_get_paired_camera_options,
        ufp_set_method_fn=_set_paired_camera,
        ufp_perm=PermRequired.WRITE,
    ),
)

VIEWER_SELECTS: tuple[ProtectSelectEntityDescription, ...] = (
    ProtectSelectEntityDescription[Viewer](
        key="viewer",
        translation_key="liveview",
        ufp_options_fn=_get_viewer_options,
        ufp_value_fn=_get_viewer_current,
        ufp_set_method_fn=_set_liveview,
        ufp_perm=PermRequired.WRITE,
    ),
)

_MODEL_DESCRIPTIONS: dict[ModelType, Sequence[ProtectEntityDescription]] = {
    ModelType.CAMERA: CAMERA_SELECTS,
    ModelType.LIGHT: LIGHT_SELECTS,
    ModelType.SENSOR: SENSE_SELECTS,
    ModelType.VIEWPORT: VIEWER_SELECTS,
    ModelType.DOORLOCK: DOORLOCK_SELECTS,
}


async def async_setup_entry(
    hass: HomeAssistant,
    entry: UFPConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up number entities for UniFi Protect integration."""
    data = entry.runtime_data

    @callback
    def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:
        entities = list(
            async_all_device_entities(
                data,
                ProtectSelects,
                model_descriptions=_MODEL_DESCRIPTIONS,
                ufp_device=device,
            )
        )
        if isinstance(device, Camera) and device.feature_flags.is_ptz:
            patrols = data.ptz_patrols.get(device.id, [])
            entities.append(ProtectPTZPatrolSelect(data, device, patrols))
        async_add_entities(entities)

    data.async_subscribe_adopt(_add_new_device)

    entities = list(
        async_all_device_entities(
            data, ProtectSelects, model_descriptions=_MODEL_DESCRIPTIONS
        )
    )

    for camera in data.api.bootstrap.cameras.values():
        if camera.feature_flags.is_ptz and camera.is_adopted_by_us:
            patrols = data.ptz_patrols.get(camera.id, [])
            entities.append(ProtectPTZPatrolSelect(data, camera, patrols))

    api = data.api
    if (
        api.has_public_bootstrap
        and api.public_bootstrap.arm_mode is not None
        and api.public_bootstrap.arm_profiles
    ):
        entities.append(ProtectNVRArmProfileSelect(data, device=api.bootstrap.nvr))

    async_add_entities(entities)


class ProtectSelects(ProtectDeviceEntity, SelectEntity):
    """A UniFi Protect Select Entity."""

    device: Camera | Light | Viewer
    entity_description: ProtectSelectEntityDescription
    _state_attrs = ("_attr_available", "_attr_options", "_attr_current_option")

    def __init__(
        self,
        data: ProtectData,
        device: Camera | Light | Viewer,
        description: ProtectSelectEntityDescription,
    ) -> None:
        """Initialize the unifi protect select entity."""
        self._async_set_options(data, description)
        super().__init__(data, device, description)

    @callback
    def _async_update_device_from_protect(self, device: ProtectDeviceType) -> None:
        super()._async_update_device_from_protect(device)
        entity_description = self.entity_description
        # entities with categories are not exposed for voice
        # and safe to update dynamically
        if (
            entity_description.entity_category is not None
            and entity_description.ufp_options_fn is not None
        ):
            _LOGGER.debug("Updating dynamic select options for %s", self.entity_id)
            self._async_set_options(self.data, entity_description)
        if (unifi_value := entity_description.get_ufp_value(device)) is None:
            unifi_value = TYPE_EMPTY_VALUE
        self._attr_current_option = self._unifi_to_hass_options.get(
            unifi_value, unifi_value
        )

    @callback
    def _async_set_options(
        self, data: ProtectData, description: ProtectSelectEntityDescription
    ) -> None:
        """Set options attributes from UniFi Protect device."""
        if (ufp_options := description.ufp_options) is not None:
            options = ufp_options
        else:
            assert description.ufp_options_fn is not None
            options = description.ufp_options_fn(data.api)

        self._attr_options = [item["name"] for item in options]
        self._hass_to_unifi_options = {item["name"]: item["id"] for item in options}
        self._unifi_to_hass_options = {item["id"]: item["name"] for item in options}

    @async_ufp_instance_command
    async def async_select_option(self, option: str) -> None:
        """Change the Select Entity Option."""

        # Light Motion is a bit different
        if self.entity_description.key == _KEY_LIGHT_MOTION:
            assert self.entity_description.ufp_set_method_fn is not None
            await self.entity_description.ufp_set_method_fn(self.device, option)
            return

        unifi_value = self._hass_to_unifi_options[option]
        if self.entity_description.ufp_enum_type is not None:
            unifi_value = self.entity_description.ufp_enum_type(unifi_value)
        await self.entity_description.ufp_set(self.device, unifi_value)


class ProtectPTZPatrolSelect(ProtectDeviceEntity, SelectEntity):
    """A UniFi Protect PTZ Patrol Select Entity."""

    device: Camera
    _attr_current_option: str | None = None
    _state_attrs = ("_attr_available", "_attr_options", "_attr_current_option")

    def __init__(
        self,
        data: ProtectData,
        device: Camera,
        patrols: list[PTZPatrol],
    ) -> None:
        """Initialize the PTZ patrol select entity."""
        # Build options from cached patrols
        self._hass_to_unifi_options: dict[str, str] = {PTZ_PATROL_STOP: PTZ_PATROL_STOP}
        self._hass_to_unifi_options.update(
            {patrol.name: str(patrol.slot) for patrol in patrols}
        )
        self._unifi_to_hass_options = {
            v: k for k, v in self._hass_to_unifi_options.items()
        }
        self._attr_options = list(self._hass_to_unifi_options)

        super().__init__(data, device, PTZ_PATROL_DESCRIPTION)
        # Set initial state based on active patrol
        self._update_patrol_state()

    def _update_patrol_state(self) -> None:
        """Update the patrol state based on active_patrol_slot."""
        if self.device.active_patrol_slot is not None:
            # A patrol is running - show which one
            slot_str = str(self.device.active_patrol_slot)
            self._attr_current_option = self._unifi_to_hass_options.get(slot_str)
        else:
            # No patrol running - show Stop
            self._attr_current_option = PTZ_PATROL_STOP

    @callback
    def _async_update_device_from_protect(self, device: ProtectDeviceType) -> None:
        super()._async_update_device_from_protect(device)
        # Update patrol state from websocket updates
        self._update_patrol_state()

    @async_ufp_instance_command
    async def async_select_option(self, option: str) -> None:
        """Start or stop a PTZ patrol."""
        # Home Assistant validates options before calling this method,
        # so we can safely assume the option is valid
        unifi_value = self._hass_to_unifi_options[option]
        await _set_ptz_patrol(self.device, unifi_value)
        # State will be updated via websocket when active_patrol_slot changes


class ProtectNVRArmProfileSelect(ProtectNVREntity, SelectEntity):
    """UniFi Protect NVR arm profile select entity."""

    _attr_translation_key = "nvr_arm_profile"
    _attr_current_option: str | None = None
    _state_attrs = ("_attr_available", "_attr_options", "_attr_current_option")

    def __init__(self, data: ProtectData, device: NVR) -> None:
        """Initialize the NVR arm profile select entity."""
        self._id_to_name: dict[str, str] = {}
        self._name_to_id: dict[str, str] = {}
        super().__init__(data, device, EntityDescription(key="nvr_arm_profile"))
        self._refresh_arm_profile_state()

    @callback
    def _refresh_arm_profile_state(self) -> None:
        """Update options and current option from the public bootstrap cache."""
        api = self.data.api
        pb = api.public_bootstrap if api.has_public_bootstrap else None
        arm_mode = pb.arm_mode if pb is not None else None

        if pb is None or arm_mode is None:
            self._attr_available = False
            self._attr_options = []
            self._attr_current_option = None
            return

        # Always append a short id suffix so every option label is unique
        # and stable even if another profile with the same name is added later.
        self._id_to_name = {}
        self._name_to_id = {}
        for pid, profile in pb.arm_profiles.items():
            label = f"{profile.name} ({pid[-6:]})"
            self._id_to_name[pid] = label
            self._name_to_id[label] = pid
        self._attr_options = sorted(self._name_to_id)
        profile_id = arm_mode.arm_profile_id
        self._attr_current_option = (
            self._id_to_name.get(profile_id) if profile_id else None
        )

    @callback
    def _async_update_device_from_protect(self, device: ProtectDeviceType) -> None:
        super()._async_update_device_from_protect(device)
        self._refresh_arm_profile_state()

    @async_ufp_instance_command
    async def async_select_option(self, option: str) -> None:
        """Change the currently active arm profile."""
        profile_id = self._name_to_id[option]
        try:
            await self.data.api.set_current_arm_profile_public(profile_id)
        except GlobalAlarmManagerError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="global_alarm_manager",
            ) from err