Implement MQTT support for BambuVirtualPrinter, including connection, message handling, and publishing commands

octoprint_bambu_printer/printer/bambu_virtual_printer.py

@@ -14,6 +14,9 @@ from octoprint_bambu_printer.printer.print_job import PrintJob
 from pybambu import BambuClient, commands
 import logging
 import logging.handlers
+import paho.mqtt.client as mqtt
+import json
+import ssl
 
 from octoprint.util import RepeatedTimer
 
@@ -105,6 +108,10 @@ class BambuVirtualPrinter:
         self._serial_io.start()
         self._printer_thread.start()
 
+        self._mqtt_client = None
+        self._mqtt_connected = False
+        self._bambu_client = None
+
         self._bambu_client: BambuClient = self._create_client_connection_async()
 
     @property
@@ -220,6 +227,48 @@ class BambuVirtualPrinter:
         self._log.debug(f"on connect called")
         return on_connect
 
+    def _on_mqtt_connect(self, client, userdata, flags, rc):
+        self._log.debug(f"MQTT connected with result code: {rc}")
+        if rc == 0:
+            self._mqtt_connected = True
+
+            # Subscribe to the relevant topics for the Bambu printer
+            device_topic = f"device/{self._settings.get(['serial'])}/report"
+            client.subscribe(device_topic)
+
+            self._log.debug(f"Subscribed to topic: {device_topic}")
+
+            # Notify that we're connected
+            self.sendOk()
+        else:
+            self._mqtt_connected = False
+            self._log.error(f"Failed to connect to MQTT broker with result code: {rc}")
+
+    def _on_mqtt_disconnect(self, client, userdata, rc):
+        self._mqtt_connected = False
+        self._log.debug(f"MQTT disconnected with result code: {rc}")
+
+    def _on_mqtt_message(self, client, userdata, msg):
+        try:
+            # Decode message and update client data
+            payload = json.loads(msg.payload.decode('utf-8'))
+
+            # If this is a Bambu Lab printer message, process it
+            if 'print' in payload or 'info' in payload:
+                # Forward the message to pybambu for processing
+                if self._bambu_client:
+                    self._bambu_client._process_message(msg.topic, payload)
+
+                    # Trigger our update handler
+                    if 'print' in payload:
+                        self.new_update("event_printer_data_update")
+                    if 'info' in payload and 'hms' in payload['info']:
+                        self.new_update("event_hms_errors")
+
+            self._log.debug(f"MQTT message received on topic {msg.topic}")
+        except Exception as e:
+            self._log.error(f"Error processing MQTT message: {e}")
+
     def _create_client_connection_async(self):
         self._create_client_connection()
         if self._bambu_client is None:
@@ -237,27 +286,82 @@ class BambuVirtualPrinter:
             self._log.debug(msg)
             raise ValueError(msg)
 
-        self._log.debug(
+        use_local_mqtt = self._settings.get_boolean(['local_mqtt'])
-            f"connecting via local mqtt: {self._settings.get_boolean(['local_mqtt'])}"
+        self._log.debug(f"connecting via local mqtt: {use_local_mqtt}")
-        )
+
+        # Create a BambuClient but don't let it handle the MQTT connection
         bambu_client = BambuClient(
             device_type=self._settings.get(["device_type"]),
             serial=self._settings.get(["serial"]),
             host=self._settings.get(["host"]),
-            username=("bambuocto"),
+            username="bambuocto",
             access_code=self._settings.get(["access_code"]),
-            local_mqtt=self._settings.get_boolean(["local_mqtt"]),
+            local_mqtt=use_local_mqtt,
             region=self._settings.get(["region"]),
             email=self._settings.get(["email"]),
             auth_token=self._settings.get(["auth_token"]),
         )
-        bambu_client.on_disconnect = self.on_disconnect(bambu_client.on_disconnect)
+
-        bambu_client.on_connect = self.on_connect(bambu_client.on_connect)
+        # Set up our own MQTT client
-        bambu_client.connect(callback=self.new_update)
+        self._mqtt_client = mqtt.Client()
-        self._log.info(f"bambu connection status: {bambu_client.connected}")
+        self._mqtt_client.on_connect = self._on_mqtt_connect
-        self.sendOk()
+        self._mqtt_client.on_disconnect = self._on_mqtt_disconnect
+        self._mqtt_client.on_message = self._on_mqtt_message
+
+        # Configure connection based on local or cloud
+        if use_local_mqtt:
+            host = self._settings.get(["host"])
+            port = 1883
+            username = "octobambu"
+
+            self._mqtt_client.username_pw_set(username)
+        else:
+            # Cloud connection settings
+            region = self._settings.get(["region"])
+            host = f"mqtt-{region}.bambulab.com"
+            port = 8883
+            username = self._settings.get(["email"])
+            password = self._settings.get(["auth_token"])
+
+            self._mqtt_client.username_pw_set(username, password)
+            self._mqtt_client.tls_set()
+
+        # Connect MQTT
+        try:
+            self._mqtt_client.connect(host, port, 60)
+            self._mqtt_client.loop_start()
+            self._log.info(f"MQTT client started with {host}:{port}")
+        except Exception as e:
+            self._log.error(f"Failed to connect to MQTT broker: {e}")
+            raise
+
+        # Inject our MQTT client into the BambuClient
+        bambu_client._mqtt_client = self._mqtt_client
+        bambu_client.connected = True
+
+        # Store the Bambu client
         self._bambu_client = bambu_client
 
+        self._log.info(f"Bambu connection status: {bambu_client.connected}")
+
+    def publish_mqtt(self, topic, payload):
+        """Publish a message to the MQTT broker"""
+        if self._mqtt_client and self._mqtt_connected:
+            return self._mqtt_client.publish(topic, json.dumps(payload))
+        return False
+
+    # Override BambuClient's publish method to use our MQTT client
+    def publish(self, command):
+        """Publish a command using our MQTT client"""
+        if not self._mqtt_connected:
+            self._log.error("Cannot publish command: MQTT not connected")
+            return False
+
+        serial = self._settings.get(["serial"])
+        topic = f"device/{serial}/request"
+
+        return self.publish_mqtt(topic, command)
+
     def __str__(self):
         return "BAMBU(read_timeout={read_timeout},write_timeout={write_timeout},options={options})".format(
             read_timeout=self.timeout,
@@ -630,6 +734,9 @@ class BambuVirtualPrinter:
             return False
 
     def close(self):
+        if self._mqtt_client and self._mqtt_connected:
+            self._mqtt_client.loop_stop()
+            self._mqtt_client.disconnect()
         if self.bambu_client.connected:
             self.bambu_client.disconnect()
         self.change_state(self._state_idle)