dumpsterChat/internal/gateway/hub.go

package gateway

import (
	"context"
	"database/sql"
	"encoding/json"
	"log/slog"
	"sync"
	"time"
)

const (
	// idleTimeout is the duration of inactivity before a user is set to idle.
	idleTimeout = 5 * time.Minute
	// idleCheckInterval is how often the hub checks for idle users.
	idleCheckInterval = 30 * time.Second
)

// PresenceData is the data payload for PRESENCE_UPDATE events.
type PresenceData struct {
	UserID   string `json:"user_id"`
	Username string `json:"username"`
	Status   string `json:"status"`
}

// Hub maintains the set of active clients and broadcasts messages to them.
type Hub struct {
	mu           sync.RWMutex
	clients      map[*Client]struct{}
	presence     map[string]string              // userID -> status
	lastActivity map[string]time.Time           // userID -> last activity time
	userServers  map[string]map[string]struct{} // userID -> set of serverIDs
	register     chan *Client
	unregister   chan *Client
	broadcast    chan []byte
	db           *sql.DB
	logger       *slog.Logger
}

// NewHub creates a new Hub.
func NewHub(db *sql.DB, logger *slog.Logger) *Hub {
	return &Hub{
		clients:      make(map[*Client]struct{}),
		presence:     make(map[string]string),
		lastActivity: make(map[string]time.Time),
		userServers:  make(map[string]map[string]struct{}),
		register:     make(chan *Client),
		unregister:   make(chan *Client),
		broadcast:    make(chan []byte, 256),
		db:           db,
		logger:       logger,
	}
}

// Run starts the hub's main event loop. Call this in a goroutine.
func (h *Hub) Run() {
	idleTicker := time.NewTicker(idleCheckInterval)
	defer idleTicker.Stop()

	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			h.clients[client] = struct{}{}
			h.presence[client.UserID] = "online"
			h.lastActivity[client.UserID] = time.Now()
			h.mu.Unlock()

			h.RefreshUserServers(client.UserID)
			h.setUserStatus(client.UserID, "online")
			h.broadcastPresence(client.UserID, "online")
			h.logger.Info("client connected", "user_id", client.UserID)

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

			// Check if user has any remaining connected clients.
			hasOther := h.hasClientForUser(client.UserID)
			if !hasOther {
				delete(h.presence, client.UserID)
				delete(h.lastActivity, client.UserID)
				delete(h.userServers, client.UserID)
			}
			h.mu.Unlock()

			if !hasOther {
				h.setUserStatus(client.UserID, "offline")
				h.broadcastPresence(client.UserID, "offline")
			}
			h.logger.Info("client disconnected", "user_id", client.UserID)

		case message := <-h.broadcast:
			h.mu.RLock()
			for client := range h.clients {
				select {
				case client.send <- message:
				default:
					go func(c *Client) {
						h.unregister <- c
					}(client)
				}
			}
			h.mu.RUnlock()

		case <-idleTicker.C:
			h.checkIdleUsers()
		}
	}
}

// hasClientForUser returns true if any connected client belongs to the given
// user. Caller must hold h.mu (at least for reading).
func (h *Hub) hasClientForUser(userID string) bool {
	for c := range h.clients {
		if c.UserID == userID {
			return true
		}
	}
	return false
}

// RecordActivity updates the last-activity timestamp for a user. If the user
// was idle, it resets their status back to online and broadcasts the change.
// Safe to call from any goroutine.
func (h *Hub) RecordActivity(userID string) {
	h.mu.Lock()
	h.lastActivity[userID] = time.Now()
	wasIdle := h.presence[userID] == "idle"
	if wasIdle {
		h.presence[userID] = "online"
	}
	h.mu.Unlock()

	if wasIdle {
		h.setUserStatus(userID, "online")
		h.broadcastPresence(userID, "online")
	}
}

