ccrm/server/internal/ws/hub.go

package ws

import (
	"encoding/json"
	"log"
	"sync"

	"github.com/gorilla/websocket"
)

// ClientType identifies whether a connection is from a daemon or dashboard.
type ClientType int

const (
	ClientDaemon    ClientType = iota
	ClientDashboard
)

// Client represents a connected WebSocket client.
type Client struct {
	Type    ClientType
	ID      string // machine ID for daemons, user ID for dashboards
	Conn    *websocket.Conn
	Send    chan []byte
	Hub     *Hub
	mu      sync.Mutex
}

// Hub manages all WebSocket connections and routes messages.
type Hub struct {
	mu               sync.RWMutex
	daemonClients    map[string]*Client // keyed by machine ID
	dashboardClients map[*Client]bool
	register         chan *Client
	unregister       chan *Client
	onDaemonMessage  func(machineID string, msgType string, raw json.RawMessage)
	onDashboardMessage func(msgType string, raw json.RawMessage)
}

// NewHub creates a new WebSocket hub.
func NewHub() *Hub {
	return &Hub{
		daemonClients:    make(map[string]*Client),
		dashboardClients: make(map[*Client]bool),
		register:         make(chan *Client),
		unregister:       make(chan *Client),
	}
}

// SetDaemonMessageHandler sets the handler for daemon messages.
func (h *Hub) SetDaemonMessageHandler(handler func(machineID string, msgType string, raw json.RawMessage)) {
	h.onDaemonMessage = handler
}

// SetDashboardMessageHandler sets the handler for dashboard messages.
func (h *Hub) SetDashboardMessageHandler(handler func(msgType string, raw json.RawMessage)) {
	h.onDashboardMessage = handler
}

// Run starts the hub's main loop.
func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			switch client.Type {
			case ClientDaemon:
				h.daemonClients[client.ID] = client
				log.Printf("Daemon connected: %s", client.ID)
			case ClientDashboard:
				h.dashboardClients[client] = true
				log.Printf("Dashboard client connected")
			}
			h.mu.Unlock()

		case client := <-h.unregister:
			h.mu.Lock()
			switch client.Type {
			case ClientDaemon:
				if _, ok := h.daemonClients[client.ID]; ok {
					delete(h.daemonClients, client.ID)
					close(client.Send)
					log.Printf("Daemon disconnected: %s", client.ID)
				}
			case ClientDashboard:
				if _, ok := h.dashboardClients[client]; ok {
					delete(h.dashboardClients, client)
					close(client.Send)
					log.Printf("Dashboard client disconnected")
				}
			}
			h.mu.Unlock()
		}
	}
}

// Register adds a client to the hub.
func (h *Hub) Register(client *Client) {
	h.register <- client
}

// Unregister removes a client from the hub.
func (h *Hub) Unregister(client *Client) {
	h.unregister <- client
}

// SendToDaemon sends a message to a specific daemon.
func (h *Hub) SendToDaemon(machineID string, msg any) {
	data, err := json.Marshal(msg)
	if err != nil {
		log.Printf("failed to marshal message for daemon: %v", err)
		return
	}

	h.mu.RLock()
	client, ok := h.daemonClients[machineID]
	h.mu.RUnlock()

	if !ok {
		log.Printf("daemon %s not connected", machineID)
		return
	}

	select {
	case client.Send <- data:
	default:
		log.Printf("daemon %s send buffer full", machineID)
	}
}

// BroadcastToDashboards sends a message to all connected dashboards.
func (h *Hub) BroadcastToDashboards(msg any) {
	data, err := json.Marshal(msg)
	if err != nil {
		log.Printf("failed to marshal broadcast message: %v", err)
		return
	}

	h.mu.RLock()
	defer h.mu.RUnlock()

	for client := range h.dashboardClients {
		select {
		case client.Send <- data:
		default:
			log.Printf("dashboard client send buffer full, skipping")
		}
	}
}

// HandleMessage processes an incoming message from a client.
func (h *Hub) HandleMessage(client *Client, message []byte) {
	var base Message
	if err := json.Unmarshal(message, &base); err != nil {
		log.Printf("failed to parse message: %v", err)
		return
	}

	switch client.Type {
	case ClientDaemon:
		if h.onDaemonMessage != nil {
			h.onDaemonMessage(client.ID, base.Type, json.RawMessage(message))
		}
	case ClientDashboard:
		if h.onDashboardMessage != nil {
			h.onDashboardMessage(base.Type, json.RawMessage(message))
		}
	}
}

// GetConnectedDaemonIDs returns a list of connected daemon machine IDs.
func (h *Hub) GetConnectedDaemonIDs() []string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	ids := make([]string, 0, len(h.daemonClients))
	for id := range h.daemonClients {
		ids = append(ids, id)
	}
	return ids
}