Nukumizu/internal/komari/ws.go

package komari

import (
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"nukumizu-backend/config"
	"nukumizu-backend/internal/node"
	"nukumizu-backend/postLog"

	"github.com/gorilla/websocket"
)

// WSMessage is the structure of messages received from Komari WebSocket.
type WSMessage struct {
	Status string `json:"status"`
	Data   struct {
		Online []string                `json:"online"`
		Data   map[string]node.Report  `json:"data"`
	} `json:"data"`
}

// WSClient manages a WebSocket connection to Komari for real-time status updates.
type WSClient struct {
	komariURL      string
	conn           *websocket.Conn
	connMu         sync.Mutex
	reconnectCount atomic.Int32
	maxRetries     int
	stopCh         chan struct{}
	running        atomic.Bool
	onReconnectFail func() // Callback when all reconnection attempts fail.
}

// NewWSClient creates a new WebSocket client for Komari.
func NewWSClient(komariURL string, maxRetries int) *WSClient {
	return &WSClient{
		komariURL:  komariURL,
		maxRetries: maxRetries,
		stopCh:     make(chan struct{}),
	}
}

// SetOnReconnectFail sets the callback invoked when all reconnection attempts fail.
func (w *WSClient) SetOnReconnectFail(cb func()) {
	w.onReconnectFail = cb
}

// Connect establishes the WebSocket connection and starts the read loop.
func (w *WSClient) Connect() error {
	postLog.Info("Connecting to Komari WebSocket at " + w.komariURL)

	conn, err := w.dial()
	if err != nil {
		postLog.Error("Failed to connect to Komari WebSocket: " + err.Error())
		return err
	}

	w.connMu.Lock()
	w.conn = conn
	w.connMu.Unlock()

	w.reconnectCount.Store(0)
	w.running.Store(true)

	// Request initial state.
	if err := conn.WriteMessage(websocket.TextMessage, []byte("get")); err != nil {
		postLog.Error("Failed to send 'get' message: " + err.Error())
		return err
	}

	postLog.Info("Connected to Komari WebSocket")
	go w.readLoop()

	return nil
}

func (w *WSClient) dial() (*websocket.Conn, error) {
	// Build the WebSocket URL from the Komari HTTP URL.
	wsURL := w.komariURL
	if len(wsURL) > 7 && wsURL[:7] == "http://" {
		wsURL = "ws://" + wsURL[7:]
	} else if len(wsURL) > 8 && wsURL[:8] == "https://" {
		wsURL = "wss://" + wsURL[8:]
	}
	wsURL = wsURL + "/api/clients"

	header := http.Header{}
	dialer := websocket.Dialer{
		HandshakeTimeout: 10 * time.Second,
	}

	conn, _, err := dialer.Dial(wsURL, header)
	return conn, err
}

func (w *WSClient) readLoop() {
	defer func() {
		if r := recover(); r != nil {
			postLog.Error(fmt.Sprintf("WebSocket read loop panic recovered: %v", r))
		}
		w.running.Store(false)
	}()

	for {
		select {
		case <-w.stopCh:
			return
		default:
		}

		w.connMu.Lock()
		conn := w.conn
		w.connMu.Unlock()

		if conn == nil {
			return
		}

		_, data, err := conn.ReadMessage()
		if err != nil {
			postLog.Warning("Komari WebSocket read error: " + err.Error())
			w.tryReconnect()
			return
		}

		var msg WSMessage
		if err := json.Unmarshal(data, &msg); err != nil {
			postLog.Warning("Failed to parse Komari WebSocket message: " + err.Error())
			continue
		}

		if msg.Status != "success" {
			postLog.Warning("Komari WebSocket message with non-success status")
			continue
		}

		// Update the node tracker with the received data.
		tracker := node.GetTracker()
		if tracker != nil {
			tracker.UpdateStatus(msg.Data.Online, msg.Data.Data)
		}
	}
}

