pekoe/internal/server/handlers.go

// @tag server
// @tag handlers:pubsub
package server

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/eleksir/pekoe/internal/glob"
	"github.com/eleksir/pekoe/internal/logger"
	"github.com/eleksir/pekoe/internal/protocol"
	"github.com/eleksir/pekoe/internal/pubsub"
)

// handleSubscribe processes SUBSCRIBE command.
// SUBSCRIBE channel [channel ...] - Subscribe to one or more channels.
// Returns immediately after subscription (doesn't block like PSUBSCRIBE).
func (s *Server) handleSubscribe(client *pubsub.Client, channels [][]byte, writer *ConnWriter) error {
	if len(channels) == 0 {
		writer.WriteError("ERR wrong number of arguments for 'subscribe' command")
		return nil
	}

	// Subscribe to each channel specified
	for _, ch := range channels {
		channel := string(ch)
		if s.hub.Subscribe(client, channel) {
			logger.Debug("client=%s event=subscribe channel=%s", client.RemoteAddr(), channel)
			// Send confirmation in format: ["subscribe", channel, count]
			writer.WriteArray([]protocol.Value{
				protocol.Bulk("subscribe"),
				protocol.Bulk(channel),
				protocol.Int(int64(client.SubsCount())),
			})
		}
	}
	return nil
}

// handleUnsubscribe processes UNSUBSCRIBE command.
// UNSUBSCRIBE [channel [channel ...]] - Unsubscribe from channels.
// If no channels specified, unsubscribe from all.
func (s *Server) handleUnsubscribe(client *pubsub.Client, channels [][]byte, writer *ConnWriter) error {
	if len(channels) == 0 {
		// Unsubscribe from all channels
		chans := client.GetChannels()
		if len(chans) == 0 {
			// No subscriptions - send single unsubscribe message with -1 (Redis behavior)
			writer.WriteArray([]protocol.Value{
				protocol.Bulk("unsubscribe"),
				protocol.Bulk("-1"),
				protocol.Int(0),
			})
			return nil
		}
		for _, ch := range chans {
			s.hub.Unsubscribe(client, ch)
			writer.WriteArray([]protocol.Value{
				protocol.Bulk("unsubscribe"),
				protocol.Bulk(ch),
				protocol.Int(int64(client.SubsCount())),
			})
		}
	} else {
		// Unsubscribe from specified channels - only send response if actually subscribed
		for _, ch := range channels {
			channel := string(ch)
			removed := s.hub.Unsubscribe(client, channel)
			if removed {
				logger.Debug("client=%s event=unsubscribe channel=%s", client.RemoteAddr(), channel)
				writer.WriteArray([]protocol.Value{
					protocol.Bulk("unsubscribe"),
					protocol.Bulk(channel),
					protocol.Int(int64(client.SubsCount())),
				})
			}
		}
	}
	return nil
}

// handlePublish processes PUBLISH command.
// PUBLISH channel message - Publish a message to a channel.
// Returns the number of subscribers who received the message.
func (s *Server) handlePublish(args [][]byte, writer *ConnWriter) error {
	if len(args) != 2 {
		writer.WriteError("ERR wrong number of arguments for 'publish' command")
		return nil
	}

	channel := string(args[0])
	message := args[1]

	regularCount := s.hub.Publish(channel, message)
	patternCount := s.hub.PublishToPattern(channel, message)

	totalCount := regularCount + patternCount
	logger.Debug("event=publish channel=%s regular=%d pattern=%d", channel, regularCount, patternCount)
	writer.WriteInt(int64(totalCount))
	return nil
}

