gochat/internal/handler/ws/hub.go

// Package ws (handler/ws) provides the WebSocket Hub that manages connected
// clients, room subscriptions, and message routing. Enhanced with Redis
// Pub/Sub relay for cross-instance broadcasting.
package ws

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	wspkg "github.com/gochat/gochat/internal/ws"
	"github.com/gochat/gochat/pkg/logger"
)

// Client represents a connected WebSocket client.
// Each client has a unique ID, user identity from auth claims,
// and a buffered Send channel for outgoing messages.
type Client struct {
	ID              string            // unique connection ID (uuid)
	UserID          uint              // from WSClaims.UserID
	AccountID       uint              // from WSClaims.AccountID
	Role            string            // from WSClaims.Role
	IsContact       bool              // from WSClaims.IsContact
	PubsubToken     string            // from WSClaims.PubsubToken
	ContactID       uint              // from WSClaims.ContactID (only for contacts)
	InboxID         uint              // from WSClaims.InboxID (only for contacts)
	Conn            *websocket.Conn   // gorilla/websocket connection
	Send            chan []byte       // buffered outgoing message channel (256 capacity)
	Hub             *Hub              // reference back to Hub
	SubscribedRooms map[string]bool   // rooms this client is subscribed to
	CancelPresence  context.CancelFunc // cancel presence refresh on disconnect
}

// NewClient creates a new WebSocket client with the given identity and connection.
func NewClient(userID, accountID uint, conn *websocket.Conn, hub *Hub) *Client {
	return &Client{
		ID:              generateClientID(),
		UserID:          userID,
		AccountID:       accountID,
		Conn:            conn,
		Send:            make(chan []byte, SendChannelSize),
		Hub:             hub,
		SubscribedRooms: make(map[string]bool),
	}
}

// Subscribe adds the client to a room and registers with the Hub.
func (c *Client) Subscribe(room string) {
	c.Hub.mu.Lock()
	c.SubscribedRooms[room] = true
	c.Hub.subscribeClient(c.ID, room)
	c.Hub.mu.Unlock()
}

// Unsubscribe removes the client from a room and unregisters with the Hub.
func (c *Client) Unsubscribe(room string) {
	c.Hub.mu.Lock()
	delete(c.SubscribedRooms, room)
	c.Hub.unsubscribeClient(c.ID, room)
	c.Hub.mu.Unlock()
}

// generateClientID creates a unique client ID using crypto/rand.
func generateClientID() string {
	b := make([]byte, 16)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}

const (
	// MaxMessageSize is the maximum size of a message that can be read from a client.
	MaxMessageSize = 512 * 1024 // 512KB

	// SendChannelSize is the buffer size for client Send channels.
	SendChannelSize = 256

	// WriteWait is the time allowed to write a message to the peer.
	WriteWait = 10 * time.Second

	// PongWait is the time allowed to read the next pong message from the peer.
	PongWait = 60 * time.Second

	// PingPeriod is how often pings are sent to peers. Must be < PongWait.
	PingPeriod = (PongWait * 9) / 10
)

// Hub maintains the set of active clients and broadcasts messages to them.
// Enhanced with Redis Pub/Sub relay for cross-instance message delivery,
// typing indicator support, and agent/contact presence tracking.
type Hub struct {
	// Registered clients indexed by connection ID
	clients map[string]*Client

	// Room subscriptions: room name → set of client IDs in that room
	rooms map[string]map[string]bool

	// Account subscriptions: account ID → set of client IDs
	accounts map[uint]map[string]bool

	// Broadcast relay for cross-instance message delivery via Redis Pub/Sub
	relay *wspkg.BroadcastRelay

	// Typing tracker for typing indicator management
	typing *wspkg.TypingTracker

	// Presence tracker for online/offline status
	presence *wspkg.PresenceTracker

	// Presence lifecycle manager
	presenceMgr *wspkg.PresenceManager

	// Inbound messages from clients (commands)
	commandChan chan *ClientCommand

	mu sync.RWMutex
}

// ClientCommand wraps a command from a client with the client reference.
type ClientCommand struct {
	Client *Client
	Cmd    wspkg.WSCommand
}

// NewHub creates a new Hub with Redis Pub/Sub relay and presence subsystems.
func NewHub(relay *wspkg.BroadcastRelay, typing *wspkg.TypingTracker,
	presence *wspkg.PresenceTracker, presenceMgr *wspkg.PresenceManager) *Hub {
	return &Hub{
		clients:     make(map[string]*Client),
		rooms:       make(map[string]map[string]bool),
		accounts:    make(map[uint]map[string]bool),
		relay:       relay,
		typing:      typing,
		presence:    presence,
		presenceMgr: presenceMgr,
		commandChan: make(chan *ClientCommand, 256),
	}
}

