python-kasa/kasa/smart/modules/clean.py

"""Implementation of vacuum clean module."""

from __future__ import annotations

import base64
import logging
from dataclasses import dataclass
from datetime import timedelta
from enum import IntEnum, StrEnum
from typing import Annotated, Literal

from ...feature import Feature
from ...module import FeatureAttribute
from ..smartmodule import SmartModule

_LOGGER = logging.getLogger(__name__)

# Only known value for start_type in setSwitchClean; required for
# targeted cleaning modes (Room) but not for StandardHome.
_START_TYPE_DEFAULT = 1

# Only known value for the type parameter in getMapData.
_MAP_DATA_TYPE_DEFAULT = 0


class Status(IntEnum):
    """Status of vacuum."""

    Idle = 0
    Cleaning = 1
    Mapping = 2
    GoingHome = 4
    Charging = 5
    Charged = 6
    Paused = 7
    Undocked = 8
    Error = 100

    UnknownInternal = -1000


class ErrorCode(IntEnum):
    """Error codes for vacuum."""

    Ok = 0
    SideBrushStuck = 2
    MainBrushStuck = 3
    WheelBlocked = 4
    Trapped = 6
    TrappedCliff = 7
    DustBinRemoved = 14
    UnableToMove = 15
    LidarBlocked = 16
    UnableToFindDock = 21
    BatteryLow = 22

    UnknownInternal = -1000


class FanSpeed(IntEnum):
    """Fan speed level."""

    Quiet = 1
    Standard = 2
    Turbo = 3
    Max = 4
    Ultra = 5


class CleanMode(IntEnum):
    """Clean mode for ``setSwitchClean`` and ``getCleanStatus``.

    Used as ``clean_mode`` in commands and ``clean_status`` in status responses.
    """

    #: Clean all rooms with uniform settings.
    StandardHome = 0
    #: Clean all rooms with per-room settings and custom order.
    AdvancedHome = 1
    #: Clean a small area around the vacuum's current position.
    Spot = 2
    #: Clean selected rooms only.
    Room = 3
    #: Clean user-defined rectangular areas.
    Zone = 4
    #: Run a saved custom cleaning preset.
    Custom = 5


@dataclass
class CleanAreaSettings:
    """Per-area cleaning settings shared by rooms and zones."""

    #: Suction power level, or ``None`` when not configured per-area.
    suction: FanSpeed | None = None
    #: Water level for mopping.
    cistern: int = 0
    #: Number of cleaning passes.
    clean_number: int = 0


@dataclass
class RoomInfo(CleanAreaSettings):
    """Information about a room on the vacuum's map."""

    #: Room ID used in cleaning commands.
    id: int = 0
    #: Human-readable room name (base64-decoded from the device).
    name: str | None = None
    #: Color index used for map rendering.
    color: int = 0


class AreaType(StrEnum):
    """Type of area entry in map data."""

    #: A named room.
    Room = "room"
    #: A user-defined rectangular cleaning zone.
    Area = "area"
    #: A virtual wall boundary.
    VirtualWall = "virtual_wall"
    #: A no-go zone.
    Forbid = "forbid"
    #: A detected carpet region.
    CarpetRectangle = "carpet_rectangle"


class AreaUnit(IntEnum):
    """Area unit."""

    #: Square meter
    Sqm = 0
    #: Square feet
    Sqft = 1
    #: Taiwanese unit: https://en.wikipedia.org/wiki/Taiwanese_units_of_measurement#Area
    Ping = 2