func (w *WSClient) tryReconnect() {
	count := w.reconnectCount.Add(1)
	if int(count) > w.maxRetries {
		postLog.Error(fmt.Sprintf("Komari WebSocket reconnection failed after %d attempts", w.maxRetries))
		if w.onReconnectFail != nil {
			w.onReconnectFail()
		}
		// Reset count and start a long-delay retry.
		w.reconnectCount.Store(0)
		time.Sleep(30 * time.Second)
		go w.tryReconnect()
		return
	}

	// Exponential backoff.
	delay := time.Duration(1<<(count-1)) * time.Second
	if delay > 30*time.Second {
		delay = 30 * time.Second
	}

	postLog.Info(fmt.Sprintf("Attempting Komari WebSocket reconnect %d/%d in %v...",
		count, w.maxRetries, delay))

	time.Sleep(delay)

	conn, err := w.dial()
	if err != nil {
		postLog.Warning(fmt.Sprintf("Komari WebSocket reconnect attempt %d failed: %v", count, err))
		go w.tryReconnect()
		return
	}

	w.connMu.Lock()
	if w.conn != nil {
		w.conn.Close()
	}
	w.conn = conn
	w.connMu.Unlock()

	w.reconnectCount.Store(0)

	// Request full state after reconnect.
	if err := conn.WriteMessage(websocket.TextMessage, []byte("get")); err != nil {
		postLog.Error("Failed to send 'get' after reconnect: " + err.Error())
		return
	}

	// Re-fetch node list after reconnect.
	client := GetClient()
	if client != nil {
		nodes, err := client.FetchNodes()
		if err != nil {
			postLog.Error("Failed to refresh node list after reconnect: " + err.Error())
		} else {
			tracker := node.GetTracker()
			if tracker != nil {
				nodeNames := make(map[string]string, len(nodes))
				for _, n := range nodes {
					nodeNames[n.UUID] = n.Name
				}
				tracker.UpdateNodeList(nodeNames)
			}
		}
	}

	postLog.Info("Komari WebSocket reconnected successfully")
	go w.readLoop()
}

// Stop closes the WebSocket connection and stops the read loop.
func (w *WSClient) Stop() {
	w.running.Store(false)
	close(w.stopCh)

	w.connMu.Lock()
	defer w.connMu.Unlock()
	if w.conn != nil {
		w.conn.Close()
		w.conn = nil
	}
}

// IsRunning returns whether the WebSocket client is actively running.
func (w *WSClient) IsRunning() bool {
	return w.running.Load()
}

// --- Global Komari client management ---

var globalClient *Client
var globalWSClient *WSClient
var globalClientMu sync.Mutex

// InitClient initializes the global Komari HTTP client.
func InitClient(komariURL string) {
	globalClientMu.Lock()
	defer globalClientMu.Unlock()
	globalClient = NewClient(komariURL)
}

// GetClient returns the global Komari HTTP client.
func GetClient() *Client {
	return globalClient
}

// InitWSClient initializes the global Komari WebSocket client.
func InitWSClient(komariURL string, maxRetries int) {
	globalClientMu.Lock()
	defer globalClientMu.Unlock()
	globalWSClient = NewWSClient(komariURL, maxRetries)
}

// GetWSClient returns the global Komari WebSocket client.
func GetWSClient() *WSClient {
	return globalWSClient
}

// LoginAndStart performs the Komari login and returns an error if it fails.
func LoginAndStart() error {
	cfg := config.GetConfig()
	client := GetClient()
	if client == nil {
		return fmt.Errorf("komari client not initialized")
	}

	if err := client.Login(cfg.Komari.Account.Username, cfg.Komari.Account.Password); err != nil {
		return fmt.Errorf("komari login failed: %w", err)
	}

	// Fetch initial node list.
	nodes, err := client.FetchNodes()
	if err != nil {
		return fmt.Errorf("failed to fetch nodes from komari: %w", err)
	}

	tracker := node.GetTracker()
	if tracker != nil {
		nodeNames := make(map[string]string, len(nodes))
		for _, n := range nodes {
			nodeNames[n.UUID] = n.Name
		}
		tracker.UpdateNodeList(nodeNames)
	}

	return nil
}