MWSE/internal/ws/hub.go

package ws

import (
	"log"
	"sync"

	"git.saqut.com/saqut/mwse/internal/protocol"
)

// Handler processes one inbound message and returns the value to reply with.
// Returning nil is allowed (it becomes JSON null, as `undefined` did in Node).
type Handler func(c *Client, m protocol.Message) any

// Listener is notified of a connection lifecycle event.
type Listener func(c *Client)

// Hub is the engine's shared state: the client and room registries, the message
// router, and the connect/disconnect event bus. Every map is guarded by its own
// RWMutex so unrelated subsystems never contend on one another.
type Hub struct {
	cmu     sync.RWMutex
	clients map[string]*Client

	rmu   sync.RWMutex
	rooms map[string]*Room

	hmu      sync.RWMutex
	handlers map[string]Handler

	lmu          sync.RWMutex
	onConnect    []Listener
	onDisconnect []Listener
}

// NewHub returns an empty hub.
func NewHub() *Hub {
	return &Hub{
		clients:  make(map[string]*Client),
		rooms:    make(map[string]*Room),
		handlers: make(map[string]Handler),
	}
}

// ---- router --------------------------------------------------------------

// Register binds a message type to a handler. Re-registering a type overwrites it.
func (h *Hub) Register(msgType string, handler Handler) {
	h.hmu.Lock()
	h.handlers[msgType] = handler
	h.hmu.Unlock()
}

// HasHandler reports whether a type is registered (used by tests).
func (h *Hub) HasHandler(msgType string) bool {
	h.hmu.RLock()
	_, ok := h.handlers[msgType]
	h.hmu.RUnlock()
	return ok
}

// Handle routes a message to its handler, mirroring the Node MessageRouter
// (MISSING_TYPE / UNKNOWN_TYPE / HANDLER_ERROR), and recovers from handler panics
// so one bad message can never take down the connection or the process.
func (h *Hub) Handle(c *Client, m protocol.Message) (result any) {
	t := m.Type()
	if t == "" {
		return failMsg("MISSING_TYPE")
	}
	h.hmu.RLock()
	handler, ok := h.handlers[t]
	h.hmu.RUnlock()
	if !ok {
		return failMsg("UNKNOWN_TYPE")
	}

	defer func() {
		if r := recover(); r != nil {
			log.Printf("handler panic [%s]: %v", t, r)
			result = map[string]any{"status": "fail", "message": "HANDLER_ERROR"}
		}
	}()
	return handler(c, m)
}

func failMsg(msg string) map[string]any {
	return map[string]any{"status": "fail", "message": msg}
}

// ---- client registry -----------------------------------------------------

func (h *Hub) addClient(c *Client) {
	h.cmu.Lock()
	h.clients[c.ID] = c
	h.cmu.Unlock()
}

func (h *Hub) removeClient(id string) {
	h.cmu.Lock()
	delete(h.clients, id)
	h.cmu.Unlock()
}

// Client looks up a connected client by id.
func (h *Hub) Client(id string) (*Client, bool) {
	h.cmu.RLock()
	c, ok := h.clients[id]
	h.cmu.RUnlock()
	return c, ok
}

// Clients returns a snapshot of all connected clients.
func (h *Hub) Clients() []*Client {
	h.cmu.RLock()
	out := make([]*Client, 0, len(h.clients))
	for _, c := range h.clients {
		out = append(out, c)
	}
	h.cmu.RUnlock()
	return out
}

// ClientCount returns the number of connected clients.
func (h *Hub) ClientCount() int {
	h.cmu.RLock()
	n := len(h.clients)
	h.cmu.RUnlock()
	return n
}

// ---- room registry -------------------------------------------------------

func (h *Hub) addRoom(r *Room) {
	h.rmu.Lock()
	h.rooms[r.ID] = r
	h.rmu.Unlock()
}

func (h *Hub) removeRoom(id string) {
	h.rmu.Lock()
	delete(h.rooms, id)
	h.rmu.Unlock()
}

// Room looks up a room by id.
func (h *Hub) Room(id string) (*Room, bool) {
	h.rmu.RLock()
	r, ok := h.rooms[id]
	h.rmu.RUnlock()
	return r, ok
}

// Rooms returns a snapshot of all rooms.
func (h *Hub) Rooms() []*Room {
	h.rmu.RLock()
	out := make([]*Room, 0, len(h.rooms))
	for _, r := range h.rooms {
		out = append(out, r)
	}
	h.rmu.RUnlock()
	return out
}

// RoomByName returns the first room with the given name. Room names are not
// unique by construction, so this matches the original "first match wins" lookup.
func (h *Hub) RoomByName(name string) (*Room, bool) {
	h.rmu.RLock()
	defer h.rmu.RUnlock()
	for _, r := range h.rooms {
		if r.Name == name {
			return r, true
		}
	}
	return nil, false
}

// ---- event bus -----------------------------------------------------------

// OnConnect registers a listener fired after a client connects.
func (h *Hub) OnConnect(l Listener) {
	h.lmu.Lock()
	h.onConnect = append(h.onConnect, l)
	h.lmu.Unlock()
}

// OnDisconnect registers a listener fired when a client disconnects.
func (h *Hub) OnDisconnect(l Listener) {
	h.lmu.Lock()
	h.onDisconnect = append(h.onDisconnect, l)
	h.lmu.Unlock()
}

func (h *Hub) emitConnect(c *Client) {
	h.lmu.RLock()
	listeners := append([]Listener(nil), h.onConnect...)
	h.lmu.RUnlock()
	for _, l := range listeners {
		l(c)
	}
}

func (h *Hub) emitDisconnect(c *Client) {
	h.lmu.RLock()
	listeners := append([]Listener(nil), h.onDisconnect...)
	h.lmu.RUnlock()
	for _, l := range listeners {
		l(c)
	}
}