class Clean(SmartModule):
    """Implementation of vacuum clean module."""

    REQUIRED_COMPONENT = "clean"
    _error_code = ErrorCode.Ok
    _logged_error_code_warnings: set | None = None
    _logged_status_code_warnings: set

    def _initialize_features(self) -> None:
        """Initialize features."""
        self._add_feature(
            Feature(
                self._device,
                id="vacuum_return_home",
                name="Return home",
                container=self,
                attribute_setter="return_home",
                category=Feature.Category.Primary,
                type=Feature.Action,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="vacuum_start",
                name="Start cleaning",
                container=self,
                attribute_setter="start",
                category=Feature.Category.Primary,
                type=Feature.Action,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="vacuum_pause",
                name="Pause",
                container=self,
                attribute_setter="pause",
                category=Feature.Category.Primary,
                type=Feature.Action,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="vacuum_status",
                name="Vacuum status",
                container=self,
                attribute_getter="status",
                category=Feature.Category.Primary,
                type=Feature.Type.Sensor,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="vacuum_error",
                name="Error",
                container=self,
                attribute_getter="error",
                category=Feature.Category.Info,
                type=Feature.Type.Sensor,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="battery_level",
                name="Battery level",
                container=self,
                attribute_getter="battery",
                icon="mdi:battery",
                unit_getter=lambda: "%",
                category=Feature.Category.Info,
                type=Feature.Type.Sensor,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="vacuum_fan_speed",
                name="Fan speed",
                container=self,
                attribute_getter="fan_speed_preset",
                attribute_setter="set_fan_speed_preset",
                icon="mdi:fan",
                choices_getter=lambda: list(FanSpeed.__members__),
                category=Feature.Category.Primary,
                type=Feature.Type.Choice,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="clean_count",
                name="Clean count",
                container=self,
                attribute_getter="clean_count",
                attribute_setter="set_clean_count",
                range_getter=lambda: (1, 3),
                category=Feature.Category.Config,
                type=Feature.Type.Number,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="carpet_boost",
                name="Carpet boost",
                container=self,
                attribute_getter="carpet_boost",
                attribute_setter="set_carpet_boost",
                icon="mdi:rug",
                category=Feature.Category.Config,
                type=Feature.Type.Switch,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="clean_area",
                name="Cleaning area",
                container=self,
                attribute_getter="clean_area",
                unit_getter="area_unit",
                category=Feature.Category.Info,
                type=Feature.Type.Sensor,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="clean_time",
                name="Cleaning time",
                container=self,
                attribute_getter="clean_time",
                category=Feature.Category.Info,
                type=Feature.Type.Sensor,
            )
        )
        self._add_feature(
            Feature(
                self._device,
                id="clean_progress",
                name="Cleaning progress",
                container=self,
                attribute_getter="clean_progress",
                unit_getter=lambda: "%",
                category=Feature.Category.Info,
                type=Feature.Type.Sensor,
            )
        )

    async def _post_update_hook(self) -> None:
        """Set error code after update."""
        if self._logged_error_code_warnings is None:
            self._logged_error_code_warnings = set()
            self._logged_status_code_warnings = set()

        errors = self._vac_status.get("err_status")
        if errors is None or not errors:
            self._error_code = ErrorCode.Ok
            return

        if len(errors) > 1 and "multiple" not in self._logged_error_code_warnings:
            self._logged_error_code_warnings.add("multiple")
            _LOGGER.warning(
                "Multiple error codes, using the first one only: %s", errors
            )

        error = errors.pop(0)
        try:
            self._error_code = ErrorCode(error)
        except ValueError:
            if error not in self._logged_error_code_warnings:
                self._logged_error_code_warnings.add(error)
                _LOGGER.warning(
                    "Unknown error code, please create an issue "
                    "describing the error: %s",
                    error,
                )
            self._error_code = ErrorCode.UnknownInternal

    def query(self) -> dict:
        """Query to execute during the update cycle."""
        return {
            "getVacStatus": {},
            "getCleanInfo": {},
            "getCarpetClean": {},
            "getAreaUnit": {},
            "getBatteryInfo": {},
            "getCleanStatus": {},
            "getCleanAttr": {"type": "global"},
            "getMapInfo": {},
        }

    async def start(self) -> dict:
        """Start cleaning."""
        # If we are paused, do not restart cleaning

        if self.status is Status.Paused:
            return await self.resume()

        return await self.call(
            "setSwitchClean",
            {
                "clean_mode": CleanMode.StandardHome,
                "clean_on": True,
                "clean_order": True,
                "force_clean": False,
            },
        )

    async def pause(self) -> dict:
        """Pause cleaning."""
        if self.status is Status.GoingHome:
            return await self.set_return_home(False)

        return await self.set_pause(True)

    async def resume(self) -> dict:
        """Resume cleaning."""
        return await self.set_pause(False)

    async def set_pause(self, enabled: bool) -> dict:
        """Pause or resume cleaning."""
        return await self.call("setRobotPause", {"pause": enabled})

    async def return_home(self) -> dict:
        """Return home."""
        return await self.set_return_home(True)

    async def set_return_home(self, enabled: bool) -> dict:
        """Return home / pause returning."""
        return await self.call("setSwitchCharge", {"switch_charge": enabled})

    @property
    def error(self) -> ErrorCode:
        """Return error."""
        return self._error_code

    @property
    def fan_speed_preset(self) -> Annotated[str, FeatureAttribute()]:
        """Return fan speed preset."""
        return FanSpeed(self._settings["suction"]).name

    async def set_fan_speed_preset(
        self, speed: str
    ) -> Annotated[dict, FeatureAttribute]:
        """Set fan speed preset."""
        name_to_value = {x.name: x.value for x in FanSpeed}
        if speed not in name_to_value:
            raise ValueError("Invalid fan speed %s, available %s", speed, name_to_value)
        return await self._change_setting("suction", name_to_value[speed])

    async def _change_setting(
        self, name: str, value: int, *, scope: Literal["global", "pose"] = "global"
    ) -> dict:
        """Change device setting."""
        params = {
            name: value,
            "type": scope,
        }
        return await self.call("setCleanAttr", params)

    @property
    def battery(self) -> int:
        """Return battery level."""
        return self.data["getBatteryInfo"]["battery_percentage"]

    @property
    def _vac_status(self) -> dict:
        """Return vac status container."""
        return self.data["getVacStatus"]

    @property
    def _info(self) -> dict:
        """Return current cleaning info."""
        return self.data["getCleanInfo"]

    @property
    def _settings(self) -> dict:
        """Return cleaning settings."""
        return self.data["getCleanAttr"]

    @property
    def status(self) -> Status:
        """Return current status."""
        if self._error_code is not ErrorCode.Ok:
            return Status.Error

        status_code = self._vac_status["status"]
        try:
            return Status(status_code)
        except ValueError:
            if status_code not in self._logged_status_code_warnings:
                self._logged_status_code_warnings.add(status_code)
                _LOGGER.warning(
                    "Got unknown status code: %s (%s)", status_code, self.data
                )
            return Status.UnknownInternal

    @property
    def carpet_boost(self) -> bool:
        """Return carpet boost mode."""
        return self.data["getCarpetClean"]["carpet_clean_prefer"] == "boost"

    async def set_carpet_boost(self, on: bool) -> dict:
        """Set carpet clean mode."""
        mode = "boost" if on else "normal"
        return await self.call("setCarpetClean", {"carpet_clean_prefer": mode})

    @property
    def area_unit(self) -> AreaUnit:
        """Return area unit."""
        return AreaUnit(self.data["getAreaUnit"]["area_unit"])

    @property
    def clean_area(self) -> Annotated[int, FeatureAttribute()]:
        """Return currently cleaned area."""
        return self._info["clean_area"]

    @property
    def clean_time(self) -> timedelta:
        """Return current cleaning time."""
        return timedelta(minutes=self._info["clean_time"])

    @property
    def clean_progress(self) -> int:
        """Return amount of currently cleaned area."""
        return self._info["clean_percent"]

    @property
    def clean_count(self) -> Annotated[int, FeatureAttribute()]:
        """Return number of times to clean."""
        return self._settings["clean_number"]

    async def set_clean_count(self, count: int) -> Annotated[dict, FeatureAttribute()]:
        """Set number of times to clean."""
        return await self._change_setting("clean_number", count)

    @property
    def current_map_id(self) -> int:
        """Return the ID of the currently active map."""
        return self.data["getMapInfo"]["current_map_id"]

    @property
    def current_map_name(self) -> str | None:
        """Return the name of the currently active map, or ``None``."""
        map_info = self.data["getMapInfo"]
        current_id = map_info["current_map_id"]
        for m in map_info.get("map_list", []):
            if m.get("map_id") == current_id:
                if raw_name := m.get("map_name"):
                    try:
                        return base64.b64decode(raw_name).decode()
                    except Exception:
                        return raw_name
                return None
        return None

    @property
    def clean_type(self) -> CleanMode | None:
        """Return the active cleaning mode, or ``None`` if unavailable."""
        cs = self.data.get("getCleanStatus")
        if cs is None or "clean_status" not in cs:
            return None
        return CleanMode(cs["clean_status"])

    async def clean_rooms(
        self, room_ids: list[int], *, map_id: int | None = None
    ) -> dict:
        """Start cleaning specific rooms.

        Per-room settings are not supported; the device uses the global
        suction / cistern / clean_number values for room cleaning.

        :param room_ids: List of room IDs to clean.
        :param map_id: Map ID to clean on. Defaults to the current active map.
        """
        if not room_ids:
            raise ValueError("room_ids must not be empty")
        if map_id is None:
            map_id = self.current_map_id
        return await self.call(
            "setSwitchClean",
            {
                "clean_mode": CleanMode.Room,
                "clean_on": True,
                "clean_order": True,
                "force_clean": False,
                "map_id": map_id,
                "room_list": list(room_ids),
                "start_type": _START_TYPE_DEFAULT,
            },
        )

    async def get_rooms(self, map_id: int | None = None) -> list[RoomInfo]:
        """Return the list of rooms for the given map.

        Room names are base64-decoded when present.

        :param map_id: Map ID to query. Defaults to the current active map.
        """
        if map_id is None:
            map_id = self.current_map_id
        resp = await self.call(
            "getMapData", {"map_id": map_id, "type": _MAP_DATA_TYPE_DEFAULT}
        )
        map_data = resp.get("getMapData", resp)

        rooms: list[RoomInfo] = []
        for area in map_data.get("area_list", []):
            if area.get("type") != AreaType.Room:
                _LOGGER.debug(
                    "Skipping non-room area: type=%s, id=%s",
                    area.get("type"),
                    area.get("id"),
                )
                continue
            name = None
            if raw_name := area.get("name"):
                try:
                    name = base64.b64decode(raw_name).decode()
                except Exception:
                    name = raw_name

            suction_val = area.get("suction")
            rooms.append(
                RoomInfo(
                    id=area["id"],
                    name=name,
                    color=area.get("color", 0),
                    suction=FanSpeed(suction_val) if suction_val else None,
                    cistern=area.get("cistern", 0),
                    clean_number=area.get("clean_number", 0),
                )
            )
        return rooms