cloud_api.py

from __future__ import annotations

import asyncio
import logging
import threading
from abc import ABC
from asyncio import Lock
from typing import Any

import paho.mqtt.client as mqtt
from paho.mqtt.enums import MQTTErrorCode

# Mypy is not seeing this for some reason. It wants me to use the depreciated ReasonCodes
from paho.mqtt.reasoncodes import ReasonCode  # type: ignore

from .api import KEEPALIVE, RoborockClient
from .containers import DeviceData, UserData
from .exceptions import RoborockException, VacuumError
from .protocol import (
    Decoder,
    Encoder,
    create_mqtt_decoder,
    create_mqtt_encoder,
    create_mqtt_params,
)
from .roborock_future import RoborockFuture

_LOGGER = logging.getLogger(__name__)
CONNECT_REQUEST_ID = 0
DISCONNECT_REQUEST_ID = 1


class _Mqtt(mqtt.Client):
    """Internal MQTT client.

    This is a subclass of the Paho MQTT client that adds some additional functionality
    for error cases where things get stuck.
    """

    _thread: threading.Thread

    def __init__(self) -> None:
        """Initialize the MQTT client."""
        super().__init__(protocol=mqtt.MQTTv5)

    def maybe_restart_loop(self) -> None:
        """Ensure that the MQTT loop is running in case it previously exited."""
        if not self._thread or not self._thread.is_alive():
            if self._thread:
                _LOGGER.info("Stopping mqtt loop")
                super().loop_stop()
            _LOGGER.info("Starting mqtt loop")
            super().loop_start()


