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"`
	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
	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),
		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.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)
			}
			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) {
	if h.db == nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

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

// 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,
		},
	})
}

// 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)
}