__init__.py

"""Create traits for V1 devices.

Traits are modular components that encapsulate specific features of a Roborock
device. This module provides a factory function to create and initialize the
appropriate traits for V1 devices based on their capabilities.

Using Traits
------------
Traits are accessed via the `v1_properties` attribute on a device. Each trait
represents a specific capability, such as `status`, `consumables`, or `rooms`.

Traits serve two main purposes:
1.  **State**: Traits are dataclasses that hold the current state of the device
    feature. You can access attributes directly (e.g., `device.v1_properties.status.battery`).
2.  **Commands**: Traits provide methods to control the device. For example,
    `device.v1_properties.volume.set_volume()`.

Additionally, the `command` trait provides a generic way to send any command to the
device (e.g. `device.v1_properties.command.send("app_start")`). This is often used
for basic cleaning operations like starting, stopping, or docking the vacuum.

Most traits have a `refresh()` method that must be called to update their state
from the device. The state is not updated automatically in real-time unless
specifically implemented by the trait or via polling.

Adding New Traits
-----------------
When adding a new trait, the most common pattern is to subclass `V1TraitMixin`
and a `RoborockBase` dataclass. You must define a `command` class variable that
specifies the `RoborockCommand` used to fetch the trait data from the device.
See `common.py` for more details on common patterns used across traits.

There are some additional decorators in `common.py` that can be used to specify which
RPC channel to use for the trait (standard, MQTT/cloud, or map-specific).

  - `@common.mqtt_rpc_channel` - Use the MQTT RPC channel for this trait.
  - `@common.map_rpc_channel` - Use the map RPC channel for this trait.

There are also some attributes that specify device feature dependencies for
optional traits:

    - `requires_feature` - The string name of the device feature that must be supported
        for this trait to be enabled. See `DeviceFeaturesTrait` for a list of
        available features.
    - `requires_dock_type` - If set, this is a function that accepts a `RoborockDockTypeCode`
        and returns a boolean indicating whether the trait is supported for that dock type.

Additionally, DeviceFeaturesTrait has a method `is_field_supported` that is used to
check individual trait field values. This is a more fine grained version to allow
optional fields in a dataclass, vs the above feature checks that apply to an entire
trait. The `requires_schema_code` field metadata attribute is a string of the schema
code in HomeDataProduct Schema that is required for the field to be supported.
"""

import logging
from collections.abc import Callable
from dataclasses import dataclass, field, fields
from typing import Any, get_args

from roborock.data.containers import HomeData, HomeDataProduct, RoborockBase
from roborock.data.v1.v1_code_mappings import RoborockDockTypeCode
from roborock.devices.cache import DeviceCache
from roborock.devices.traits import Trait
from roborock.exceptions import RoborockException
from roborock.map.map_parser import MapParserConfig
from roborock.protocols.v1_protocol import V1RpcChannel, decode_data_protocol_message
from roborock.roborock_message import RoborockDataProtocol, RoborockMessage
from roborock.web_api import UserWebApiClient

from . import (
    child_lock,
    clean_summary,
    command,
    common,
    consumeable,
    device_features,
    do_not_disturb,
    dust_collection_mode,
    flow_led_status,
    home,
    led_status,
    map_content,
    maps,
    network_info,
    rooms,
    routines,
    smart_wash_params,
    status,
    valley_electricity_timer,
    volume,
    wash_towel_mode,
)
from .child_lock import ChildLockTrait
from .clean_summary import CleanSummaryTrait
from .command import CommandTrait
from .common import V1TraitMixin
from .consumeable import ConsumableTrait
from .device_features import DeviceFeaturesTrait
from .do_not_disturb import DoNotDisturbTrait
from .dust_collection_mode import DustCollectionModeTrait
from .flow_led_status import FlowLedStatusTrait
from .home import HomeTrait
from .led_status import LedStatusTrait
from .map_content import MapContentTrait
from .maps import MapsTrait
from .network_info import NetworkInfoTrait
from .rooms import RoomsTrait
from .routines import RoutinesTrait
from .smart_wash_params import SmartWashParamsTrait
from .status import StatusTrait
from .valley_electricity_timer import ValleyElectricityTimerTrait
from .volume import SoundVolumeTrait
from .wash_towel_mode import WashTowelModeTrait

_LOGGER = logging.getLogger(__name__)

__all__ = [
    "PropertiesApi",
    "child_lock",
    "clean_summary",
    "command",
    "common",
    "consumeable",
    "device_features",
    "do_not_disturb",
    "dust_collection_mode",
    "flow_led_status",
    "home",
    "led_status",
    "map_content",
    "maps",
    "network_info",
    "rooms",
    "routines",
    "smart_wash_params",
    "status",
    "valley_electricity_timer",
    "volume",
    "wash_towel_mode",
]