// NewHubSimple creates a minimal Hub without Redis subsystems.
// Used for development/testing when Redis is not available.
func NewHubSimple() *Hub {
	return &Hub{
		clients:     make(map[string]*Client),
		rooms:       make(map[string]map[string]bool),
		accounts:    make(map[uint]map[string]bool),
		commandChan: make(chan *ClientCommand, 256),
	}
}

// Run starts the Hub's main event loop. Processes client registrations,
// unregistrations, and commands. Must be called in a goroutine.
func (h *Hub) Run(ctx context.Context) {
	logger.L().Info("ws hub: starting event loop")

	for {
		select {
		case <-ctx.Done():
			logger.L().Info("ws hub: shutting down event loop")
			h.shutdown()
			return
		case cmd := <-h.commandChan:
			h.processCommand(cmd)
		}
	}
}

// Register adds a client to the Hub and sets up presence.
func (h *Hub) Register(c *Client) {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.clients[c.ID] = c

	// Auto-subscribe to account room on connect
	roomName := accountRoomName(c.AccountID)
	h.subscribeClient(c.ID, roomName)
	c.SubscribedRooms[roomName] = true

	// Set up presence tracking
	if h.presenceMgr != nil {
		if c.IsContact {
			h.presenceMgr.OnContactConnect(context.Background(), c.ContactID, c.AccountID)
			// Also auto-subscribe to pubsub_token room (Chatwoot RoomChannel pattern)
			if c.PubsubToken != "" {
				tokenRoom := pubsubTokenRoomName(c.PubsubToken)
				h.subscribeClient(c.ID, tokenRoom)
				c.SubscribedRooms[tokenRoom] = true
			}
		} else {
			cancelFn := h.presenceMgr.OnAgentConnect(context.Background(), c.UserID, c.AccountID)
			c.CancelPresence = cancelFn
		}
	}

	logger.L().Infof("ws hub: client registered (id=%s, user_id=%d, account_id=%d, is_contact=%v)",
		c.ID, c.UserID, c.AccountID, c.IsContact)

	// Send welcome message
	welcomeMsg, _ := json.Marshal(wspkg.WSMessage{
		Event: wspkg.EventWelcome,
		Data:  map[string]any{"client_id": c.ID},
	})
	c.Send <- welcomeMsg
}

// Unregister removes a client from the Hub and cleans up presence.
func (h *Hub) Unregister(c *Client) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if _, ok := h.clients[c.ID]; !ok {
		return // already unregistered
	}

	// Remove from all subscribed rooms
	for room := range c.SubscribedRooms {
		h.unsubscribeClient(c.ID, room)
	}

	// Remove from Hub's client map
	delete(h.clients, c.ID)

	// Clean up presence tracking
	if h.presenceMgr != nil {
		if c.IsContact {
			h.presenceMgr.OnContactDisconnect(context.Background(), c.ContactID, c.AccountID)
		} else {
			if c.CancelPresence != nil {
				c.CancelPresence() // stop presence refresh loop
			}
			h.presenceMgr.OnAgentDisconnect(context.Background(), c.UserID, c.AccountID)
		}
	}

	// Close Send channel
	close(c.Send)

	logger.L().Infof("ws hub: client unregistered (id=%s, user_id=%d)", c.ID, c.UserID)
}

// SendToAccount sends a message to all clients subscribed to an account room.
func (h *Hub) SendToAccount(accountID uint, data []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	roomName := accountRoomName(accountID)
	if clientIDs, ok := h.rooms[roomName]; ok {
		for clientID := range clientIDs {
			if client, ok := h.clients[clientID]; ok {
				select {
				case client.Send <- data:
				default:
					logger.L().Warnf("ws hub: dropping message for slow client %s", clientID)
				}
			}
		}
	}
}

// SendToAccountConversation sends a message to all clients subscribed to
// a specific conversation within an account.
func (h *Hub) SendToAccountConversation(accountID uint, conversationID uint, data []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	roomName := conversationRoomName(accountID, conversationID)
	if clientIDs, ok := h.rooms[roomName]; ok {
		for clientID := range clientIDs {
			if client, ok := h.clients[clientID]; ok {
				select {
				case client.Send <- data:
				default:
					logger.L().Warnf("ws hub: dropping message for slow client %s", clientID)
				}
			}
		}
	}
}

// SendToRoom sends a message to all clients in a named room.
func (h *Hub) SendToRoom(room string, data []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	if clientIDs, ok := h.rooms[room]; ok {
		for clientID := range clientIDs {
			if client, ok := h.clients[clientID]; ok {
				select {
				case client.Send <- data:
				default:
					logger.L().Warnf("ws hub: dropping message for slow client %s", clientID)
				}
			}
		}
	}
}

