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
 21
 22from roborock.data import CombinedMapInfo, MultiMapsListMapInfo, NamedRoomMapping, RoborockBase
 23from roborock.data.v1.v1_code_mappings import RoborockStateCode
 24from roborock.devices.cache import DeviceCache
 25from roborock.devices.traits.v1 import common
 26from roborock.exceptions import RoborockDeviceBusy, RoborockException, RoborockInvalidStatus
 27from roborock.roborock_typing import RoborockCommand
 28
 29from .map_content import MapContent, MapContentTrait
 30from .maps import MapsTrait
 31from .rooms import RoomsTrait
 32from .status import StatusTrait
 33
 34_LOGGER = logging.getLogger(__name__)
 35
 36MAP_SLEEP = 3
 37
 38
 39class HomeTrait(RoborockBase, common.V1TraitMixin):
 40    """Trait that represents a full view of the home layout."""
 41
 42    command = RoborockCommand.GET_MAP_V1  # This is not used
 43    converter = common.DefaultConverter(RoborockBase)  # 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.converter.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            _LOGGER.debug("Cannot perform home discovery without current map info")
111            self._discovery_completed = True
112            await self._update_home_cache({}, {})
113            return
114
115        home_map_info, home_map_content = await self._build_home_map_info()
116        _LOGGER.debug("Home discovery complete, caching data for %d maps", len(home_map_info))
117        self._discovery_completed = True
118        await self._update_home_cache(home_map_info, home_map_content)
119
120    async def _refresh_map_info(self, map_info: MultiMapsListMapInfo) -> CombinedMapInfo:
121        """Collect room data for a specific map and return CombinedMapInfo."""
122        await self._rooms_trait.refresh()
123
124        # We have room names from multiple sources:
125        # - The map_info.rooms which we just received from the MultiMapsList
126        # - RoomsTrait rooms come from the GET_ROOM_MAPPING command for the current device (only)
127        # - RoomsTrait rooms that are pulled from the cloud API
128        # We always prefer the RoomsTrait room names since they are always newer and
129        # just refreshed above.
130        rooms_map: dict[int, NamedRoomMapping] = {
131            **map_info.rooms_map,
132            **{room.segment_id: room for room in self._rooms_trait.rooms or ()},
133        }
134        return CombinedMapInfo(
135            map_flag=map_info.map_flag,
136            name=map_info.name,
137            rooms=list(rooms_map.values()),
138        )
139
140    async def _refresh_map_content(self) -> MapContent:
141        """Refresh the map content trait to get the latest map data."""
142        await self._map_content.refresh()
143        return MapContent(
144            image_content=self._map_content.image_content,
145            map_data=self._map_content.map_data,
146            raw_api_response=self._map_content.raw_api_response,
147        )
148
149    async def _build_home_map_info(self) -> tuple[dict[int, CombinedMapInfo], dict[int, MapContent]]:
150        """Perform the actual discovery and caching of home map info and content."""
151        home_map_info: dict[int, CombinedMapInfo] = {}
152        home_map_content: dict[int, MapContent] = {}
153
154        # Sort map_info to process the current map last, reducing map switching.
155        # False (non-original maps) sorts before True (original map). We ensure
156        # we load the original map last.
157        sorted_map_infos = sorted(
158            self._maps_trait.map_info or [],
159            key=lambda mi: mi.map_flag == self._maps_trait.current_map,
160            reverse=False,
161        )
162        _LOGGER.debug("Building home cache for maps: %s", [mi.map_flag for mi in sorted_map_infos])
163        for map_info in sorted_map_infos:
164            # We need to load each map to get its room data
165            if len(sorted_map_infos) > 1:
166                _LOGGER.debug("Loading map %s", map_info.map_flag)
167                try:
168                    await self._maps_trait.set_current_map(map_info.map_flag)
169                except RoborockInvalidStatus as ex:
170                    # Device is in a state that forbids map switching. Translate to
171                    # "busy" so callers can fall back to refreshing the current map only.
172                    raise RoborockDeviceBusy("Cannot switch maps right now (device action locked)") from ex
173                await asyncio.sleep(MAP_SLEEP)
174
175            map_content = await self._refresh_map_content()
176            home_map_content[map_info.map_flag] = map_content
177
178            combined_map_info = await self._refresh_map_info(map_info)
179            home_map_info[map_info.map_flag] = combined_map_info
180        return home_map_info, home_map_content
181
182    async def refresh(self) -> None:
183        """Refresh current map's underlying map and room data, updating cache as needed.
184
185        This will only refresh the current map's data and will not populate non
186        active maps or re-discover the home. It is expected that this will keep
187        information up to date for the current map as users switch to that map.
188        """
189        if not self._discovery_completed:
190            # Running initial discovery also populates all of the same information
191            # as below so we can just call that method. If the device is busy
192            # then we'll fall through below to refresh the current map only.
193            try:
194                await self.discover_home()
195                return
196            except RoborockDeviceBusy:
197                _LOGGER.debug("Cannot refresh home data while device is busy cleaning")
198
199        # Refresh the list of map names/info
200        await self._maps_trait.refresh()
201        if (current_map_info := self._maps_trait.current_map_info) is None or (
202            map_flag := self._maps_trait.current_map
203        ) is None:
204            _LOGGER.debug("Cannot refresh home data without current map info")
205            return
206
207        # Refresh the map content to ensure we have the latest image and object positions
208        new_map_content = await self._refresh_map_content()
209        # Refresh the current map's room data
210        combined_map_info = await self._refresh_map_info(current_map_info)
211        await self._update_current_map(
212            map_flag, combined_map_info, new_map_content, update_cache=self._discovery_completed
213        )
214
215    @property
216    def home_map_info(self) -> dict[int, CombinedMapInfo] | None:
217        """Returns the map information for all cached maps."""
218        return self._home_map_info
219
220    @property
221    def current_map_data(self) -> CombinedMapInfo | None:
222        """Returns the map data for the current map."""
223        current_map_flag = self._maps_trait.current_map
224        if current_map_flag is None or self._home_map_info is None:
225            return None
226        return self._home_map_info.get(current_map_flag)
227
228    @property
229    def current_rooms(self) -> list[NamedRoomMapping]:
230        """Returns the room names for the current map."""
231        if self.current_map_data is None:
232            return []
233        return self.current_map_data.rooms
234
235    @property
236    def home_map_content(self) -> dict[int, MapContent] | None:
237        """Returns the map content for all cached maps."""
238        return self._home_map_content
239
240    async def _update_home_cache(
241        self, home_map_info: dict[int, CombinedMapInfo], home_map_content: dict[int, MapContent]
242    ) -> None:
243        """Update the entire home cache with new map info and content."""
244        device_cache_data = await self._device_cache.get()
245        device_cache_data.home_map_info = home_map_info
246        device_cache_data.home_map_content_base64 = {
247            k: base64.b64encode(v.raw_api_response).decode("utf-8")
248            for k, v in home_map_content.items()
249            if v.raw_api_response
250        }
251        await self._device_cache.set(device_cache_data)
252        self._home_map_info = home_map_info
253        self._home_map_content = home_map_content
254
255    async def _update_current_map(
256        self,
257        map_flag: int,
258        map_info: CombinedMapInfo,
259        map_content: MapContent,
260        update_cache: bool,
261    ) -> None:
262        """Update the cache for the current map only."""
263        # Update the persistent cache if requested e.g. home discovery has
264        # completed and we want to keep it fresh. Otherwise just update the
265        # in memory map below.
266        if update_cache:
267            device_cache_data = await self._device_cache.get()
268            if device_cache_data.home_map_info is None:
269                device_cache_data.home_map_info = {}
270            device_cache_data.home_map_info[map_flag] = map_info
271            if map_content.raw_api_response:
272                if device_cache_data.home_map_content_base64 is None:
273                    device_cache_data.home_map_content_base64 = {}
274                device_cache_data.home_map_content_base64[map_flag] = base64.b64encode(
275                    map_content.raw_api_response
276                ).decode("utf-8")
277            await self._device_cache.set(device_cache_data)
278
279        if self._home_map_info is None:
280            self._home_map_info = {}
281        self._home_map_info[map_flag] = map_info
282
283        if self._home_map_content is None:
284            self._home_map_content = {}
285        self._home_map_content[map_flag] = map_content