// handlePubsub processes PUBSUB subcommands for introspection.
// nolint: cyclop
func (s *Server) handlePubsub(args [][]byte, writer *ConnWriter) error {
	if len(args) == 0 {
		writer.WriteError("ERR wrong number of arguments for 'pubsub' command")
		return nil
	}

	subcmd := strings.ToLower(string(args[0]))
	switch subcmd {
	case "numsub":
		// PUBSUB NUMSUB [channel ...] - Get subscriber count for channels
		if len(args) < 2 {
			writer.WriteError("ERR wrong number of arguments for 'pubsub numsub' command")
			return nil
		}
		// Build response: [[channel, count], ...]
		channels := args[1:]
		result := make([]protocol.Value, 0, len(channels)*2)
		for _, ch := range channels {
			channel := string(ch)
			count := s.hub.ChannelCount(channel)
			result = append(result, protocol.Bulk(ch), protocol.Int(int64(count)))
		}
		writer.WriteArray(result)

	case "channels":
		// PUBSUB CHANNELS [pattern] - List active channels, optionally filtered
		var pattern string
		if len(args) > 1 {
			pattern = string(args[1])
		}
		channels := s.channelList(pattern)
		result := make([]protocol.Value, 0, len(channels))
		for _, ch := range channels {
			if pattern == "" || glob.Match(pattern, ch) {
				result = append(result, protocol.Bulk(ch))
			}
		}
		writer.WriteArray(result)

	case "shardchannels":
		// PUBSUB SHARDCHANNELS [pattern] - List active sharded channels
		// Pekoe doesn't support sharded pubsub
		writer.WriteError("ERR command not supported in current context")

	case "shardnumsub":
		// PUBSUB SHARDNUMSUB [shardchannel ...] - Get subscriber count for sharded channels
		// Pekoe doesn't support sharded pubsub
		writer.WriteError("ERR command not supported in current context")

	case "numpat":
		// PUBSUB NUMPAT - Count of unique pattern subscriptions
		writer.WriteInt(int64(s.hub.PatternCount()))

	case "help":
		// PUBSUB HELP - Returns helpful text about pubsub subcommands
		helpText := []string{
			"PUBSUB <subcommand> [<argument> [<argument> ...]]>",
			"  PUBSUB CHANNELS [pattern] - List active channels (optionally filtered by pattern)",
			"  PUBSUB SHARDCHANNELS [pattern] - List active sharded channels (not supported)",
			"  PUBSUB NUMSUB [channel ...] - Get subscriber count for channels",
			"  PUBSUB SHARDNUMSUB [shardchannel ...] - Get subscriber count for sharded channels (not supported)",
			"  PUBSUB NUMPAT - Get the number of pattern subscriptions",
			"  PUBSUB HELP - Show this help message",
		}
		result := make([]protocol.Value, len(helpText))
		for i, line := range helpText {
			result[i] = protocol.Bulk(line)
		}
		writer.WriteArray(result)

	default:
		writer.WriteError(fmt.Sprintf("ERR unknown pubsub subcommand '%s'", subcmd))
	}
	return nil
}

// channelList returns all active channel names, optionally filtered by pattern.
func (s *Server) channelList(pattern string) []string {
	channels := s.hub.ChannelList(pattern)
	if channels == nil {
		return []string{}
	}
	return channels
}

// handlePing processes PING command.
// PING [message] - Returns PONG or the optional message.
// Redis returns simple string for message, not bulk string.
func (s *Server) handlePing(args [][]byte, writer *ConnWriter) error {
	if len(args) > 0 {
		writer.WriteStatus(string(args[0]))
	} else {
		writer.WriteStatus("PONG")
	}
	return nil
}

// handleQuit processes QUIT command.
// QUIT - Close the connection gracefully.
func (s *Server) handleQuit(client *pubsub.Client, writer *ConnWriter) error {
	writer.WriteStatus("OK")
	return ErrQuit
}

