pekoe/internal/pubsub/client.go

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

import (
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/eleksir/pekoe/internal/protocol"
)

// ReplyMode controls client reply behavior
type ReplyMode int

const (
	ReplyOn   ReplyMode = iota // Normal reply mode - send all responses
	ReplyOff                   // Reply off - suppress all responses
	ReplySkip                  // Reply skip - skip the next response only
)

// Global client ID counter - each client gets a unique ID
var clientIDCounter atomic.Uint64

// Client represents a subscriber connection to the pubsub system.
// It maintains the connection (net.Conn) and tracks which channels/patterns the client is subscribed to.
type Client struct {
	id         uint64 // Unique client ID
	conn       net.Conn
	name       string    // Connection name set via CLIENT SETNAME
	replyMode  ReplyMode // Reply mode (ON/OFF/SKIP)
	flags      uint64    // Client flags (e.g., paused, no-touch)
	created    time.Time // When the client connected
	lastActive time.Time // Last activity time
	subscribed map[string]struct{}
	patterns   map[string]struct{}
	mu         sync.RWMutex
	writeMu    sync.Mutex // Protects concurrent writes to conn
}

// NewClient creates a new pubsub client wrapping a connection.
func NewClient(conn net.Conn) *Client {
	now := time.Now()
	return &Client{
		id:         clientIDCounter.Add(1),
		conn:       conn,
		name:       "",
		replyMode:  ReplyOn,
		flags:      0,
		created:    now,
		lastActive: now,
		subscribed: make(map[string]struct{}),
		patterns:   make(map[string]struct{}),
	}
}

// ID returns the unique client ID
func (c *Client) ID() uint64 {
	return c.id
}

// SetName sets the connection name
func (c *Client) SetName(name string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.name = name
}

// GetName returns the connection name
func (c *Client) GetName() string {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.name
}

// SetReplyMode sets the reply mode
func (c *Client) SetReplyMode(mode ReplyMode) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.replyMode = mode
}

// GetReplyMode returns the current reply mode
func (c *Client) GetReplyMode() ReplyMode {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.replyMode
}

// UpdateActivity updates the last activity timestamp
func (c *Client) UpdateActivity() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.lastActive = time.Now()
}

// GetChannels returns a copy of the channels this client is subscribed to.
// Returns a slice to prevent external modification.
func (c *Client) GetChannels() []string {
	c.mu.RLock()
	defer c.mu.RUnlock()

	channels := make([]string, 0, len(c.subscribed))
	for ch := range c.subscribed {
		channels = append(channels, ch)
	}
	return channels
}

// GetPatterns returns a copy of the patterns this client is subscribed to.
func (c *Client) GetPatterns() []string {
	c.mu.RLock()
	defer c.mu.RUnlock()

	patterns := make([]string, 0, len(c.patterns))
	for p := range c.patterns {
		patterns = append(patterns, p)
	}
	return patterns
}

// SubsCount returns the number of channels this client is subscribed to.
func (c *Client) SubsCount() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.subscribed)
}

// PatternSubsCount returns the number of patterns this client is subscribed to.
func (c *Client) PatternSubsCount() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.patterns)
}

// AddSubscription tracks that this client has subscribed to a channel.
func (c *Client) AddSubscription(channel string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.subscribed[channel] = struct{}{}
	c.lastActive = time.Now()
}

// RemoveSubscription removes a channel from the client's subscriptions.
func (c *Client) RemoveSubscription(channel string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.subscribed, channel)
	c.lastActive = time.Now()
}

// AddPatternSubscription tracks that this client has subscribed to a pattern.
func (c *Client) AddPatternSubscription(pattern string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.patterns[pattern] = struct{}{}
	c.lastActive = time.Now()
}

// RemovePatternSubscription removes a pattern from the client's subscriptions.
func (c *Client) RemovePatternSubscription(pattern string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.patterns, pattern)
	c.lastActive = time.Now()
}

// SendMessage sends a pubsub message to the client in RESP format.
// Format: ["message", channel, payload]
// Returns error if the write fails - caller will handle cleanup.
func (c *Client) sendMessage(channel string, message []byte) error {
	response := protocol.MakeMultiBulk([]protocol.Value{
		protocol.Bulk("message"),
		protocol.Bulk(channel),
		protocol.Bulk(message),
	})

	c.writeMu.Lock()
	_, err := c.conn.Write(response)
	c.writeMu.Unlock()
	c.UpdateActivity()
	return err
}

// SendPMessage sends a pubsub message matching a pattern.
// Format: ["pmessage", pattern, channel, payload]
func (c *Client) sendPMessage(pattern string, channel string, message []byte) error {
	response := protocol.MakeMultiBulk([]protocol.Value{
		protocol.Bulk("pmessage"),
		protocol.Bulk(pattern),
		protocol.Bulk(channel),
		protocol.Bulk(message),
	})

	c.writeMu.Lock()
	_, err := c.conn.Write(response)
	c.writeMu.Unlock()
	c.UpdateActivity()
	return err
}

// RemoteAddr returns the remote address of the client for logging.
func (c *Client) RemoteAddr() string {
	return c.conn.RemoteAddr().String()
}

// LocalAddr returns the local address of the client connection.
func (c *Client) LocalAddr() string {
	return c.conn.LocalAddr().String()
}

// InfoString returns CLIENT INFO formatted string (Redis format)
// Format: id= addr= laddr= fd= age= idle= flags= db= sub= psub= multi= qbuf= qbuf-free= obl= oll= omem= tot-mem= cmd= user= redir=
func (c *Client) InfoString() string {
	c.mu.RLock()
	defer c.mu.RUnlock()

	age := int(time.Since(c.created).Seconds())
	idle := int(time.Since(c.lastActive).Seconds())

	// Determine flags
	flags := ""
	if c.replyMode == ReplyOff {
		flags += "O"
	}
	if c.replyMode == ReplySkip {
		flags += "S"
	}

	// Format: Redis CLIENT INFO format
	return fmt.Sprintf("id=%d addr=%s laddr=%s fd=-1 age=%d idle=%d flags=%s db=0 sub=%d psub=%d multi=-1 qbuf=0 qbuf-free=0 obl=0 oll=0 omem=0 tot-mem=0 cmd=client user=default redir=-1",
		c.id,
		c.RemoteAddr(),
		c.LocalAddr(),
		age,
		idle,
		flags,
		len(c.subscribed),
		len(c.patterns),
	)
}

// ListString returns CLIENT LIST formatted string
// Format: id= addr= ... (same as INFO but one line per client)
func (c *Client) ListString() string {
	return c.InfoString()
}

// Close closes the client's connection.
// Called when the client disconnects or on error.
func (c *Client) Close() error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	return c.conn.Close()
}