Extrudex/backend/internal/clients/moonraker_ws.go

package clients

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// ── WebSocket message types ─────────────────────────────────────────────────

// moonrakerWSMessage is a single JSON-RPC frame from the Moonraker WebSocket.
type moonrakerWSMessage struct {
	JSONRPC string          `json:"jsonrpc"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params"`
	ID      *int            `json:"id"`
}

// MoonrakerPrintEvent is the payload delivered by the "notify_status_update"
// subscription when print_stats or display_status change.
type MoonrakerPrintEvent struct {
	PrintStats    *MoonrakerPrintStats    `json:"print_stats"`
	DisplayStatus *MoonrakerDisplayStatus `json:"display_status"`
}

// MoonrakerDisplayStatus carries progress and the LCD message.
type MoonrakerDisplayStatus struct {
	Progress float64 `json:"progress"`
	Message  string  `json:"message"`
}

// MoonrakerStatusHandler is called for every status update received from the
// Moonraker WebSocket. It receives the parsed event and the raw JSON.
type MoonrakerStatusHandler func(event MoonrakerPrintEvent) error

// ── WebSocket client ────────────────────────────────────────────────────────

// MoonrakerWSClient maintains a persistent WebSocket connection to the
// Moonraker server and delivers parsed status updates to a handler.
type MoonrakerWSClient struct {
	wsURL   string
	handler MoonrakerStatusHandler
	dialer  *websocket.Dialer

	mu     sync.Mutex
	conn   *websocket.Conn
	done   chan struct{}
	once   sync.Once
}

// NewMoonrakerWSClient creates a WebSocket client for the given Moonraker base
// URL. The handler is invoked on every status update.
func NewMoonrakerWSClient(baseURL string, handler MoonrakerStatusHandler) *MoonrakerWSClient {
	baseURL = strings.TrimRight(baseURL, "/")
	wsURL := strings.Replace(baseURL, "http://", "ws://", 1)
	wsURL = strings.Replace(wsURL, "https://", "wss://", 1)
	wsURL += "/websocket"

	return &MoonrakerWSClient{
		wsURL:   wsURL,
		handler: handler,
		dialer: &websocket.Dialer{
			Proxy:            http.ProxyFromEnvironment,
			HandshakeTimeout: 10 * time.Second,
		},
		done: make(chan struct{}),
	}
}

// Connect establishes the WebSocket, subscribes to status updates, and
// starts the read loop in a background goroutine. It retries on failure
// with exponential backoff up to a 60-second cap.
func (c *MoonrakerWSClient) Connect(ctx context.Context) {
	go c.run(ctx)
}

// Shutdown gracefully closes the WebSocket and stops the read loop.
func (c *MoonrakerWSClient) Shutdown() {
	c.once.Do(func() {
		close(c.done)
	})
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}
}

// run is the main connection loop with reconnect backoff.
func (c *MoonrakerWSClient) run(ctx context.Context) {
	backoff := 1 * time.Second
	const maxBackoff = 60 * time.Second

	for {
		select {
		case <-ctx.Done():
			slog.Info("moonraker ws: context cancelled, stopping")
			return
		case <-c.done:
			slog.Info("moonraker ws: shutdown requested")
			return
		default:
		}

		if err := c.connectAndRead(ctx); err != nil {
			slog.Error("moonraker ws: connection error, retrying", "error", err, "backoff", backoff)
		}

		// Exponential backoff.
		select {
		case <-ctx.Done():
			return
		case <-c.done:
			return
		case <-time.After(backoff):
		}
		backoff *= 2
		if backoff > maxBackoff {
			backoff = maxBackoff
		}
	}
}

func (c *MoonrakerWSClient) connectAndRead(ctx context.Context) error {
	slog.Info("moonraker ws: connecting", "url", c.wsURL)

	conn, _, err := c.dialer.DialContext(ctx, c.wsURL, nil)
	if err != nil {
		return fmt.Errorf("dial failed: %w", err)
	}

	c.mu.Lock()
	if c.conn != nil {
		c.conn.Close()
	}
	c.conn = conn
	c.mu.Unlock()

	defer func() {
		c.mu.Lock()
		if c.conn == conn {
			c.conn = nil
		}
		c.mu.Unlock()
		conn.Close()
	}()

	// Subscribe to status updates.
	subReq := map[string]interface{}{
		"jsonrpc": "2.0",
		"method":  "printer.objects.subscribe",
		"params": map[string]interface{}{
			"objects": map[string]interface{}{
				"print_stats":    nil,
				"display_status": nil,
			},
		},
		"id": 1,
	}
	if err := conn.WriteJSON(subReq); err != nil {
		return fmt.Errorf("subscribe failed: %w", err)
	}

	slog.Info("moonraker ws: subscribed to status updates")

	// Set read deadline to detect stale connections.
	// 120s is long enough to avoid false positives.
	pingPeriod := 60 * time.Second

	for {
		// Set read deadline.
		if err := conn.SetReadDeadline(time.Now().Add(150 * time.Second)); err != nil {
			return fmt.Errorf("set read deadline: %w", err)
		}

		_, raw, err := conn.ReadMessage()
		if err != nil {
			return fmt.Errorf("read message: %w", err)
		}

		// Send periodic pings to keep the connection alive.
		go func() {
			time.Sleep(pingPeriod)
			c.mu.Lock()
			if c.conn == conn {
				c.conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(10*time.Second))
			}
			c.mu.Unlock()
		}()

		var msg moonrakerWSMessage
		if err := json.Unmarshal(raw, &msg); err != nil {
			slog.Warn("moonraker ws: failed to parse message", "error", err)
			continue
		}

		// Only process notify_status_update messages.
		if msg.Method != "notify_status_update" {
			continue
		}

		var statusWrapper []MoonrakerPrintEvent
		if err := json.Unmarshal(msg.Params, &statusWrapper); err != nil {
			// Params might be an object, not an array.
			var singleEvent MoonrakerPrintEvent
			if err2 := json.Unmarshal(msg.Params, &singleEvent); err2 != nil {
				slog.Warn("moonraker ws: failed to unmarshal status params", "error", err2)
				continue
			}
			statusWrapper = []MoonrakerPrintEvent{singleEvent}
		}

		for _, ev := range statusWrapper {
			if c.handler != nil {
				if err := c.handler(ev); err != nil {
					slog.Error("moonraker ws: handler error", "error", err)
				}
			}
		}
	}
}