boc/backend/websocket/hub.go

package websocket

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

	"github.com/gorilla/websocket"
	"github.com/rs/zerolog"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true // Allow all origins in development
	},
}

// Client represents a WebSocket client connection
type Client struct {
	hub      *Hub
	conn     *websocket.Conn
	send     chan []byte
	tenantID string
	userID   string
}

// Hub maintains the set of active clients and broadcasts messages
type Hub struct {
	clients    map[*Client]bool
	broadcast  chan []byte
	register   chan *Client
	unregister chan *Client
	logger     zerolog.Logger
	mu         sync.RWMutex
}

// NewHub creates a new WebSocket hub
func NewHub(logger zerolog.Logger) *Hub {
	return &Hub{
		clients:    make(map[*Client]bool),
		broadcast:  make(chan []byte),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		logger:     logger.With().Str("component", "websocket").Logger(),
	}
}

// Run starts the hub's event loop
func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			h.clients[client] = true
			h.mu.Unlock()
			h.logger.Info().Str("tenant", client.tenantID).Msg("client connected")

		case client := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				close(client.send)
			}
			h.mu.Unlock()
			h.logger.Info().Str("tenant", client.tenantID).Msg("client disconnected")

		case message := <-h.broadcast:
			h.mu.RLock()
			for client := range h.clients {
				select {
				case client.send <- message:
				default:
					// Client's send channel is full, close it
					close(client.send)
					delete(h.clients, client)
				}
			}
			h.mu.RUnlock()
		}
	}
}

// HandleWebSocket upgrades HTTP connection to WebSocket
// Requires JWT token in query param ?token=<jwt>
func (h *Hub) HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	// Verify JWT from query parameter
	tokenString := r.URL.Query().Get("token")
	if tokenString == "" {
		h.logger.Warn().Msg("websocket connection rejected: missing token")
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	// TODO: Parse and validate JWT against cfg.JWTSecret
	// For now, reject all unauthenticated connections
	h.logger.Warn().Msg("websocket connection rejected: JWT validation not implemented")
	w.WriteHeader(http.StatusUnauthorized)
	return

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		h.logger.Error().Err(err).Msg("websocket upgrade failed")
		return
	}

	tenantID := r.URL.Query().Get("tenant_id")
	userID := r.URL.Query().Get("user_id")

	if tenantID == "" {
		tenantID = "default"
	}

	client := &Client{
		hub:      h,
		conn:     conn,
		send:     make(chan []byte, 256),
		tenantID: tenantID,
		userID:   userID,
	}

	client.hub.register <- client

	// Start goroutines for reading and writing
	go client.writePump()
	go client.readPump()
}

// Broadcast sends a message to all connected clients
func (h *Hub) Broadcast(message []byte) {
	h.broadcast <- message
}

// BroadcastToTenant sends a message to clients of a specific tenant
func (h *Hub) BroadcastToTenant(tenantID string, message []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	for client := range h.clients {
		if client.tenantID == tenantID {
			select {
			case client.send <- message:
			default:
				// Channel full, skip
			}
		}
	}
}

// readPump handles incoming messages from the client
func (c *Client) readPump() {
	defer func() {
		c.hub.unregister <- c
		c.conn.Close()
	}()

	c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	c.conn.SetPongHandler(func(string) error {
		c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				c.hub.logger.Error().Err(err).Msg("websocket read error")
			}
			break
		}

		// Handle incoming messages (e.g., subscribe to events)
		c.hub.logger.Debug().Str("message", string(message)).Msg("received websocket message")
	}
}

// writePump handles outgoing messages to the client
func (c *Client) writePump() {
	ticker := time.NewTicker(54 * time.Second)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()

	for {
		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			c.conn.WriteMessage(websocket.TextMessage, message)

		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

// Event types for WebSocket messages
type EventType string

const (
	EventDealUpdated      EventType = "deal_updated"
	EventInvoiceCreated   EventType = "invoice_created"
	EventTicketUpdated    EventType = "ticket_updated"
	EventEmployeeUpdated  EventType = "employee_updated"
	EventContractReminder EventType = "contract_reminder"
	EventReportReady      EventType = "report_ready"
	EventWorkflowRun      EventType = "workflow_run"
)

// Event represents a real-time event
type Event struct {
	Type      EventType   `json:"type"`
	TenantID  string      `json:"tenant_id"`
	EntityID  string      `json:"entity_id"`
	Data      interface{} `json:"data"`
	Timestamp time.Time   `json:"timestamp"`
}

// SendEvent sends an event to relevant clients
func (h *Hub) SendEvent(event Event) {
	// In production, filter by tenant and user permissions
	message, _ := json.Marshal(event)
	h.BroadcastToTenant(event.TenantID, message)
}