// handleClient processes CLIENT command.
// CLIENT subcommand [argument [argument ...]]
// nolint: cyclop
func (s *Server) handleClient(client *pubsub.Client, args [][]byte, writer *ConnWriter) error {
	if len(args) == 0 {
		writer.WriteError("ERR wrong number of arguments for 'client' command")
		return nil
	}

	subcmd := strings.ToLower(string(args[0]))
	switch subcmd {
	case "id":
		// CLIENT ID - return the connection ID
		//nolint: gosec // client ID is sequential, overflow is theoretical
		writer.WriteInt(int64(client.ID()))

	case "setname":
		// CLIENT SETNAME connection-name
		if len(args) < 2 {
			writer.WriteError("ERR wrong number of arguments for 'client setname' command")
			return nil
		}
		name := string(args[1])
		// Validate name - must be printable characters and no spaces
		if !isValidClientName(name) {
			writer.WriteError("ERR Client names cannot contain spaces, null characters or names longer than 255 characters")
			return nil
		}
		client.SetName(name)
		writer.WriteStatus("OK")

	case "getname":
		// CLIENT GETNAME - return the connection name
		name := client.GetName()
		if name == "" {
			writer.WriteBulk(nil)
		} else {
			writer.WriteBulk([]byte(name))
		}

	case "list":
		// CLIENT LIST [TYPE normal|master|slave|pubsub] [ID min-id max-id]
		// For pekoe, we only support pubsub clients
		clientList := s.buildClientList(args[1:])
		writer.WriteBulk([]byte(clientList))

	case "info":
		// CLIENT INFO - return information about the client
		writer.WriteBulk([]byte(client.InfoString()))

	case "kill":
		// CLIENT KILL [ID min-id max-id] [ADDR ip:port] [LADDR ip:port] [TYPE normal|master|slave|pubsub] [USER username] [SKIPME yes|no]
		count, err := s.handleClientKill(client, args[1:])
		if err != nil {
			writer.WriteError(err.Error())
			return nil
		}
		writer.WriteInt(int64(count))

	case "pause":
		// CLIENT PAUSE timeout [WRITE|ALL]
		// Acknowledge but don't actually implement - pubsub doesn't need this
		writer.WriteStatus("OK")

	case "unpause":
		// CLIENT UNPAUSE
		writer.WriteStatus("OK")

	case "reply":
		// CLIENT REPLY ON|OFF|SKIP
		if len(args) < 2 {
			writer.WriteError("ERR wrong number of arguments for 'client reply' command")
			return nil
		}
		mode := strings.ToUpper(string(args[1]))
		switch mode {
		case "ON":
			client.SetReplyMode(pubsub.ReplyOn)
			writer.WriteStatus("OK")
		case "OFF":
			client.SetReplyMode(pubsub.ReplyOff)
			// Don't send reply when setting to OFF
			_ = writer.Flush()
		case "SKIP":
			client.SetReplyMode(pubsub.ReplySkip)
			// Send one skip reply, then next will be skipped
			writer.WriteStatus("OK")
		default:
			writer.WriteError("ERR wrong argument for 'client reply' command")
		}

	case "tracking":
		// CLIENT TRACKING ON/OFF [REDIRECT client-id] [PREFIX prefix ...] [BCAST] [OPTIN] [OPTOUT] [NOLOOP]
		// Acknowledge but don't implement - pekoe is pub/sub only
		writer.WriteStatus("OK")

	case "no-evict":
		// CLIENT NO-EVICT ON|OFF
		writer.WriteStatus("OK")

	case "no-touch":
		// CLIENT NO-TOUCH ON|OFF
		writer.WriteStatus("OK")

	case "caching":
		// CLIENT CACHING YES|NO
		writer.WriteStatus("OK")

	default:
		// Unknown client subcommand
		writer.WriteError(fmt.Sprintf("ERR Unknown subcommand or wrong number of arguments for '%s'. Try CLIENT HELP.", subcmd))
	}
	return nil
}

// isValidClientName checks if a client name is valid per Redis spec
func isValidClientName(name string) bool {
	if len(name) == 0 || len(name) > 255 {
		return false
	}
	for _, c := range name {
		if c < 32 || c > 126 || c == ' ' {
			return false
		}
	}
	return true
}

// buildClientList builds CLIENT LIST response for all or filtered clients
func (s *Server) buildClientList(args [][]byte) string {
	clients := s.hub.GetAllClients()

	// Parse optional filter arguments
	var filterType string
	for i := 0; i < len(args); i++ {
		arg := strings.ToLower(string(args[i]))
		if arg == "type" && i+1 < len(args) {
			filterType = strings.ToLower(string(args[i+1]))
			i++
		}
	}

	var result string
	for _, c := range clients {
		// For pekoe, all clients are "pubsub" type
		if filterType != "" && filterType != "pubsub" {
			continue
		}
		result += c.ListString() + "\n"
	}
	return result
}

