continuum-proxy/src/printer/bambu_mqtt.rs

//! The real MQTT wiring for Bambu printers - the piece `connect`/
//! `send_gcode_line` in `bambu.rs` used to just stub out with a
//! `println!`. Three things make this more than "call `publish`, get a
//! `String` back":
//!
//! 1. MQTT v3.1.1 (BambuMQTT v1) and MQTT v5 (BambuMQTT v2) are two
//!    entirely separate rumqttc client types (`rumqttc::AsyncClient` vs
//!    `rumqttc::v5::AsyncClient`) with different event/error types -
//!    `BambuMqttClient` is the enum that hides which one a connection
//!    ended up using, same pattern as `PrinterHandle` hides which vendor
//!    struct you're holding.
//! 2. A response to a published request doesn't come back from that same
//!    call - it arrives later, asynchronously, on a *subscription* to
//!    `device/{sn}/report`. That needs a background task continuously
//!    polling the connection (rumqttc doesn't do this for you - nothing
//!    happens until you poll the `EventLoop` yourself, repeatedly,
//!    forever) - this is the first real use of `tokio::spawn` in this
//!    crate, flagged as deferred since `printer/mod.rs`'s very first doc
//!    comment, and this is why: it only makes sense once there's a
//!    long-running background job that has to run independently of
//!    whatever call site is waiting on a response.
//!
//!    Correlating a reply to its request isn't as uniform as it looks -
//!    confirmed live, not assumed. `info.get_version`'s reply echoes back
//!    the `sequence_id` you sent; `request` relies on that. A tempting
//!    generalization - add your own extra field instead of `sequence_id`,
//!    since most replies are supposed to echo back whatever extra fields
//!    a request carried - turned out false for this exact command: a real
//!    H2C's `get_version` reply had precisely `["command", "module",
//!    "sequence_id"]`, nothing else, so a homemade correlation field never
//!    survived and every request timed out. `sequence_id` - proven to
//!    round-trip, not just assumed to - is what's actually used.
//!    `print.push_status` doesn't echo `sequence_id` either (the printer
//!    stamps that with its own free-running counter regardless of what
//!    triggered it) - but it was never a request/reply case to begin with
//!    (`pushing.pushall` triggers an ongoing *stream*, not a single
//!    reply) - see `PrintStatusCache` for how that's handled instead.
//! 3. `print.push_status` also needs to be kept as running in-memory
//!    state, not just handed to whoever happens to be waiting at the
//!    moment it arrives - see `PrintStatusCache` below for why (short
//!    version: v1 only ever sends partial updates, so something has to
//!    remember what the fields *were* to produce a coherent whole
//!    object).
//!
//! Unverified against a live printer (no way to force it in a test): the
//! exact `ConnectReturnCode`/error shape a real Bambu printer's embedded
//! broker gives when it doesn't understand MQTT v5's CONNECT at all
//! (rather than a clean version-rejection CONNACK) - `open()` below
//! treats *any* failure of the v5 attempt as "fall back to v3.1.1", not
//! just a precisely-matched version-rejection code, since a broker that
//! can't parse a v5 CONNECT packet may just drop the connection instead
//! of replying with one.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use rumqttc::v5::mqttbytes::v5::Packet as PacketV5;
use rumqttc::v5::mqttbytes::QoS as QoSV5;
use rumqttc::{Event, Packet, QoS, Transport, TlsConfiguration};
use serde_json::Value;
use tokio::sync::{oneshot, Mutex};

use super::bambu::BambuTls;
use super::bambu_commands::PrintReport;
use super::BambuProtocolVersion;
use super::PrinterError;

/// Requests waiting on a report with a matching `sequence_id` - only
/// valid for a command empirically confirmed to echo it back, like
/// `get_version` (see this module's doc comment - a homemade correlation
/// field was tried instead and confirmed *not* to survive on this
/// command's reply). Shared between whichever call registered one (in
/// `request`) and the background poll task that completes it.
type PendingTable = Arc<Mutex<HashMap<String, oneshot::Sender<Value>>>>;

const REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
/// Both rumqttc clients default to a 10KB max packet size in *both*
/// directions - fine for a request, nowhere near enough for a reply:
/// confirmed live against a real H2C, whose `get_version` reply alone
/// (25625 bytes - it reports a lot of modules: dual AMS, AMS HT, six WTM
/// units, ...) blew straight through that limit and killed the whole
/// connection outright (a real MQTT v5 protocol error, not a timeout).
/// `print.push_status` reports (full AMS/xcam state) can get large too.
/// 1MB has real headroom for that.
const MAX_PACKET_SIZE: usize = 1024 * 1024;

enum BambuMqttClient {
    V1(rumqttc::AsyncClient),
    V2(rumqttc::v5::AsyncClient),
}

/// The running in-memory `print.push_status` object, plus enough state to
/// know how to fold the next update into it:
///
/// - BambuMQTT v2 sends the *entire* object every time - always replace.
/// - BambuMQTT v1 only sends the fields that changed since last time -
///   merge into the existing object, or the picture goes incomplete the
///   instant anything not present in the latest push gets treated as
///   "unset". The exception is a `pushing.pushall`-triggered push, which
///   is always the full object - and v1 self-describes which kind a given
///   message is via its own `msg` field (`1` automatic/partial, `0`
///   forced/full), confirmed against a real printer. `push_status` also
///   doesn't echo `sequence_id` (see this module's doc comment), so it
///   couldn't be identified via `request`'s correlation even if the
///   timing worked out; `msg` needs none of that: it's on the message
///   itself, correct regardless of which client (or nothing at all)
///   triggered it, and regardless of timing.
/// - `last_sequence_id` exists to notice when v1's organic stream skipped
///   an update (network hiccup, a message this client wasn't fast enough
///   to read): if the incoming `sequence_id` isn't exactly one more than
///   the last one seen, the merged object is missing whatever changed in
///   the gap. There's no way to know what that baseline number even is
///   ahead of time - the printer runs its own counter, not one this
///   client started - so the very first `push_status` received after
///   connecting just seeds `last_sequence_id` rather than being checked
///   against anything.
struct PrintStatusCache {
    report: Option<PrintReport>,
    last_sequence_id: Option<u64>,
}

impl PrintStatusCache {
    fn new() -> Self {
        Self { report: None, last_sequence_id: None }
    }

    fn apply(&mut self, incoming: PrintReport, version: BambuProtocolVersion) {
        let incoming_seq = incoming.envelope.sequence_id.parse::<u64>().ok();
        // A forced/full update isn't missing anything by definition, so a
        // sequence gap right after one (or on v2, where every message is
        // full) isn't a real problem worth warning about.
        let is_full_object = matches!(version, BambuProtocolVersion::V2) || incoming.msg != Some(1);

        if let (Some(last), Some(incoming_seq), false) = (self.last_sequence_id, incoming_seq, is_full_object) {
            if incoming_seq != last.wrapping_add(1) {
                eprintln!(
                    "print.push_status sequence gap: expected {}, got {incoming_seq} - the in-memory \
                     object may be missing whatever changed in between",
                    last.wrapping_add(1)
                );
            }
        }
        if let Some(incoming_seq) = incoming_seq {
            self.last_sequence_id = Some(incoming_seq);
        }

        match (&mut self.report, is_full_object) {
            (_, true) => self.report = Some(incoming),
            (Some(existing), false) => existing.merge_from(incoming),
            (None, false) => self.report = Some(incoming),
        }
    }
}

/// One live MQTT connection to a Bambu printer - whichever protocol
/// version `open()` settled on.
pub struct BambuMqttSession {
    client: BambuMqttClient,
    pending: PendingTable,
    print_status: Arc<Mutex<PrintStatusCache>>,
    sn: String,
    next_sequence_id: AtomicU64,
}

