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