@dataclass
class PropertiesApi(Trait):
    """Common properties for V1 devices.

    This class holds all the traits that are common across all V1 devices.
    """

    # All v1 devices have these traits
    status: StatusTrait
    command: CommandTrait
    dnd: DoNotDisturbTrait
    clean_summary: CleanSummaryTrait
    sound_volume: SoundVolumeTrait
    rooms: RoomsTrait
    maps: MapsTrait
    map_content: MapContentTrait
    consumables: ConsumableTrait
    home: HomeTrait
    device_features: DeviceFeaturesTrait
    network_info: NetworkInfoTrait
    routines: RoutinesTrait

    # Optional features that may not be supported on all devices
    child_lock: ChildLockTrait | None = None
    led_status: LedStatusTrait | None = None
    flow_led_status: FlowLedStatusTrait | None = None
    valley_electricity_timer: ValleyElectricityTimerTrait | None = None
    dust_collection_mode: DustCollectionModeTrait | None = None
    wash_towel_mode: WashTowelModeTrait | None = None
    smart_wash_params: SmartWashParamsTrait | None = None

    def __init__(
        self,
        device_uid: str,
        product: HomeDataProduct,
        home_data: HomeData,
        rpc_channel: V1RpcChannel,
        mqtt_rpc_channel: V1RpcChannel,
        map_rpc_channel: V1RpcChannel,
        add_dps_listener: Callable[[Callable[[dict[RoborockDataProtocol, Any]], None]], Callable[[], None]],
        web_api: UserWebApiClient,
        device_cache: DeviceCache,
        map_parser_config: MapParserConfig | None = None,
        region: str | None = None,
    ) -> None:
        """Initialize the V1TraitProps."""
        self._device_uid = device_uid
        self._rpc_channel = rpc_channel
        self._mqtt_rpc_channel = mqtt_rpc_channel
        self._map_rpc_channel = map_rpc_channel
        self._web_api = web_api
        self._device_cache = device_cache
        self._region = region
        self._unsub: Callable[[], None] | None = None
        self._add_dps_listener = add_dps_listener

        self.device_features = DeviceFeaturesTrait(product, self._device_cache)
        self.status = StatusTrait(self.device_features, region=self._region)
        self.consumables = ConsumableTrait()
        self.rooms = RoomsTrait(home_data, web_api)
        self.maps = MapsTrait(self.status)
        self.map_content = MapContentTrait(map_parser_config)
        self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, self._device_cache)
        self.network_info = NetworkInfoTrait(device_uid, self._device_cache)
        self.routines = RoutinesTrait(device_uid, web_api)

        # Dynamically create any traits that need to be populated
        for item in fields(self):
            if (trait := getattr(self, item.name, None)) is None:
                # We exclude optional features and them via discover_features
                if (union_args := get_args(item.type)) is None or len(union_args) > 0:
                    continue
                _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
                if not callable(item.type):
                    continue
                trait = item.type()
                setattr(self, item.name, trait)
            # This is a hack to allow setting the rpc_channel on all traits. This is
            # used so we can preserve the dataclass behavior when the values in the
            # traits are updated, but still want to allow them to have a reference
            # to the rpc channel for sending commands.
            trait._rpc_channel = self._get_rpc_channel(trait)

    def _get_rpc_channel(self, trait: V1TraitMixin) -> V1RpcChannel:
        # The decorator `@common.mqtt_rpc_channel` means that the trait needs
        # to use the mqtt_rpc_channel (cloud only) instead of the rpc_channel (adaptive)
        if hasattr(trait, "mqtt_rpc_channel"):
            return self._mqtt_rpc_channel
        elif hasattr(trait, "map_rpc_channel"):
            return self._map_rpc_channel
        else:
            return self._rpc_channel

    async def start(self) -> None:
        """Start the properties API and discover features."""
        if self._unsub:
            return
        await self.discover_features()
        self._unsub = self._add_dps_listener(self._on_dps_update)

    def close(self) -> None:
        if self._unsub:
            self._unsub()
            self._unsub = None

    def _on_dps_update(self, dps: dict[RoborockDataProtocol, Any]) -> None:
        """Handle incoming messages from the device.

        This will notify all traits of the new values.
        """
        _LOGGER.debug("Received message from device: %s", dps)
        self.status.update_from_dps(dps)
        self.consumables.update_from_dps(dps)

    async def discover_features(self) -> None:
        """Populate any supported traits that were not initialized in __init__."""
        _LOGGER.debug("Starting optional trait discovery")
        await self.device_features.refresh()
        # Dock type also acts like a device feature for some traits.
        dock_type = await self._dock_type()

        # Initialize traits with special arguments before the generic loop
        if self.wash_towel_mode is None and self._is_supported(WashTowelModeTrait, "wash_towel_mode", dock_type):
            wash_towel_mode = WashTowelModeTrait(self.device_features)
            wash_towel_mode._rpc_channel = self._get_rpc_channel(wash_towel_mode)  # type: ignore[assignment]
            self.wash_towel_mode = wash_towel_mode

        # Dynamically create any traits that need to be populated
        for item in fields(self):
            if (trait := getattr(self, item.name, None)) is not None:
                continue
            if (union_args := get_args(item.type)) is None:
                raise ValueError(f"Unexpected non-union type for trait {item.name}: {item.type}")
            if len(union_args) != 2 or type(None) not in union_args:
                raise ValueError(f"Unexpected non-optional type for trait {item.name}: {item.type}")

            # Union args may not be in declared order
            item_type = union_args[0] if union_args[1] is type(None) else union_args[1]
            if not self._is_supported(item_type, item.name, dock_type):
                _LOGGER.debug("Trait '%s' not supported, skipping", item.name)
                continue
            _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
            trait = item_type()
            setattr(self, item.name, trait)
            trait._rpc_channel = self._get_rpc_channel(trait)

    def _is_supported(self, trait_type: type[V1TraitMixin], name: str, dock_type: RoborockDockTypeCode) -> bool:
        """Check if a trait is supported by the device."""

        if (requires_dock_type := getattr(trait_type, "requires_dock_type", None)) is not None:
            return requires_dock_type(dock_type)

        if (feature_name := getattr(trait_type, "requires_feature", None)) is None:
            _LOGGER.debug("Optional trait missing 'requires_feature' attribute %s, skipping", name)
            return False
        if (is_supported := getattr(self.device_features, feature_name)) is None:
            raise ValueError(f"Device feature '{feature_name}' on trait '{name}' is unknown")
        return is_supported

    async def _dock_type(self) -> RoborockDockTypeCode:
        """Get the dock type from the status trait or cache."""
        dock_type = await self._get_cached_trait_data("dock_type")
        if dock_type is not None:
            _LOGGER.debug("Using cached dock type: %s", dock_type)
            try:
                return RoborockDockTypeCode(dock_type)
            except ValueError:
                _LOGGER.debug("Cached dock type %s is invalid, refreshing", dock_type)

        _LOGGER.debug("Starting dock type discovery")
        await self.status.refresh()
        _LOGGER.debug("Fetched dock type: %s", self.status.dock_type)
        if self.status.dock_type is None:
            # Explicitly set so we reuse cached value next type
            dock_type = RoborockDockTypeCode.no_dock
        else:
            dock_type = self.status.dock_type
        await self._set_cached_trait_data("dock_type", dock_type)
        return dock_type

    async def _get_cached_trait_data(self, name: str) -> Any:
        """Get the dock type from the status trait or cache."""
        cache_data = await self._device_cache.get()
        if cache_data.trait_data is None:
            cache_data.trait_data = {}
        _LOGGER.debug("Cached trait data: %s", cache_data.trait_data)
        return cache_data.trait_data.get(name)

    async def _set_cached_trait_data(self, name: str, value: Any) -> None:
        """Set trait-specific cached data."""
        cache_data = await self._device_cache.get()
        if cache_data.trait_data is None:
            cache_data.trait_data = {}
        cache_data.trait_data[name] = value
        _LOGGER.debug("Updating cached trait data: %s", cache_data.trait_data)
        await self._device_cache.set(cache_data)

    def as_dict(self) -> dict[str, Any]:
        """Return the trait data as a dictionary."""
        result: dict[str, Any] = {}
        for item in fields(self):
            trait = getattr(self, item.name, None)
            if trait is None or not isinstance(trait, RoborockBase):
                continue
            data = trait.as_dict()
            if data:  # Don't omit unset traits
                result[item.name] = data
        return result


def create(
    device_uid: str,
    product: HomeDataProduct,
    home_data: HomeData,
    rpc_channel: V1RpcChannel,
    mqtt_rpc_channel: V1RpcChannel,
    map_rpc_channel: V1RpcChannel,
    add_dps_listener: Callable[[Callable[[dict[RoborockDataProtocol, Any]], None]], Callable[[], None]],
    web_api: UserWebApiClient,
    device_cache: DeviceCache,
    map_parser_config: MapParserConfig | None = None,
    region: str | None = None,
) -> PropertiesApi:
    """Create traits for V1 devices."""
    return PropertiesApi(
        device_uid,
        product,
        home_data,
        rpc_channel,
        mqtt_rpc_channel,
        map_rpc_channel,
        add_dps_listener,
        web_api,
        device_cache,
        map_parser_config,
        region=region,
    )