impl BambuMqttSession {
    /// `ForceV1`/`ForceV2` connect with that protocol directly. `Auto`
    /// tries v5 first and falls back to v3.1.1 if that attempt fails -
    /// see this module's doc comment for why "any failure", not just a
    /// precisely-matched version-rejection code.
    pub async fn open(
        host: &str,
        access_code: &str,
        sn: &str,
        tls: &BambuTls,
        preference: super::bambu::BambuProtocolPreference,
    ) -> Result<(Self, BambuProtocolVersion), PrinterError> {
        use super::bambu::BambuProtocolPreference as Pref;

        match preference {
            Pref::ForceV1 => {
                let session = Self::open_v1(host, access_code, sn, tls).await?;
                Ok((session, BambuProtocolVersion::V1))
            }
            Pref::ForceV2 => {
                let session = Self::open_v2(host, access_code, sn, tls).await?;
                Ok((session, BambuProtocolVersion::V2))
            }
            Pref::Auto => match Self::open_v2(host, access_code, sn, tls).await {
                Ok(session) => Ok((session, BambuProtocolVersion::V2)),
                Err(_) => {
                    let session = Self::open_v1(host, access_code, sn, tls).await?;
                    Ok((session, BambuProtocolVersion::V1))
                }
            },
        }
    }

    async fn open_v1(host: &str, access_code: &str, sn: &str, tls: &BambuTls) -> Result<Self, PrinterError> {
        let connector = tls.build_connector().map_err(|err| PrinterError(format!("TLS setup failed: {err}")))?;

        let mut options = rumqttc::MqttOptions::new(format!("continuum-{sn}"), host, 8883);
        options.set_credentials("bblp", access_code);
        options.set_transport(Transport::tls_with_config(TlsConfiguration::NativeConnector(connector)));
        options.set_keep_alive(Duration::from_secs(30));
        options.set_max_packet_size(MAX_PACKET_SIZE, MAX_PACKET_SIZE);

        let (client, mut event_loop) = rumqttc::AsyncClient::new(options, 10);
        client
            .subscribe(format!("device/{sn}/report"), QoS::AtMostOnce)
            .await
            .map_err(|err| PrinterError(format!("subscribe failed: {err}")))?;

        // Nothing actually happens on the wire until the eventloop is
        // polled - this first poll is what makes the TCP+TLS+CONNECT
        // handshake happen, and its result is the only way to know
        // whether it worked.
        match event_loop.poll().await {
            Ok(_) => {}
            Err(err) => return Err(PrinterError(format!("MQTT v3.1.1 connect failed: {err}"))),
        }

        let pending: PendingTable = Arc::new(Mutex::new(HashMap::new()));
        let print_status = Arc::new(Mutex::new(PrintStatusCache::new()));
        let sn_owned = sn.to_string();
        tokio::spawn(poll_loop_v1(event_loop, pending.clone(), print_status.clone(), BambuProtocolVersion::V1, sn_owned));

        Ok(Self {
            client: BambuMqttClient::V1(client),
            pending,
            print_status,
            sn: sn.to_string(),
            next_sequence_id: AtomicU64::new(1),
        })
    }

    async fn open_v2(host: &str, access_code: &str, sn: &str, tls: &BambuTls) -> Result<Self, PrinterError> {
        let connector = tls.build_connector().map_err(|err| PrinterError(format!("TLS setup failed: {err}")))?;

        let mut options = rumqttc::v5::MqttOptions::new(format!("continuum-{sn}"), host, 8883);
        options.set_credentials("bblp", access_code);
        options.set_transport(Transport::tls_with_config(TlsConfiguration::NativeConnector(connector)));
        options.set_keep_alive(Duration::from_secs(30));
        options.set_max_packet_size(Some(MAX_PACKET_SIZE as u32));

        let (client, mut event_loop) = rumqttc::v5::AsyncClient::new(options, 10);
        client
            .subscribe(format!("device/{sn}/report"), QoSV5::AtMostOnce)
            .await
            .map_err(|err| PrinterError(format!("subscribe failed: {err}")))?;

        match event_loop.poll().await {
            Ok(_) => {}
            Err(err) => return Err(PrinterError(format!("MQTT v5 connect failed: {err}"))),
        }

        let pending: PendingTable = Arc::new(Mutex::new(HashMap::new()));
        let print_status = Arc::new(Mutex::new(PrintStatusCache::new()));
        let sn_owned = sn.to_string();
        tokio::spawn(poll_loop_v2(event_loop, pending.clone(), print_status.clone(), BambuProtocolVersion::V2, sn_owned));

        Ok(Self {
            client: BambuMqttClient::V2(client),
            pending,
            print_status,
            sn: sn.to_string(),
            next_sequence_id: AtomicU64::new(1),
        })
    }