// SendToClient sends a message to a specific client by ID.
func (h *Hub) SendToClient(clientID string, data []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	if client, ok := h.clients[clientID]; ok {
		select {
		case client.Send <- data:
		default:
			logger.L().Warnf("ws hub: dropping message for slow client %s", clientID)
		}
	}
}

// SubmitCommand enqueues a client command for the Hub's event loop to process.
func (h *Hub) SubmitCommand(client *Client, cmd wspkg.WSCommand) {
	h.commandChan <- &ClientCommand{Client: client, Cmd: cmd}
}

// processCommand handles a client command in the Hub's event loop.
func (h *Hub) processCommand(cmd *ClientCommand) {
	switch cmd.Cmd.Command {
	case "subscribe":
		h.handleSubscribe(cmd)
	case "unsubscribe":
		h.handleUnsubscribe(cmd)
	case "ping":
		h.handlePing(cmd)
	case "typing_on":
		h.handleTypingOn(cmd)
	case "typing_off":
		h.handleTypingOff(cmd)
	case "update_presence":
		h.handleUpdatePresence(cmd)
	default:
		logger.L().Warnf("ws hub: unknown command '%s' from client %s", cmd.Cmd.Command, cmd.Client.ID)
	}
}

// handleSubscribe processes a subscribe command.
// Mirrors Chatwoot's RoomChannel.subscribe flow.
func (h *Hub) handleSubscribe(cmd *ClientCommand) {
	var data wspkg.SubscribeData
	if err := json.Unmarshal([]byte(cmd.Cmd.Data), &data); err != nil {
		logger.L().Warnf("ws hub: invalid subscribe data from client %s: %v", cmd.Client.ID, err)
		return
	}

	// Verify account access
	if data.AccountID != cmd.Client.AccountID {
		rejectMsg, _ := json.Marshal(wspkg.WSMessage{
			Event: wspkg.EventSubscribeReject,
			Data:  map[string]any{"reason": "account_id mismatch"},
		})
		cmd.Client.Send <- rejectMsg
		return
	}

	var roomName string
	switch data.Channel {
	case wspkg.ChannelAccount:
		roomName = accountRoomName(data.AccountID)
	case wspkg.ChannelConversation:
		if data.ConversationID == 0 {
			rejectMsg, _ := json.Marshal(wspkg.WSMessage{
				Event: wspkg.EventSubscribeReject,
				Data:  map[string]any{"reason": "conversation_id required for ConversationChannel"},
			})
			cmd.Client.Send <- rejectMsg
			return
		}
		roomName = conversationRoomName(data.AccountID, data.ConversationID)
	default:
		rejectMsg, _ := json.Marshal(wspkg.WSMessage{
			Event: wspkg.EventSubscribeReject,
			Data:  map[string]any{"reason": "unknown channel type"},
		})
		cmd.Client.Send <- rejectMsg
		return
	}

	h.mu.Lock()
	h.subscribeClient(cmd.Client.ID, roomName)
	cmd.Client.SubscribedRooms[roomName] = true
	h.mu.Unlock()

	confirmMsg, _ := json.Marshal(wspkg.WSMessage{
		Event: wspkg.EventSubscribeConfirm,
		Data:  map[string]any{"room": roomName, "channel": data.Channel},
	})
	cmd.Client.Send <- confirmMsg

	logger.L().Infof("ws hub: client %s subscribed to room %s", cmd.Client.ID, roomName)
}

// handleUnsubscribe processes an unsubscribe command.
func (h *Hub) handleUnsubscribe(cmd *ClientCommand) {
	var data wspkg.SubscribeData
	if err := json.Unmarshal([]byte(cmd.Cmd.Data), &data); err != nil {
		logger.L().Warnf("ws hub: invalid unsubscribe data from client %s: %v", cmd.Client.ID, err)
		return
	}

	var roomName string
	switch data.Channel {
	case wspkg.ChannelAccount:
		roomName = accountRoomName(data.AccountID)
	case wspkg.ChannelConversation:
		roomName = conversationRoomName(data.AccountID, data.ConversationID)
	default:
		return
	}

	h.mu.Lock()
	h.unsubscribeClient(cmd.Client.ID, roomName)
	delete(cmd.Client.SubscribedRooms, roomName)
	h.mu.Unlock()

	confirmMsg, _ := json.Marshal(wspkg.WSMessage{
		Event: wspkg.EventUnsubscribeConfirm,
		Data:  map[string]any{"room": roomName},
	})
	cmd.Client.Send <- confirmMsg

	logger.L().Infof("ws hub: client %s unsubscribed from room %s", cmd.Client.ID, roomName)
}

// handlePing responds to a ping command from the client.
func (h *Hub) handlePing(cmd *ClientCommand) {
	pingMsg, _ := json.Marshal(wspkg.WSMessage{
		Event: wspkg.EventPingResponse,
		Data:  map[string]any{"timestamp": time.Now().UnixMilli()},
	})
	cmd.Client.Send <- pingMsg
}