// handleClientKill handles CLIENT KILL subcommand
// nolint: cyclop
func (s *Server) handleClientKill(currentClient *pubsub.Client, args [][]byte) (int, error) {
	clients := s.hub.GetAllClients()
	killed := 0

	// Parse arguments
	var killID uint64
	var killIDMin uint64
	var killIDMax uint64
	var killAddr string
	var skipMe = true
	var maxAgeMs int64

	for i := 0; i < len(args); i++ {
		arg := strings.ToUpper(string(args[i]))
		switch arg {
		case "ID":
			if i+2 < len(args) {
				// ID min max format - kill all IDs in range
				if _, err := fmt.Sscanf(string(args[i+1]), "%d", &killIDMin); err != nil {
					return 0, fmt.Errorf("invalid client ID: %s", args[i+1])
				}
				if _, err := fmt.Sscanf(string(args[i+2]), "%d", &killIDMax); err != nil {
					return 0, fmt.Errorf("invalid client ID: %s", args[i+2])
				}
				i += 2
			} else if i+1 < len(args) {
				// Single ID
				if _, err := fmt.Sscanf(string(args[i+1]), "%d", &killID); err != nil {
					return 0, fmt.Errorf("invalid client ID: %s", args[i+1])
				}
				i++
			}
		case "ADDR":
			if i+1 < len(args) {
				killAddr = string(args[i+1])
				i++
			}
		case "LADDR":
			// LADDR not fully supported, just skip
			if i+1 < len(args) {
				i++
			}
		case "SKIPME":
			if i+1 < len(args) {
				skipMe = strings.ToUpper(string(args[i+1])) == "YES"
				i++
			}
		case "MAXAGE":
			if i+1 < len(args) {
				if _, err := fmt.Sscanf(string(args[i+1]), "%d", &maxAgeMs); err != nil {
					return 0, fmt.Errorf("invalid maxage: %s", args[i+1])
				}
				if maxAgeMs <= 0 {
					return 0, fmt.Errorf("maxage must be greater than 0")
				}
				i++
			}
		case "TYPE":
			// TYPE filter handled separately in handleClient
			if i+1 < len(args) {
				i++
			}
		}
	}

	now := time.Now()
	for _, c := range clients {
		// Skip current client if requested
		if skipMe && c.ID() == currentClient.ID() {
			continue
		}

		shouldKill := false

		// Check ID (single or range)
		if killID > 0 {
			shouldKill = c.ID() == killID
		} else if killIDMin > 0 || killIDMax > 0 {
			id := c.ID()
			if killIDMin > 0 && id < killIDMin {
				shouldKill = false
			} else if killIDMax > 0 && id > killIDMax {
				shouldKill = false
			} else {
				shouldKill = true
			}
		}

		// Check ADDR
		if !shouldKill && killAddr != "" {
			shouldKill = c.RemoteAddr() == killAddr
		}

		// Check MAXAGE - kill clients that have been connected longer than maxAge
		if shouldKill && maxAgeMs > 0 {
			age := now.Sub(c.CreatedAt()).Milliseconds()
			shouldKill = age <= maxAgeMs
		}

		if shouldKill {
			// Unsubscribe and close
			s.hub.UnsubscribeAll(c)
			s.hub.UnregisterClient(c)
			_ = c.Close()
			killed++
		}
	}

	// For single ID/ADDR that wasn't found, Redis returns error
	if killed == 0 && (killID > 0 || killAddr != "") {
		return 0, fmt.Errorf("No such client")
	}

	return killed, nil
}

