From a6ed79340f4e38f582c2d93b3e29e97318fcfa8b Mon Sep 17 00:00:00 2001
From: Mike Blouin <mike@michaelblouin.ca>
Date: Thu, 30 Dec 2021 12:49:48 -0500
Subject: [PATCH] add support for targeting a specific room

---
 sharkiqpy/sharkiq.py | 67 +++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 60 insertions(+), 7 deletions(-)

diff --git a/sharkiqpy/sharkiq.py b/sharkiqpy/sharkiq.py
index 536f754..1a2bae4 100644
--- a/sharkiqpy/sharkiq.py
+++ b/sharkiqpy/sharkiq.py
@@ -51,6 +51,7 @@ class OperatingModes(enum.IntEnum):
 @enum.unique
 class Properties(enum.Enum):
     """Useful properties"""
+    AREAS_TO_CLEAN = "Areas_To_Clean"
     BATTERY_CAPACITY = "Battery_Capacity"
     CHARGING_STATUS = "Charging_Status"
     CLEAN_COMPLETE = "CleanComplete"
@@ -66,6 +67,7 @@ class Properties(enum.Enum):
     RECHARGE_RESUME = "Recharge_Resume"
     RECHARGING_TO_RESUME = "Recharging_To_Resume"
     ROBOT_FIRMWARE_VERSION = "Robot_Firmware_Version"
+    ROBOT_ROOM_LIST = "Robot_Room_List"
     RSSI = "RSSI"
 
 
@@ -340,16 +342,67 @@ async def async_get_file_property(self, property_name: PropertyName) -> bytes:
     def _encode_room_list(self, rooms: List[str]):
         """Base64 encode the list of rooms to clean"""
         if not rooms:
-            raise ValueError('Room list must not be empty')
-        if len(rooms) > 3:
-            raise ValueError('At most three rooms may be given')
-        # These are a mystery to me, but they seem constant
-        header = b'\x80\x01\x0b\xca\x02'
-        footer = b'\x1a\x08155B43C4'
+            # By default, clean all rooms
+            return '*'
+
+        room_list = self._get_device_room_list()
+        print(f'Room list identifier is: {room_list["identifier"]}')
+
+        # Header explained:
+        # 0x80: Control character - some mode selection
+        # 0x01: Start of Heading Character
+        # 0x0B: Use Line Tabulation (entries separated by newlines)
+        # 0xca: Control character - purpose unknown
+        # 0x02: Start of text (indicates start of room list)
+        header = '\x80\x01\x0b\xca\x02'
+
+        # For each room in the list:
+        # - Insert a byte representing the length of the room name string
+        # - Add the room name
+        # - Join with newlines (presumably because of the 0x0B in the header)
+        rooms_enc = "\n".join([chr(len(room)) + room for room in rooms])
+
+        # The footer starts with control character 0x1A
+        # Then add the length indicator for the room list identifier
+        # Then add the room list identifier
+        footer = '\x1a' + chr(len(room_list['identifier'])) + room_list['identifier']
+
+        # Now that we've computed the room list and footer and know their lengths, finish building the header
+        # This character denotes the length of the remaining input
+        header += chr(0
+            + 1 # Add one for a newline following the length specifier
+            + len(rooms_enc)
+            + len(footer)
+        )
+        header += '\n' # This is the newline reference above
+
+        # Finally, join and base64 encode the parts
+        return base64.b64encode(
+            # First encode the string as latin_1 to get the right endianness
+            (header + rooms_enc + footer).encode('latin_1')
+            # Then return as a utf8 string for ease of handling
+        ).decode('utf8')
+
+    def _get_device_room_list(self):
+        """Gets the list of known rooms from the device, including the map identifier"""
+        room_list = self.get_property_value(Properties.ROBOT_ROOM_LIST)
+        split = room_list.split(':')
+        return {
+            # The room list is preceded by an identifier, which I believe identifies the list of rooms with the
+            # onboard map in the robot
+            'identifier': split[0],
+            'rooms': split[1:],
+        }
+
+    def get_room_list(self) -> List[str]:
+        """Gets the list of rooms known by the device"""
+        return self._get_device_room_list()['rooms']
 
     def clean_rooms(self, rooms: List[str]) -> None:
         payload = self._encode_room_list(rooms)
-        raise NotImplementedError
+        print('Room list payload: ' + payload)
+        self.set_property_value(Properties.AREAS_TO_CLEAN, payload)
+        self.set_operating_mode(OperatingModes.START)
 
 
 class SharkPropertiesView(abc.Mapping):