    /// Publishes `body` under `device/{sn}/request`, wrapped as
    /// `{category: body}` (Bambu's top-level key names the *category*,
    /// e.g. "print"/"info" - `body`'s own `command` field says which
    /// specific one), and waits up to `REQUEST_TIMEOUT` for a report
    /// whose `sequence_id` matches the one this call sent.
    ///
    /// This crate tried adding its own correlation field
    /// (`continuum_correlation_id`) instead of relying on `sequence_id`,
    /// on the understanding that extra fields survive into most replies.
    /// Confirmed live that this doesn't hold for `get_version` specifically
    /// - a real H2C's reply had exactly `["command", "module",
    /// "sequence_id"]`, nothing extra at all, so the correlation field
    /// never came back and every `get_version` call just timed out.
    /// `sequence_id` is what's actually proven to round-trip for it
    /// (confirmed earlier, live, on both protocol versions), so that's
    /// what this uses. Whether a custom field survives on some *other*
    /// command (`calibration`, say) is unverified, not confirmed false -
    /// this was only ever checked for `get_version`.
    ///
    /// Not for `print.push_status` - that's not a single reply, see
    /// `PrintStatusCache`. Not for a command with no reply at all either,
    /// like `gcode_line` - use `publish` (fire-and-forget) for those.
    pub async fn request(&self, category: &str, mut body: Value) -> Result<Value, PrinterError> {
        let sequence_id = self.next_sequence_id.fetch_add(1, Ordering::Relaxed).to_string();
        body["sequence_id"] = Value::String(sequence_id.clone());

        let (tx, rx) = oneshot::channel();
        self.pending.lock().await.insert(sequence_id.clone(), tx);

        if let Err(err) = self.publish(category, &body).await {
            self.pending.lock().await.remove(&sequence_id);
            return Err(err);
        }

        match tokio::time::timeout(REQUEST_TIMEOUT, rx).await {
            Ok(Ok(response)) => Ok(response),
            Ok(Err(_)) => Err(PrinterError(format!("{} closed before replying", self.sn))),
            Err(_) => {
                self.pending.lock().await.remove(&sequence_id);
                Err(PrinterError(format!("{} did not reply within {REQUEST_TIMEOUT:?}", self.sn)))
            }
        }
    }

    /// Publishes `body` under `device/{sn}/request`, wrapped as
    /// `{category: body}`, without waiting for a reply - for commands
    /// like `gcode_line` and `pushing.pushall` that don't get a direct
    /// correlated reply (an implicit ack the printer just executes, or an
    /// ongoing stream instead of a single answer, respectively).
    pub async fn publish(&self, category: &str, body: &Value) -> Result<(), PrinterError> {
        self.publish_raw(serde_json::json!({ category: body }).to_string()).await
    }

    /// Publishes an already-serialized payload under `device/{sn}/request`
    /// as-is, with no category wrapping - the low-level primitive both
    /// `publish` above and `GenericPrinter::publish` (a raw pass-through,
    /// no envelope opinion at all) are built on.
    pub async fn publish_raw(&self, payload: String) -> Result<(), PrinterError> {
        let topic = format!("device/{}/request", self.sn);

        match &self.client {
            BambuMqttClient::V1(client) => client
                .publish(topic, QoS::AtLeastOnce, false, payload)
                .await
                .map_err(|err| PrinterError(format!("publish failed: {err}"))),
            BambuMqttClient::V2(client) => client
                .publish(topic, QoSV5::AtLeastOnce, false, payload)
                .await
                .map_err(|err| PrinterError(format!("publish failed: {err}"))),
        }
    }

    /// The current in-memory `print.push_status` object, if at least one
    /// has arrived since connecting - `None` until then. Cloned out from
    /// behind the lock rather than handing out a guard, so callers aren't
    /// holding the poll loop's own lock while they use it.
    pub async fn print_status(&self) -> Option<PrintReport> {
        self.print_status.lock().await.report.clone()
    }
}