// handleSelect processes SELECT command.
// SELECT db - Switch to a different database (0-15).
// For pekoe (pubsub-only), this only tracks the selected DB for client compatibility.
// Redis default max DB is 16 (0-15), configurable via max-databases in config.
func (s *Server) handleSelect(client *pubsub.Client, args [][]byte, writer *ConnWriter) error {
	if len(args) != 1 {
		writer.WriteError("ERR wrong number of arguments for 'select' command")
		return nil
	}

	dbIndexStr := strings.TrimSpace(string(args[0]))
	dbIndex, err := strconv.Atoi(dbIndexStr)
	if err != nil {
		writer.WriteError("ERR invalid DB index")
		return nil
	}

	if dbIndex < 0 || dbIndex >= MaxDatabases {
		writer.WriteError("ERR DB index is out of range")
		return nil
	}

	client.SetSelectedDB(dbIndex)
	writer.WriteStatus("OK")
	return nil
}

// handleInfo processes INFO command - returns server information.
// Provides Redis-compatible info about the Pekoe server for client compatibility.
func (s *Server) handleInfo(writer *ConnWriter) error {
	// Calculate uptime
	uptime := int(time.Since(s.startTime).Seconds())

	info := "# Server\r\n" +
		"redis_version:" + s.version + "\r\n" +
		"redis_mode:standalone\r\n" +
		"os:linux 5.15.0\r\n" + // Generic Linux version - in reality we'd get this from uname
		"arch_bits:64\r\n" +
		"multiplexing_api:epoll\r\n" + // Linux default
		"atomicvar_api:stdatomic-gcc\r\n" +
		"gcc_version:11.2.0\r\n" +
		"process_id:" + strconv.Itoa(os.Getpid()) + "\r\n" +
		"process_supervised:no\r\n" +
		"run_id:" + s.runID + "\r\n" +
		"tcp_port:" + strconv.Itoa(s.cfg.Port) + "\r\n" +
		"uptime_in_seconds:" + strconv.Itoa(uptime) + "\r\n" +
		"uptime_in_days:" + strconv.Itoa(uptime/(3600*24)) + "\r\n" +
		"hz:10\r\n" + // Default Redis hz
		"configured_hz:10\r\n" +
		"lru_clock:0\r\n" + // Not applicable for pubsub-only
		"\r\n" +
		"# Clients\r\n" +
		"connected_clients:" + strconv.Itoa(s.hub.ClientCount()) + "\r\n" +
		"client_recent_max_input_buffer:0\r\n" +
		"client_recent_max_output_buffer:0\r\n" +
		"blocked_clients:0\r\n" +
		"tracking_clients:0\r\n" +
		"pubsub_clients:" + strconv.Itoa(s.hub.ClientCount()) + "\r\n" + // Approximation: all clients are pubsub
		"watching_clients:0\r\n" +
		"\r\n" +
		"# Memory\r\n" +
		"used_memory:0\r\n" +
		"used_memory_human:0B\r\n" +
		"used_memory_rss:0\r\n" +
		"used_memory_peak:0\r\n" +
		"used_memory_lua:0\r\n" + // Not applicable
		"\r\n" +
		"# Persistence\r\n" +
		"loading:0\r\n" +
		"rdb_changes_since_last_save:0\r\n" +
		"rdb_bgsave_in_progress:0\r\n" +
		"rdb_last_save_time:0\r\n" +
		"rdb_last_bgsave_status:ok\r\n" +
		"rdb_last_bgsave_time_sec:0\r\n" +
		"rdb_current_bgsave_time_sec:-1\r\n" +
		"aof_enabled:0\r\n" +
		"aof_rewrite_in_progress:0\r\n" +
		"aof_rewrite_scheduled:0\r\n" +
		"aof_last_rewrite_time_sec:0\r\n" +
		"aof_last_bgrewrite_status:ok\r\n" +
		"aof_last_write_status:ok\r\n" +
		"\r\n" +
		"# Stats\r\n" +
		"total_connections_received:0\r\n" + // Would need to track this
		"total_commands_processed:0\r\n" + // Would need to track this
		"instantaneous_ops_per_sec:0\r\n" +
		"total_net_input_bytes:0\r\n" + // Would need to track
		"total_net_output_bytes:0\r\n" + // Would need to track
		"instantaneous_input_kbps:0\r\n" +
		"instantaneous_output_kbps:0\r\n" +
		"rejected_connections:0\r\n" +
		"sync_full:0\r\n" +
		"sync_partial:0\r\n" +
		"expired_keys:0\r\n" +
		"evicted_keys:0\r\n" +
		"keyspace_hits:0\r\n" +
		"keyspace_misses:0\r\n" +
		"pubsub_channels:" + strconv.Itoa(s.hub.TotalChannels()) + "\r\n" +
		"pubsub_patterns:" + strconv.Itoa(s.hub.PatternCount()) + "\r\n" +
		"latest_fork_usec:0\r\n" +
		"\r\n" +
		"# Replication\r\n" +
		"role:master\r\n" +
		"connected_slaves:0\r\n" +
		"master_failover_state:no-failover\r\n" +
		"master_repl_offset:0\r\n" +
		"repl_backlog_active:0\r\n" +
		"repl_backlog_size:1048576\r\n" +
		"repl_backlog_first_byte_offset:0\r\n" +
		"repl_backlog_histlen:0\r\n" +
		"\r\n" +
		"# CPU\r\n" +
		"used_cpu_sys:0.00\r\n" +
		"used_cpu_user:0.00\r\n" +
		"used_cpu_sys_children:0.00\r\n" +
		"used_cpu_user_children:0.00\r\n" +
		"\r\n" +
		"# Keyspace\r\n"

	writer.WriteBulk([]byte(info))
	return nil
}

