pekoe/internal/pubsub/hub.go

// @tag pubsub
// @tag core
package pubsub

import (
	"fmt"
	"regexp"
	"slices"
	"strings"
	"sync"
)

// Hub maintains the mapping between channels and their subscribers.
type Hub struct {
	mu       sync.RWMutex
	subs     map[string]map[*Client]struct{}
	patterns map[string]*patternEntry

	// Pattern index for fast channel-to-pattern matching
	// Maps channel name to patterns that would match it
	patternIndex map[string][]string

	// Client registry for CLIENT LIST/INFO/KILL commands
	clients map[uint64]*Client
}

// patternEntry holds compiled regex and subscribers for a pattern.
type patternEntry struct {
	regex   *regexp.Regexp
	clients map[*Client]struct{}
}

// patternCache caches compiled regex patterns to avoid recompilation.
type patternCacheType struct {
	sync.Map
}

var patternCache = patternCacheType{}

// NewHub creates a new pubsub hub instance ready for use.
func NewHub() *Hub {
	return &Hub{
		subs:         make(map[string]map[*Client]struct{}),
		patterns:     make(map[string]*patternEntry),
		patternIndex: make(map[string][]string),
		clients:      make(map[uint64]*Client),
	}
}

// RegisterClient adds a client to the registry
func (h *Hub) RegisterClient(client *Client) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.clients[client.ID()] = client
}

// UnregisterClient removes a client from the registry
func (h *Hub) UnregisterClient(client *Client) {
	h.mu.Lock()
	defer h.mu.Unlock()
	delete(h.clients, client.ID())
}

// GetClient returns a client by ID, or nil if not found
func (h *Hub) GetClient(id uint64) *Client {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.clients[id]
}

// GetAllClients returns a slice of all registered clients
func (h *Hub) GetAllClients() []*Client {
	h.mu.RLock()
	defer h.mu.RUnlock()

	clients := make([]*Client, 0, len(h.clients))
	for _, c := range h.clients {
		clients = append(clients, c)
	}
	return clients
}

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

// Subscribe adds a client to a channel.
func (h *Hub) Subscribe(client *Client, channel string) bool {
	if channel == "" {
		return false
	}

	// Must hold hub mutex when modifying the global subs map
	h.mu.Lock()
	defer h.mu.Unlock()

	clients, exists := h.subs[channel]
	if !exists {
		clients = make(map[*Client]struct{})
		h.subs[channel] = clients

		// New channel - add to pattern index
		h.addChannelToPatternIndex(channel)
	}

	if _, already := clients[client]; already {
		return false
	}

	clients[client] = struct{}{}
	client.AddSubscription(channel)
	return true
}

// addChannelToPatternIndex adds a channel to the pattern index for all matching patterns.
func (h *Hub) addChannelToPatternIndex(channel string) {
	for pattern, entry := range h.patterns {
		if entry.regex.MatchString(channel) {
			h.patternIndex[channel] = append(h.patternIndex[channel], pattern)
		}
	}
}

// Unsubscribe removes a client from a channel.
func (h *Hub) Unsubscribe(client *Client, channel string) bool {
	if channel == "" {
		return false
	}

	// Must hold hub mutex when modifying the global subs map
	h.mu.Lock()
	defer h.mu.Unlock()

	clients, exists := h.subs[channel]
	if !exists {
		return false
	}

	if _, ok := clients[client]; !ok {
		return false
	}

	delete(clients, client)

	if len(clients) == 0 {
		delete(h.subs, channel)
	}

	client.RemoveSubscription(channel)
	return true
}

// UnsubscribeAll removes a client from all subscribed channels and patterns.
func (h *Hub) UnsubscribeAll(client *Client) {
	channels := client.GetChannels()
	for _, channel := range channels {
		h.Unsubscribe(client, channel)
	}

	patterns := client.GetPatterns()
	for _, pattern := range patterns {
		h.PatternUnsubscribe(client, pattern)
	}
}

// Publish sends a message to all subscribers of a channel.
func (h *Hub) Publish(channel string, message []byte) int {
	// Copy clients under lock, then release lock before sending
	h.mu.RLock()
	clients, exists := h.subs[channel]
	if !exists || len(clients) == 0 {
		h.mu.RUnlock()
		return 0
	}

	clientList := make([]*Client, 0, len(clients))
	for client := range clients {
		clientList = append(clientList, client)
	}
	h.mu.RUnlock()

	count := 0
	for _, client := range clientList {
		if err := client.sendMessage(channel, message); err != nil {
			h.removeClientFromChannel(client, channel)
			continue
		}
		count++
	}

	return count
}

// PublishToPattern sends a message to all clients subscribed to patterns matching the channel.
// nolint: cyclop
func (h *Hub) PublishToPattern(channel string, message []byte) int {
	h.mu.RLock()

	// First check the pattern index for patterns that have been indexed
	// Copy the slice to avoid aliasing issues when appending
	matchedPatterns := slices.Clone(h.patternIndex[channel])

	// Also check all patterns to see if any match this channel
	// This ensures we don't miss patterns when publishing to channels with no direct subscribers
	for pattern, entry := range h.patterns {
		if entry.regex.MatchString(channel) {
			// Check if already in matchedPatterns to avoid duplicates
			found := false
			for _, p := range matchedPatterns {
				if p == pattern {
					found = true
					break
				}
			}
			if !found {
				matchedPatterns = append(matchedPatterns, pattern)
			}
		}
	}

	if len(matchedPatterns) == 0 {
		h.mu.RUnlock()
		return 0
	}

	type match struct {
		client  *Client
		pattern string
	}
	var matches []match

	for _, pattern := range matchedPatterns {
		if entry, ok := h.patterns[pattern]; ok {
			for client := range entry.clients {
				matches = append(matches, match{client: client, pattern: pattern})
			}
		}
	}
	h.mu.RUnlock()

	count := 0
	for _, m := range matches {
		if err := m.client.sendPMessage(m.pattern, channel, message); err != nil {
			h.removeClientFromPattern(m.client, m.pattern)
			continue
		}
		count++
	}

	return count
}