class RoborockMqttClient(RoborockClient, ABC):
    """Roborock MQTT client base class."""

    def __init__(self, user_data: UserData, device_info: DeviceData) -> None:
        """Initialize the Roborock MQTT client."""
        rriot = user_data.rriot
        if rriot is None:
            raise RoborockException("Got no rriot data from user_data")
        RoborockClient.__init__(self, device_info)
        mqtt_params = create_mqtt_params(rriot)
        self._mqtt_user = rriot.u
        self._hashed_user = mqtt_params.username
        self._mqtt_host = mqtt_params.host
        self._mqtt_port = mqtt_params.port

        self._mqtt_client = _Mqtt()
        self._mqtt_client.on_connect = self._mqtt_on_connect
        self._mqtt_client.on_message = self._mqtt_on_message
        # Due to the incorrect ReasonCode, it is confused by typing
        self._mqtt_client.on_disconnect = self._mqtt_on_disconnect  # type: ignore
        if mqtt_params.tls:
            self._mqtt_client.tls_set()

        self._mqtt_client.username_pw_set(mqtt_params.username, mqtt_params.password)
        self._waiting_queue: dict[int, RoborockFuture] = {}
        self._mutex = Lock()
        self._decoder: Decoder = create_mqtt_decoder(device_info.device.local_key)
        self._encoder: Encoder = create_mqtt_encoder(device_info.device.local_key)
        self.received_message_since_last_disconnect = False
        self._topic = f"rr/m/o/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}"

    def _mqtt_on_connect(
        self,
        client: mqtt.Client,
        userdata: object,
        flags: dict[str, int],
        rc: ReasonCode,
        properties: mqtt.Properties | None = None,
    ):
        connection_queue = self._waiting_queue.get(CONNECT_REQUEST_ID)
        if rc.is_failure:
            message = f"Failed to connect ({rc})"
            self._logger.error(message)
            if connection_queue:
                connection_queue.set_exception(VacuumError(message))
            else:
                self._logger.debug("Failed to notify connect future, not in queue")
            return
        self._logger.info(f"Connected to mqtt {self._mqtt_host}:{self._mqtt_port}")
        (result, mid) = self._mqtt_client.subscribe(self._topic)
        if result != 0:
            message = f"Failed to subscribe ({str(rc)})"
            self._logger.error(message)
            if connection_queue:
                connection_queue.set_exception(VacuumError(message))
            return
        self._logger.info(f"Subscribed to topic {self._topic}")
        if connection_queue:
            connection_queue.set_result(True)

    def _mqtt_on_message(self, *args, **kwargs):
        self.received_message_since_last_disconnect = True
        client, __, msg = args
        try:
            messages = self._decoder(msg.payload)
            super().on_message_received(messages)
        except Exception as ex:
            self._logger.exception(ex)

    def _mqtt_on_disconnect(
        self,
        client: mqtt.Client,
        data: object,
        flags: dict[str, int],
        rc: ReasonCode | None,
        properties: mqtt.Properties | None = None,
    ):
        try:
            exc = RoborockException(str(rc)) if rc is not None and rc.is_failure else None
            super().on_connection_lost(exc)
            connection_queue = self._waiting_queue.get(DISCONNECT_REQUEST_ID)
            if connection_queue:
                connection_queue.set_result(True)
        except Exception as ex:
            self._logger.exception(ex)

    def is_connected(self) -> bool:
        """Check if the mqtt client is connected."""
        return self._mqtt_client.is_connected()

    def _sync_disconnect(self) -> Any:
        if not self.is_connected():
            return None

        self._logger.info("Disconnecting from mqtt")
        disconnected_future = self._async_response(DISCONNECT_REQUEST_ID)
        rc = self._mqtt_client.disconnect()

        if rc == mqtt.MQTT_ERR_NO_CONN:
            disconnected_future.cancel()
            return None

        if rc != mqtt.MQTT_ERR_SUCCESS:
            disconnected_future.cancel()
            raise RoborockException(f"Failed to disconnect ({str(rc)})")

        return disconnected_future

    def _sync_connect(self) -> Any:
        if self.is_connected():
            self._mqtt_client.maybe_restart_loop()
            return None

        if self._mqtt_port is None or self._mqtt_host is None:
            raise RoborockException("Mqtt information was not entered. Cannot connect.")

        self._logger.debug("Connecting to mqtt")
        connected_future = self._async_response(CONNECT_REQUEST_ID)
        self._mqtt_client.connect(host=self._mqtt_host, port=self._mqtt_port, keepalive=KEEPALIVE)
        self._mqtt_client.maybe_restart_loop()
        return connected_future

    async def async_disconnect(self) -> None:
        async with self._mutex:
            if disconnected_future := self._sync_disconnect():
                # There are no errors set on this future
                await disconnected_future
            loop = asyncio.get_running_loop()
            await loop.run_in_executor(None, self._mqtt_client.loop_stop)

    async def async_connect(self) -> None:
        async with self._mutex:
            if connected_future := self._sync_connect():
                try:
                    await connected_future
                except VacuumError as err:
                    raise RoborockException(err) from err

    def _send_msg_raw(self, msg: bytes) -> None:
        info = self._mqtt_client.publish(
            f"rr/m/i/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}", msg
        )
        if info.rc != mqtt.MQTT_ERR_SUCCESS:
            raise RoborockException(f"Failed to publish ({mqtt.error_string(info.rc)})")

    async def _unsubscribe(self) -> MQTTErrorCode:
        """Unsubscribe from the topic."""
        loop = asyncio.get_running_loop()
        (result, mid) = await loop.run_in_executor(None, self._mqtt_client.unsubscribe, self._topic)

        if result != 0:
            message = f"Failed to unsubscribe ({mqtt.error_string(result)})"
            self._logger.error(message)
        else:
            self._logger.info(f"Unsubscribed from topic {self._topic}")
        return result

    async def _subscribe(self) -> MQTTErrorCode:
        """Subscribe to the topic."""
        loop = asyncio.get_running_loop()
        (result, mid) = await loop.run_in_executor(None, self._mqtt_client.subscribe, self._topic)

        if result != 0:
            message = f"Failed to subscribe ({mqtt.error_string(result)})"
            self._logger.error(message)
        else:
            self._logger.info(f"Subscribed to topic {self._topic}")
        return result

    async def _reconnect(self) -> None:
        """Reconnect to the MQTT broker."""
        await self.async_disconnect()
        await self.async_connect()

    async def _validate_connection(self) -> None:
        """Override the default validate connection to try to re-subscribe rather than disconnect.
        When something seems to be wrong with our connection, we should follow the following steps:
        1. Try to unsubscribe and resubscribe from the topic.
        2. If we don't end up getting a message, we should completely disconnect and reconnect to the MQTT broker.
        3. We will continue to try to disconnect and reconnect until we get a message.
        4. If we get a message, the next time connection is lost, We will go back to step 1.
        """
        # If we should no longer keep the current connection alive...
        if not self.should_keepalive():
            self._logger.info("Resetting Roborock connection due to keepalive timeout")
            if not self.received_message_since_last_disconnect:
                # If we have already tried to unsub and resub, and we are still in this state,
                # we should try to reconnect.
                return await self._reconnect()
            try:
                # Mark that we have tried to unsubscribe and resubscribe
                self.received_message_since_last_disconnect = False
                if await self._unsubscribe() != 0:
                    # If we fail to unsubscribe, reconnect to the broker
                    return await self._reconnect()
                if await self._subscribe() != 0:
                    # If we fail to subscribe, reconnected to the broker.
                    return await self._reconnect()

            except Exception:  # noqa
                # If we get any errors at all, we should just reconnect.
                return await self._reconnect()
        # Call connect to make sure everything is still in a good state.
        await self.async_connect()