// handlePsubscribe processes PSUBSCRIBE command.
// PSUBSCRIBE pattern [pattern ...] - Subscribe to channels matching patterns.
func (s *Server) handlePsubscribe(client *pubsub.Client, patterns [][]byte, writer *ConnWriter) error {
	if len(patterns) == 0 {
		writer.WriteError("ERR wrong number of arguments for 'psubscribe' command")
		return nil
	}

	for _, p := range patterns {
		pattern := string(p)
		ok, err := s.hub.PatternSubscribe(client, pattern)
		if err != nil {
			writer.WriteError(fmt.Sprintf("ERR %v", err))
			return nil
		}
		if ok {
			logger.Debug("client=%s event=psubscribe pattern=%s", client.RemoteAddr(), pattern)
			// Send confirmation in format: ["psubscribe", pattern, count]
			writer.WriteArray([]protocol.Value{
				protocol.Bulk("psubscribe"),
				protocol.Bulk(pattern),
				protocol.Int(int64(client.PatternSubsCount())),
			})
		}
	}
	return nil
}

// handlePunsubscribe processes PUNSUBSCRIBE command.
// PUNSUBSCRIBE [pattern [pattern ...]] - Unsubscribe from patterns.
func (s *Server) handlePunsubscribe(client *pubsub.Client, patterns [][]byte, writer *ConnWriter) error {
	if len(patterns) == 0 {
		// Unsubscribe from all patterns
		pats := client.GetPatterns()
		if len(pats) == 0 {
			// No pattern subscriptions - send single punsubscribe message with -1 (Redis behavior)
			writer.WriteArray([]protocol.Value{
				protocol.Bulk("punsubscribe"),
				protocol.Bulk("-1"),
				protocol.Int(0),
			})
			return nil
		}
		for _, pattern := range pats {
			s.hub.PatternUnsubscribe(client, pattern)
			writer.WriteArray([]protocol.Value{
				protocol.Bulk("punsubscribe"),
				protocol.Bulk(pattern),
				protocol.Int(int64(client.PatternSubsCount())),
			})
		}
	} else {
		// Unsubscribe from specified patterns - only send response if actually subscribed
		for _, p := range patterns {
			pattern := string(p)
			removed := s.hub.PatternUnsubscribe(client, pattern)
			if removed {
				logger.Debug("client=%s event=punsubscribe pattern=%s", client.RemoteAddr(), pattern)
				writer.WriteArray([]protocol.Value{
					protocol.Bulk("punsubscribe"),
					protocol.Bulk(pattern),
					protocol.Int(int64(client.PatternSubsCount())),
				})
			}
		}
	}
	return nil
}