/// Runs forever, polling the v3.1.1 eventloop and routing incoming
/// reports (pending-request correlation, `print.push_status` caching).
/// Exits - loudly, not silently like before (that silence is exactly
/// what made a real dropped connection to a live H2C look identical to
/// "the printer just isn't sending anything") - once the connection is
/// gone for good. Dropping the `BambuMqttSession` drops its `AsyncClient`,
/// which is what eventually makes that happen during normal shutdown;
/// this also fires on an unexpected disconnect, which is the case this
/// logging is actually for. No reconnect logic yet - that's the next
/// thing to add here, not something to fake by pretending the loop
/// didn't end.
async fn poll_loop_v1(
    mut event_loop: rumqttc::EventLoop,
    pending: PendingTable,
    print_status: Arc<Mutex<PrintStatusCache>>,
    version: BambuProtocolVersion,
    sn: String,
) {
    loop {
        match event_loop.poll().await {
            Ok(Event::Incoming(Packet::Publish(publish))) => {
                handle_incoming_report(&pending, &print_status, version, &publish.payload).await;
            }
            Ok(_) => {}
            Err(err) => {
                eprintln!("[{sn}] MQTT v3.1.1 connection lost, no longer receiving reports: {err}");
                return;
            }
        }
    }
}

async fn poll_loop_v2(
    mut event_loop: rumqttc::v5::EventLoop,
    pending: PendingTable,
    print_status: Arc<Mutex<PrintStatusCache>>,
    version: BambuProtocolVersion,
    sn: String,
) {
    loop {
        match event_loop.poll().await {
            Ok(rumqttc::v5::Event::Incoming(PacketV5::Publish(publish))) => {
                handle_incoming_report(&pending, &print_status, version, &publish.payload).await;
            }
            Ok(_) => {}
            Err(err) => {
                eprintln!("[{sn}] MQTT v5 connection lost, no longer receiving reports: {err}");
                return;
            }
        }
    }
}

/// Parses an incoming report, unwraps its single category key (a report
/// is `{category: {"sequence_id": "...", ...}}`, same shape a request is
/// published as - `request` publishes `body` and wants `body`-shaped data
/// back, not the wrapper around it), and:
///
/// 1. If its `sequence_id` matches something `request` is waiting on,
///    hands it the unwrapped body - but doesn't wake that waiter yet
///    (see 3 for why).
/// 2. If the unwrapped body is a `print.push_status`, folds it into the
///    running `PrintStatusCache` (which reads the message's own `msg`
///    field to decide full-vs-merge, not anything to do with step 1 - see
///    `PrintStatusCache::apply`'s doc comment for why not).
/// 3. Only now - once the cache is guaranteed up to date - actually sends
///    on whichever oneshot matched in step 1, waking `request`. The order
///    of 2 and 3 is load-bearing, caught live on an earlier version of
///    this mechanism: sending on the oneshot first let a waiter's task
///    resume and check `print_status()` before this function had gotten
///    around to updating it. Nothing here awaits between the cache write
///    and the send, so that race stays closed.
///
/// A report matching nothing pending isn't an error - most Bambu reports
/// (including every organic `push_status` push) aren't a reply to
/// anything specific.
async fn handle_incoming_report(
    pending: &PendingTable,
    print_status: &Arc<Mutex<PrintStatusCache>>,
    version: BambuProtocolVersion,
    payload: &[u8],
) {
    let Ok(value) = serde_json::from_slice::<Value>(payload) else { return };
    let Some(object) = value.as_object() else { return };
    let Some((category, body)) = object.iter().next() else { return };

    let request_sender = match body.get("sequence_id").and_then(Value::as_str) {
        Some(sequence_id) => pending.lock().await.remove(sequence_id),
        None => None,
    };

    if category == "print" && body.get("command").and_then(Value::as_str) == Some("push_status") {
        if let Ok(report) = serde_json::from_value::<PrintReport>(body.clone()) {
            print_status.lock().await.apply(report, version);
        }
    }

    if let Some(sender) = request_sender {
        let _ = sender.send(body.clone());
    }
}