// handleTypingOn processes a typing_on command.
func (h *Hub) handleTypingOn(cmd *ClientCommand) {
	if h.typing == nil {
		return
	}

	var data wspkg.TypingData
	if err := json.Unmarshal([]byte(cmd.Cmd.Data), &data); err != nil {
		logger.L().Warnf("ws hub: invalid typing_on data from client %s: %v", cmd.Client.ID, err)
		return
	}

	performerType := "user"
	if cmd.Client.IsContact {
		performerType = "contact"
	}

	performer := &wspkg.Performer{
		ID:   cmd.Client.UserID,
		Type: performerType,
	}

	if err := h.typing.SetTypingOn(context.Background(), data.AccountID, data.ConversationID, performer); err != nil {
		logger.L().Warnf("ws hub: typing_on failed: %v", err)
	}
}

// handleTypingOff processes a typing_off command.
func (h *Hub) handleTypingOff(cmd *ClientCommand) {
	if h.typing == nil {
		return
	}

	var data wspkg.TypingData
	if err := json.Unmarshal([]byte(cmd.Cmd.Data), &data); err != nil {
		logger.L().Warnf("ws hub: invalid typing_off data from client %s: %v", cmd.Client.ID, err)
		return
	}

	performerType := "user"
	if cmd.Client.IsContact {
		performerType = "contact"
	}

	performer := &wspkg.Performer{
		ID:   cmd.Client.UserID,
		Type: performerType,
	}

	if err := h.typing.SetTypingOff(context.Background(), data.AccountID, data.ConversationID, performer); err != nil {
		logger.L().Warnf("ws hub: typing_off failed: %v", err)
	}
}

// handleUpdatePresence processes an update_presence command.
func (h *Hub) handleUpdatePresence(cmd *ClientCommand) {
	if h.presence == nil {
		return
	}

	var data wspkg.PresenceData
	if err := json.Unmarshal([]byte(cmd.Cmd.Data), &data); err != nil {
		logger.L().Warnf("ws hub: invalid update_presence data from client %s: %v", cmd.Client.ID, err)
		return
	}

	ctx := context.Background()
	switch data.Status {
	case "online":
		h.presence.SetAgentOnline(ctx, cmd.Client.UserID, cmd.Client.AccountID)
	case "busy":
		h.presence.SetAgentBusy(ctx, cmd.Client.UserID, cmd.Client.AccountID)
	case "offline":
		h.presence.SetAgentOffline(ctx, cmd.Client.UserID, cmd.Client.AccountID)
	default:
		logger.L().Warnf("ws hub: unknown presence status '%s' from client %s", data.Status, cmd.Client.ID)
	}
}

// --- Internal helpers ---

// subscribeClient adds a client ID to a room's member set.
// Must be called with h.mu held.
func (h *Hub) subscribeClient(clientID string, room string) {
	if _, ok := h.rooms[room]; !ok {
		h.rooms[room] = make(map[string]bool)
	}
	h.rooms[room][clientID] = true
}

// unsubscribeClient removes a client ID from a room's member set.
// Must be called with h.mu held.
func (h *Hub) unsubscribeClient(clientID string, room string) {
	if clients, ok := h.rooms[room]; ok {
		delete(clients, clientID)
		if len(clients) == 0 {
			delete(h.rooms, room)
		}
	}
}

// shutdown disconnects all clients and cleans up.
func (h *Hub) shutdown() {
	h.mu.Lock()
	defer h.mu.Unlock()

	for id, client := range h.clients {
		if client.CancelPresence != nil {
			client.CancelPresence()
		}
		close(client.Send)
		client.Conn.Close()
		delete(h.clients, id)
	}

	h.rooms = make(map[string]map[string]bool)
	h.accounts = make(map[uint]map[string]bool)

	logger.L().Info("ws hub: all clients disconnected, shutdown complete")
}

// Shutdown gracefully closes the Hub, disconnecting all WebSocket clients.
// This is the public API called during application graceful shutdown.
// It sends a close message to each client and waits for them to finish,
// respecting the context deadline.
func (h *Hub) Shutdown(ctx context.Context) {
	logger.L().Info("ws hub: initiating graceful shutdown")
	h.shutdown()
}

// Room name helpers following Chatwoot naming conventions.

func accountRoomName(accountID uint) string {
	return fmt.Sprintf("account_%d", accountID)
}

func conversationRoomName(accountID uint, conversationID uint) string {
	return fmt.Sprintf("account_%d_conversation_%d", accountID, conversationID)
}

func pubsubTokenRoomName(token string) string {
	return fmt.Sprintf("pubsub_token_%s", token)
}