Extrudex/backend/internal/clients/mqtt.go

package clients

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"log/slog"
	"sync"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
)

// ── Bambu Lab telemetry types ───────────────────────────────────────────────

// BambuPrintReport is the JSON payload published by Bambu Lab printers
// on the MQTT report topic. The structure varies by printer model;
// we extract the common fields needed for filament tracking.
type BambuPrintReport struct {
	// Print holds the active print job data.
	Print BambuPrintData `json:"print"`

	// VtTray contains AMS tray info; the extruded length is per-tray.
	VtTray *BambuVtTray `json:"vt_tray,omitempty"`
}

// BambuPrintData carries the active print state from a Bambu report.
type BambuPrintData struct {
	// GcodeFile is the filename being printed.
	GcodeFile string `json:"gcode_file"`
	// GcodeState describes the current print state:
	// "IDLE", "RUNNING", "PAUSE", "FINISH", "FAILED".
	GcodeState string `json:"gcode_state"`
	// McPercent is the progress as a percentage (0-100).
	McPercent int `json:"mc_percent"`
	// McRemainingTime is the estimated remaining time in minutes.
	McRemainingTime int `json:"mc_remaining_time"`
}

// BambuVtTray holds AMS tray telemetry from Bambu printers.
type BambuVtTray struct {
	ID         string `json:"id"`
	TagUID     string `json:"tag_uid"`
	TrayIDName string `json:"tray_id_name"`
	// TrayInfoIdx is the hex color code for the tray's filament.
	TrayInfoIdx string `json:"tray_info_idx"`
	// TrayColor is a hex color string like "FF0000FF".
	TrayColor string `json:"tray_color"`
	// Remain is the percentage of filament remaining on this tray (0-100).
	Remain int `json:"remain"`
	// K is a temperature coefficient.
	K float64 `json:"k"`
	// N is a second temperature coefficient.
	N float64 `json:"n"`
}

// BambuReportHandler is called for each parsed Bambu telemetry message.
type BambuReportHandler func(report BambuPrintReport) error

// ── MQTT client ─────────────────────────────────────────────────────────────

// MQTTClient wraps the Eclipse Paho MQTT client for Bambu Lab printer
// telemetry with optional TLS support.
type MQTTClient struct {
	broker      string
	clientID    string
	topicPrefix string
	tlsCert     string
	tlsKey      string
	handler     BambuReportHandler

	mu     sync.Mutex
	client mqtt.Client
}

// MQTTConfig holds the configuration for creating an MQTTClient.
type MQTTConfig struct {
	Broker      string // e.g., "ssl://192.168.1.50:8883"
	ClientID    string // unique MQTT client id, defaults to "extrudex"
	TopicPrefix string // topic prefix, defaults to "device/+/report"
	TLSCert     string // path to TLS client certificate (optional)
	TLSKey      string // path to TLS client key (optional)
	Handler     BambuReportHandler
}

// NewMQTTClient creates a new MQTTClient. The connection is not established
// until Connect is called.
func NewMQTTClient(cfg MQTTConfig) *MQTTClient {
	if cfg.ClientID == "" {
		cfg.ClientID = "extrudex"
	}
	if cfg.TopicPrefix == "" {
		cfg.TopicPrefix = "device/+/report"
	}
	return &MQTTClient{
		broker:      cfg.Broker,
		clientID:    cfg.ClientID,
		topicPrefix: cfg.TopicPrefix,
		tlsCert:     cfg.TLSCert,
		tlsKey:      cfg.TLSKey,
		handler:     cfg.Handler,
	}
}

// Connect establishes the MQTT connection and subscribes to the configured
// topic prefix. Returns an error if the initial connection fails.
func (c *MQTTClient) Connect() error {
	opts := mqtt.NewClientOptions().
		AddBroker(c.broker).
		SetClientID(c.clientID).
		SetAutoReconnect(true).
		SetMaxReconnectInterval(30 * time.Second).
		SetKeepAlive(30 * time.Second).
		SetPingTimeout(10 * time.Second).
		SetConnectTimeout(15 * time.Second).
		SetOnConnectHandler(func(client mqtt.Client) {
			slog.Info("mqtt: connected", "broker", c.broker)
			// Subscribe on every reconnect.
			token := client.Subscribe(c.topicPrefix, 0, c.messageHandler)
			token.Wait()
			if err := token.Error(); err != nil {
				slog.Error("mqtt: subscribe failed on reconnect", "topic", c.topicPrefix, "error", err)
			} else {
				slog.Info("mqtt: subscribed", "topic", c.topicPrefix)
			}
		}).
		SetConnectionLostHandler(func(client mqtt.Client, err error) {
			slog.Warn("mqtt: connection lost", "error", err)
		})

	// Configure TLS if cert and key are provided.
	if c.tlsCert != "" && c.tlsKey != "" {
		cert, err := tls.LoadX509KeyPair(c.tlsCert, c.tlsKey)
		if err != nil {
			return fmt.Errorf("mqtt: failed to load TLS cert/key: %w", err)
		}
		opts.SetTLSConfig(&tls.Config{
			Certificates: []tls.Certificate{cert},
			MinVersion:   tls.VersionTLS12,
		})
		slog.Info("mqtt: TLS configured", "cert", c.tlsCert)
	}

	c.client = mqtt.NewClient(opts)
	token := c.client.Connect()
	if !token.WaitTimeout(15 * time.Second) {
		return fmt.Errorf("mqtt: connect timed out to %s", c.broker)
	}
	if err := token.Error(); err != nil {
		return fmt.Errorf("mqtt: connect failed: %w", err)
	}

	slog.Info("mqtt: initial connection established", "broker", c.broker)
	return nil
}

// Disconnect gracefully closes the MQTT connection.
func (c *MQTTClient) Disconnect() {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.client != nil && c.client.IsConnected() {
		c.client.Disconnect(2500) // wait up to 2.5s
		slog.Info("mqtt: disconnected")
	}
}

// messageHandler is the MQTT callback invoked for every message received on
// the subscribed topic.
func (c *MQTTClient) messageHandler(_ mqtt.Client, msg mqtt.Message) {
	if c.handler == nil {
		return
	}

	var report BambuPrintReport
	if err := json.Unmarshal(msg.Payload(), &report); err != nil {
		slog.Warn("mqtt: failed to parse bambu report", "topic", msg.Topic(), "error", err)
		return
	}

	if err := c.handler(report); err != nil {
		slog.Error("mqtt: handler error", "topic", msg.Topic(), "error", err)
	}
}