// PatternSubscribe adds a client to a pattern subscription.
func (h *Hub) PatternSubscribe(client *Client, pattern string) (bool, error) {
	regex, err := getOrCompilePattern(pattern)
	if err != nil {
		return false, fmt.Errorf("invalid pattern: %v", err)
	}

	h.mu.Lock()
	defer h.mu.Unlock()

	entry, exists := h.patterns[pattern]
	if !exists {
		entry = &patternEntry{
			regex:   regex,
			clients: make(map[*Client]struct{}),
		}
		h.patterns[pattern] = entry
	}

	if _, already := entry.clients[client]; already {
		return false, nil
	}

	entry.clients[client] = struct{}{}

	// Add pattern to index for all existing channels
	for ch := range h.subs {
		if regex.MatchString(ch) {
			h.patternIndex[ch] = append(h.patternIndex[ch], pattern)
		}
	}

	client.AddPatternSubscription(pattern)
	return true, nil
}

// getOrCompilePattern returns a cached compiled regex or compiles and caches a new one.
func getOrCompilePattern(pattern string) (*regexp.Regexp, error) {
	if cached, ok := patternCache.Load(pattern); ok {
		return cached.(*regexp.Regexp), nil
	}

	regex, err := compilePattern(pattern)
	if err != nil {
		return nil, err
	}

	patternCache.Store(pattern, regex)
	return regex, nil
}

func compilePattern(pattern string) (*regexp.Regexp, error) {
	var buf strings.Builder
	for i := 0; i < len(pattern); i++ {
		switch pattern[i] {
		case '*':
			buf.WriteString(".*")
		case '?':
			buf.WriteByte('.')
		case '\\':
			if i+1 < len(pattern) {
				buf.WriteByte('\\')
				buf.WriteByte(pattern[i+1])
				i++
			}
		case '.', '+', '(', ')', '[', ']', '{', '}', '^', '$', '|':
			buf.WriteByte('\\')
			buf.WriteByte(pattern[i])
		default:
			buf.WriteByte(pattern[i])
		}
	}
	return regexp.Compile("^" + buf.String() + "$")
}

// PatternUnsubscribe removes a client from a pattern subscription.
func (h *Hub) PatternUnsubscribe(client *Client, pattern string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()

	entry, exists := h.patterns[pattern]
	if !exists {
		return false
	}

	if _, ok := entry.clients[client]; !ok {
		return false
	}

	delete(entry.clients, client)

	if len(entry.clients) == 0 {
		delete(h.patterns, pattern)
		// Remove pattern from index
		h.removePatternFromIndex(pattern)
	}

	client.RemovePatternSubscription(pattern)
	return true
}

// removePatternFromIndex removes a pattern from all channel indexes.
func (h *Hub) removePatternFromIndex(pattern string) {
	for ch, patterns := range h.patternIndex {
		newPatterns := make([]string, 0, len(patterns)-1)
		for _, p := range patterns {
			if p != pattern {
				newPatterns = append(newPatterns, p)
			}
		}
		if len(newPatterns) == 0 {
			delete(h.patternIndex, ch)
		} else {
			h.patternIndex[ch] = newPatterns
		}
	}
}

// PatternCount returns the total number of pattern subscriptions.
func (h *Hub) PatternCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()

	count := 0
	for _, entry := range h.patterns {
		count += len(entry.clients)
	}
	return count
}

// RemoveClientFromChannel handles cleanup when a client's connection fails.
func (h *Hub) removeClientFromChannel(client *Client, channel string) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if clients, ok := h.subs[channel]; ok {
		if _, exists := clients[client]; exists {
			delete(clients, client)
			if len(clients) == 0 {
				delete(h.subs, channel)
			}
		}
	}
	client.RemoveSubscription(channel)
}

// RemoveClientFromPattern handles cleanup when a client's connection fails for patterns.
func (h *Hub) removeClientFromPattern(client *Client, pattern string) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if entry, ok := h.patterns[pattern]; ok {
		if _, exists := entry.clients[client]; exists {
			delete(entry.clients, client)
			if len(entry.clients) == 0 {
				delete(h.patterns, pattern)
			}
		}
	}
	client.RemovePatternSubscription(pattern)
}

// ChannelCount returns the number of subscribers to a specific channel.
func (h *Hub) ChannelCount(channel string) int {
	h.mu.RLock()
	defer h.mu.RUnlock()

	if clients, ok := h.subs[channel]; ok {
		return len(clients)
	}
	return 0
}

// TotalChannels returns the total number of active channels.
func (h *Hub) TotalChannels() int {
	h.mu.RLock()
	defer h.mu.RUnlock()

	return len(h.subs)
}

// ChannelList returns a list of all active channel names, optionally filtered by pattern.
func (h *Hub) ChannelList(pattern string) []string {
	h.mu.RLock()
	defer h.mu.RUnlock()

	if pattern == "" {
		channels := make([]string, 0, len(h.subs))
		for ch := range h.subs {
			channels = append(channels, ch)
		}
		return channels
	}

	// Compile pattern once outside the loop (using cache)
	regex, err := getOrCompilePattern(pattern)
	if err != nil {
		return nil
	}

	channels := make([]string, 0, len(h.subs))
	for ch := range h.subs {
		if regex.MatchString(ch) {
			channels = append(channels, ch)
		}
	}
	return channels
}