roborock.devices.traits.v1.home

  1"""Trait that represents a full view of the home layout.
  2
  3This trait combines information about maps and rooms to provide a comprehensive
  4view of the home layout, including room names and their corresponding segment
  5on the map. It also makes it straight forward to fetch the map image and data.
  6
  7This trait depends on the MapsTrait and RoomsTrait to gather the necessary
  8information. It provides properties to access the current map, the list of
  9rooms with names, and the map image and data.
 10
 11Callers may first call `discover_home()` to populate the home layout cache by
 12iterating through all available maps on the device. This will cache the map
 13information and room names for all maps to minimize map switching and improve
 14performance. After the initial discovery, callers can call `refresh()` to update
 15the current map's information and room names as needed.
 16"""
 17
 18import asyncio
 19import base64
 20import logging
 21from typing import Self
 22
 23from roborock.data import CombinedMapInfo, NamedRoomMapping, RoborockBase
 24from roborock.data.v1.v1_code_mappings import RoborockStateCode
 25from roborock.devices.cache import DeviceCache
 26from roborock.devices.traits.v1 import common
 27from roborock.exceptions import RoborockDeviceBusy, RoborockException, RoborockInvalidStatus
 28from roborock.roborock_typing import RoborockCommand
 29
 30from .map_content import MapContent, MapContentTrait
 31from .maps import MapsTrait
 32from .rooms import RoomsTrait
 33from .status import StatusTrait
 34
 35_LOGGER = logging.getLogger(__name__)
 36
 37MAP_SLEEP = 3
 38
 39
 40class HomeTrait(RoborockBase, common.V1TraitMixin):
 41    """Trait that represents a full view of the home layout."""
 42
 43    command = RoborockCommand.GET_MAP_V1  # This is not used
 44
 45    def __init__(
 46        self,
 47        status_trait: StatusTrait,
 48        maps_trait: MapsTrait,
 49        map_content: MapContentTrait,
 50        rooms_trait: RoomsTrait,
 51        device_cache: DeviceCache,
 52    ) -> None:
 53        """Initialize the HomeTrait.
 54
 55        We keep track of the MapsTrait and RoomsTrait to provide a comprehensive
 56        view of the home layout. This also depends on the StatusTrait to determine
 57        the current map. See comments in MapsTrait for details on that dependency.
 58
 59        The cache is used to store discovered home data to minimize map switching
 60        and improve performance. The cache should be persisted by the caller to
 61        ensure data is retained across restarts.
 62
 63        After initial discovery, only information for the current map is refreshed
 64        to keep data up to date without excessive map switching. However, as
 65        users switch rooms, the current map's data will be updated to ensure
 66        accuracy.
 67        """
 68        super().__init__()
 69        self._status_trait = status_trait
 70        self._maps_trait = maps_trait
 71        self._map_content = map_content
 72        self._rooms_trait = rooms_trait
 73        self._device_cache = device_cache
 74        self._discovery_completed = False
 75        self._home_map_info: dict[int, CombinedMapInfo] | None = None
 76        self._home_map_content: dict[int, MapContent] | None = None
 77
 78    async def discover_home(self) -> None:
 79        """Iterate through all maps to discover rooms and cache them.
 80
 81        This will be a no-op if the home cache is already populated.
 82
 83        This cannot be called while the device is cleaning, as that would interrupt the
 84        cleaning process. This will raise `RoborockDeviceBusy` if the device is
 85        currently cleaning.
 86
 87        After discovery, the home cache will be populated and can be accessed via the `home_map_info` property.
 88        """
 89        device_cache_data = await self._device_cache.get()
 90        if device_cache_data and device_cache_data.home_map_info:
 91            _LOGGER.debug("Home cache already populated, skipping discovery")
 92            self._home_map_info = device_cache_data.home_map_info
 93            self._discovery_completed = True
 94            try:
 95                self._home_map_content = {
 96                    k: self._map_content.parse_map_content(base64.b64decode(v))
 97                    for k, v in (device_cache_data.home_map_content_base64 or {}).items()
 98                }
 99            except (ValueError, RoborockException) as ex:
