iloom/shared/pkg/websocket/hub.go

package websocket

import "sync"

type Hub struct {
	clients    map[string]map[*Client]bool // companyID -> clients
	broadcast  chan *Event
	register   chan *Client
	unregister chan *Client
	mu         sync.RWMutex
}

func NewHub() *Hub {
	return &Hub{
		clients:    make(map[string]map[*Client]bool),
		broadcast:  make(chan *Event, 256),
		register:   make(chan *Client),
		unregister: make(chan *Client),
	}
}

func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			if h.clients[client.companyID] == nil {
				h.clients[client.companyID] = make(map[*Client]bool)
			}
			h.clients[client.companyID][client] = true
			h.mu.Unlock()

		case client := <-h.unregister:
			h.mu.Lock()
			if clients, ok := h.clients[client.companyID]; ok {
				if _, exists := clients[client]; exists {
					delete(clients, client)
					close(client.send)
					if len(clients) == 0 {
						delete(h.clients, client.companyID)
					}
				}
			}
			h.mu.Unlock()

		case event := <-h.broadcast:
			h.mu.RLock()
			if event.CompanyID != "" {
				h.broadcastToCompanyLocked(event.CompanyID, event)
			} else {
				for companyID := range h.clients {
					h.broadcastToCompanyLocked(companyID, event)
				}
			}
			h.mu.RUnlock()
		}
	}
}

func (h *Hub) broadcastToCompanyLocked(companyID string, event *Event) {
	data, err := event.Marshal()
	if err != nil {
		return
	}
	for client := range h.clients[companyID] {
		select {
		case client.send <- data:
		default:
			close(client.send)
			delete(h.clients[companyID], client)
		}
	}
}

func (h *Hub) Broadcast(event *Event) {
	h.broadcast <- event
}

func (h *Hub) BroadcastToCompany(companyID string, event *Event) {
	event.CompanyID = companyID
	h.broadcast <- event
}