// checkIdleUsers scans all tracked users and marks as idle any who have been
// inactive for longer than idleTimeout.
func (h *Hub) checkIdleUsers() {
	h.mu.Lock()
	now := time.Now()
	var idleUsers []string
	for userID, last := range h.lastActivity {
		if h.presence[userID] == "online" && now.Sub(last) >= idleTimeout {
			h.presence[userID] = "idle"
			idleUsers = append(idleUsers, userID)
		}
	}
	h.mu.Unlock()

	for _, userID := range idleUsers {
		h.setUserStatus(userID, "idle")
		h.broadcastPresence(userID, "idle")
	}
}

// setUserStatus persists the user's status to the database.
func (h *Hub) setUserStatus(userID, status string) {
	setUserStatus(h.db, h.logger, userID, status)
}

// broadcastPresence sends a PRESENCE_UPDATE event to all connected clients.
func (h *Hub) broadcastPresence(userID, status string) {
	h.BroadcastEvent(Event{
		Type: EventPresenceUpdate,
		Data: PresenceData{
			UserID: userID,
			Status: status,
		},
	})
}

// SetUserStatus is the exported form of setUserStatus, used by non-gateway
// handlers that need to persist a status change.
func (h *Hub) SetUserStatus(userID, status string) {
	h.setUserStatus(userID, status)
}

// BroadcastPresence sends a PRESENCE_UPDATE event including the username.
func (h *Hub) BroadcastPresence(userID, username, status string) {
	h.mu.Lock()
	h.presence[userID] = status
	h.lastActivity[userID] = time.Now()
	h.mu.Unlock()

	h.BroadcastEvent(Event{
		Type: EventPresenceUpdate,
		Data: PresenceData{
			UserID:   userID,
			Username: username,
			Status:   status,
		},
	})
}

// Register queues a client for registration with the hub.
func (h *Hub) Register(client *Client) {
	h.register <- client
}

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

// BroadcastEvent serializes an Event and broadcasts it to all connected clients.
func (h *Hub) BroadcastEvent(event Event) {
	data, err := json.Marshal(event)
	if err != nil {
		h.logger.Error("failed to marshal event", "type", event.Type, "error", err)
		return
	}
	h.broadcast <- data
}

// ClientCount returns the number of connected clients.
func (h *Hub) ClientCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.clients)
}

// RefreshUserServers loads the server membership for a user from the database
// and caches it. Call when a client connects or when server membership changes.
func (h *Hub) RefreshUserServers(userID string) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	rows, err := h.db.QueryContext(ctx, `SELECT server_id FROM members WHERE user_id = $1`, userID)
	if err != nil {
		h.logger.Error("failed to load user servers", "user_id", userID, "error", err)
		return
	}
	defer rows.Close()

	servers := make(map[string]struct{})
	for rows.Next() {
		var sid string
		if err := rows.Scan(&sid); err == nil {
			servers[sid] = struct{}{}
		}
	}

	h.mu.Lock()
	h.userServers[userID] = servers
	h.mu.Unlock()
}

// BroadcastToServer sends an event only to clients who are members of the
// given server. Use this for server-scoped events (messages, channels,
// reactions) to prevent cross-server data leaks.
func (h *Hub) BroadcastToServer(serverID string, event Event) {
	data, err := json.Marshal(event)
	if err != nil {
		h.logger.Error("failed to marshal event", "type", event.Type, "error", err)
		return
	}

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

	for client := range h.clients {
		servers, ok := h.userServers[client.UserID]
		if !ok {
			continue
		}
		if _, member := servers[serverID]; member {
			select {
			case client.send <- data:
			default:
				go func(c *Client) { h.unregister <- c }(client)
			}
		}
	}
}

// ServerIDForChannel looks up the server_id that owns the given channel.
// Used by handlers to route events to the correct server scope.
func (h *Hub) ServerIDForChannel(ctx context.Context, channelID string) (string, error) {
	var serverID string
	err := h.db.QueryRowContext(ctx, `SELECT server_id FROM channels WHERE id = $1`, channelID).Scan(&serverID)
	return serverID, err
}

func setUserStatus(db *sql.DB, logger *slog.Logger, userID, status string) {
	if db == nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	_, err := db.ExecContext(ctx, `UPDATE users SET status = $1 WHERE id = $2`, status, userID)
	if err != nil {
		logger.Error("failed to update user status", "user_id", userID, "status", status, "error", err)
	}
}