100                _LOGGER.warning("Failed to parse cached home map content, will re-discover: %s", ex)
101                self._home_map_content = {}
102            else:
103                return
104
105        if self._status_trait.state == RoborockStateCode.cleaning:
106            raise RoborockDeviceBusy("Cannot perform home discovery while the device is cleaning")
107
108        await self._maps_trait.refresh()
109        if self._maps_trait.current_map_info is None:
110            raise RoborockException("Cannot perform home discovery without current map info")
111
112        home_map_info, home_map_content = await self._build_home_map_info()
113        _LOGGER.debug("Home discovery complete, caching data for %d maps", len(home_map_info))
114        self._discovery_completed = True
115        await self._update_home_cache(home_map_info, home_map_content)
116
117    async def _refresh_map_info(self, map_info) -> CombinedMapInfo:
118        """Collect room data for a specific map and return CombinedMapInfo."""
119        await self._rooms_trait.refresh()
120
121        rooms: dict[int, NamedRoomMapping] = {}
122        if map_info.rooms:
123            # Not all vacuums resopnd with rooms inside map_info.
124            for room in map_info.rooms:
125                if room.id is not None and room.iot_name_id is not None:
126                    rooms[room.id] = NamedRoomMapping(
127                        segment_id=room.id,
128                        iot_id=room.iot_name_id,
129                        name=room.iot_name or "Unknown",
130                    )
131
132        # Add rooms from rooms_trait. If room already exists and rooms_trait has "Unknown", don't override.
133        if self._rooms_trait.rooms:
134            for room in self._rooms_trait.rooms:
135                if room.segment_id is not None and room.name:
136                    if room.segment_id not in rooms or room.name != "Unknown":
137                        # Add the room to rooms if the room segment is not already in it
138                        # or if the room name isn't unknown.
139                        rooms[room.segment_id] = room
140
141        return CombinedMapInfo(
142            map_flag=map_info.map_flag,
143            name=map_info.name,
144            rooms=list(rooms.values()),
145        )
146
147    async def _refresh_map_content(self) -> MapContent:
148        """Refresh the map content trait to get the latest map data."""
149        await self._map_content.refresh()
150        return MapContent(
151            image_content=self._map_content.image_content,
152            map_data=self._map_content.map_data,
153            raw_api_response=self._map_content.raw_api_response,
154        )
155
156    async def _build_home_map_info(self) -> tuple[dict[int, CombinedMapInfo], dict[int, MapContent]]:
157        """Perform the actual discovery and caching of home map info and content."""
158        home_map_info: dict[int, CombinedMapInfo] = {}
159        home_map_content: dict[int, MapContent] = {}
160
161        # Sort map_info to process the current map last, reducing map switching.
162        # False (non-original maps) sorts before True (original map). We ensure
163        # we load the original map last.
164        sorted_map_infos = sorted(
165            self._maps_trait.map_info or [],
166            key=lambda mi: mi.map_flag == self._maps_trait.current_map,
167            reverse=False,
168        )
169        _LOGGER.debug("Building home cache for maps: %s", [mi.map_flag for mi in sorted_map_infos])
170        for map_info in sorted_map_infos:
171            # We need to load each map to get its room data
172            if len(sorted_map_infos) > 1:
173                _LOGGER.debug("Loading map %s", map_info.map_flag)
174                try:
175                    await self._maps_trait.set_current_map(map_info.map_flag)
176                except RoborockInvalidStatus as ex:
177                    # Device is in a state that forbids map switching. Translate to
178                    # "busy" so callers can fall back to refreshing the current map only.
179                    raise RoborockDeviceBusy("Cannot switch maps right now (device action locked)") from ex
180                await asyncio.sleep(MAP_SLEEP)
181
182            map_content = await self._refresh_map_content()
183            home_map_content[map_info.map_flag] = map_content
184
185            combined_map_info = await self._refresh_map_info(map_info)
186            home_map_info[map_info.map_flag] = combined_map_info
187        return home_map_info, home_map_content
188
189    async def refresh(self) -> None:
190        """Refresh current map's underlying map and room data, updating cache as needed.
191
192        This will only refresh the current map's data and will not populate non
193        active maps or re-discover the home. It is expected that this will keep
194        information up to date for the current map as users switch to that map.
195        """
196        if not self._discovery_completed:
197            # Running initial discovery also populates all of the same information
198            # as below so we can just call that method. If the device is busy
199            # then we'll fall through below to refresh the current map only.
200            try:
201                await self.discover_home()
202                return
203            except RoborockDeviceBusy:
204                _LOGGER.debug("Cannot refresh home data while device is busy cleaning")
205
206        # Refresh the list of map names/info
207        await self._maps_trait.refresh()
208        if (current_map_info := self._maps_trait.current_map_info) is None or (
209            map_flag := self._maps_trait.current_map
210        ) is None:
211            raise RoborockException("Cannot refresh home data without current map info")
212
213        # Refresh the map content to ensure we have the latest image and object positions
214        new_map_content = await self._refresh_map_content()
215        # Refresh the current map's room data
216        combined_map_info = await self._refresh_map_info(current_map_info)
217        await self._update_current_map(
218            map_flag, combined_map_info, new_map_content, update_cache=self._discovery_completed
219        )
220
221    @property
222    def home_map_info(self) -> dict[int, CombinedMapInfo] | None:
223        """Returns the map information for all cached maps."""
224        return self._home_map_info
225
226    @property
227    def current_map_data(self) -> CombinedMapInfo | None:
228        """Returns the map data for the current map."""
229        current_map_flag = self._maps_trait.current_map
230        if current_map_flag is None or self._home_map_info is None:
231            return None
232        return self._home_map_info.get(current_map_flag)
233
234    @property
235    def home_map_content(self) -> dict[int, MapContent] | None:
236        """Returns the map content for all cached maps."""
237        return self._home_map_content
238
239    def _parse_response(self, response: common.V1ResponseData) -> Self:
240        """This trait does not parse responses directly."""
241        raise NotImplementedError("HomeTrait does not support direct command responses")
242
243    async def _update_home_cache(
244        self, home_map_info: dict[int, CombinedMapInfo], home_map_content: dict[int, MapContent]
245    ) -> None:
246        """Update the entire home cache with new map info and content."""
247        device_cache_data = await self._device_cache.get()
248        device_cache_data.home_map_info = home_map_info
249        device_cache_data.home_map_content_base64 = {
250            k: base64.b64encode(v.raw_api_response).decode("utf-8")
251            for k, v in home_map_content.items()
252            if v.raw_api_response
253        }
254        await self._device_cache.set(device_cache_data)
255        self._home_map_info = home_map_info
256        self._home_map_content = home_map_content
257
258    async def _update_current_map(
259        self,
260        map_flag: int,
261        map_info: CombinedMapInfo,
262        map_content: MapContent,
263        update_cache: bool,
264    ) -> None:
265        """Update the cache for the current map only."""
266        # Update the persistent cache if requested e.g. home discovery has
267        # completed and we want to keep it fresh. Otherwise just update the
268        # in memory map below.
269        if update_cache:
270            device_cache_data = await self._device_cache.get()
271            if device_cache_data.home_map_info is None:
272                device_cache_data.home_map_info = {}
273            device_cache_data.home_map_info[map_flag] = map_info
274            if map_content.raw_api_response:
275                if device_cache_data.home_map_content_base64 is None:
276                    device_cache_data.home_map_content_base64 = {}
277                device_cache_data.home_map_content_base64[map_flag] = base64.b64encode(
278                    map_content.raw_api_response
279                ).decode("utf-8")
280            await self._device_cache.set(device_cache_data)
281
282        if self._home_map_info is None:
283            self._home_map_info = {}
284        self._home_map_info[map_flag] = map_info
285
286        if self._home_map_content is None:
287            self._home_map_content = {}
288        self._home_map_content[map_flag] = map_content