raft/cli.go

package raft

import (
	"bufio"
	"bytes"
	"fmt"
	"net"
	"os"
	"regexp"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"text/tabwriter"
	"time"
)

const (
	ColorReset  = "\033[0m"
	ColorDim    = "\033[90m" // Dark Gray
	ColorRed    = "\033[31m"
	ColorGreen  = "\033[32m"
	ColorYellow = "\033[33m"
	ColorBlue   = "\033[34m"
	ColorCyan   = "\033[36m"
)

// CLICommand defines a CLI command handler
type CLICommand struct {
	Name        string
	Description string
	Handler     func(parts []string, server *KVServer)
}

// CLI manages the command line interface
type CLI struct {
	server      *KVServer
	commands    map[string]CLICommand
	mu          sync.RWMutex
}

// NewCLI creates a new CLI instance
func NewCLI(server *KVServer) *CLI {
	cli := &CLI{
		server:   server,
		commands: make(map[string]CLICommand),
	}
	cli.registerDefaultCommands()
	return cli
}

// Helper to calculate visible length of string ignoring ANSI codes
var ansiRegex = regexp.MustCompile(`\x1b\[[0-9;]*m`)

func visibleLen(s string) int {
	clean := ansiRegex.ReplaceAllString(s, "")
	return len([]rune(clean))
}

func printBoxed(content string) {
	lines := strings.Split(strings.TrimSpace(content), "\n")
	maxWidth := 0
	for _, line := range lines {
		l := visibleLen(line)
		if l > maxWidth {
			maxWidth = l
		}
	}
	
	// Min width to look decent
	if maxWidth < 20 {
		maxWidth = 20
	}

	// Add padding
	contentWidth := maxWidth + 2 

	fmt.Println() // Start new line before box
	// Top Border
	fmt.Printf("%s╭%s╮%s\n", ColorDim, strings.Repeat("─", contentWidth), ColorReset)
	
	// Content
	for _, line := range lines {
		visLen := visibleLen(line)
		padding := contentWidth - visLen - 1 // -1 for the space after │
		// Ensure padding is not negative
		if padding < 0 { padding = 0 }
		
		fmt.Printf("%s│%s %s%s%s│%s\n", ColorDim, ColorReset, line, strings.Repeat(" ", padding), ColorDim, ColorReset)
	}
	
	// Bottom Border
	fmt.Printf("%s╰%s╯%s\n", ColorDim, strings.Repeat("─", contentWidth), ColorReset)
	fmt.Println() // End new line after box
}

func formatBytes(bytes int64) string {
	const unit = 1024
	if bytes < unit {
		return fmt.Sprintf("%d B", bytes)
	}
	div, exp := int64(unit), 0
	for n := bytes / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
}

// Start starts the command line interface
func (c *CLI) Start() {
	// Boxed startup message
	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("Node %s%s%s CLI Started\n", ColorGreen, c.server.Raft.nodeID, ColorReset))
	sb.WriteString("Type 'help' for commands.")
	printBoxed(sb.String())

	// State Monitor Loop
	go func() {
		var lastState string
		var lastTerm uint64
		stats := c.server.GetStats()
		lastState = stats.State
		lastTerm = stats.Term

		ticker := time.NewTicker(100 * time.Millisecond)
		defer ticker.Stop()

		for {
			select {
			case <-c.server.stopCh:
				return
			case <-ticker.C:
				stats := c.server.GetStats()
				if stats.State != lastState || stats.Term != lastTerm {
					msg := fmt.Sprintf("%s[State Change] %s (Term %d) -> %s (Term %d)%s",
						ColorYellow, lastState, lastTerm, stats.State, stats.Term, ColorReset)
					fmt.Printf("\n%s\n> ", msg)
					
					lastState = stats.State
					lastTerm = stats.Term
				}
			}
		}
	}()

	scanner := bufio.NewScanner(os.Stdin)
	fmt.Print("> ")
	for scanner.Scan() {
		text := strings.TrimSpace(scanner.Text())
		if text == "" {
			fmt.Print("> ")
			continue
		}
		parts := strings.Fields(text)
		cmdName := strings.ToLower(parts[0])

		c.mu.RLock()
		cmd, exists := c.commands[cmdName]
		c.mu.RUnlock()

		if exists {
			cmd.Handler(parts, c.server)
		} else {
			// Check aliases
			switch cmdName {
			case "binlog", "db", "stats":
				c.mu.RLock()
				if infoCmd, ok := c.commands["info"]; ok {
					infoCmd.Handler(parts, c.server)
				}
				c.mu.RUnlock()
			default:
				printBoxed("Unknown command. Type 'help'.")
			}
		}
		fmt.Print("> ")
	}
}

// RegisterCommand registers a new CLI command
func (c *CLI) RegisterCommand(name, description string, handler func(parts []string, server *KVServer)) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.commands[strings.ToLower(name)] = CLICommand{
		Name:        name,
		Description: description,
		Handler:     handler,
	}
}

func (c *CLI) registerDefaultCommands() {
	c.RegisterCommand("set", "Set value (set <key> <val>)", func(parts []string, server *KVServer) {
		if len(parts) != 3 {
			printBoxed("Usage: set <key> <value>")
			return
		}
		key, val := parts[1], parts[2]
		if err := server.Set(key, val); err != nil {
			printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
		} else {
			printBoxed(fmt.Sprintf("%sOK%s", ColorGreen, ColorReset))
		}
	})

	c.RegisterCommand("get", "Get value (get <key>)", func(parts []string, server *KVServer) {
		if len(parts) != 2 {
			printBoxed("Usage: get <key>")
			return
		}
		key := parts[1]
		if val, ok, err := server.GetLinear(key); err != nil {
			printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
		} else if !ok {
			printBoxed(fmt.Sprintf("%sNot Found%s", ColorYellow, ColorReset))
		} else {
			printBoxed(fmt.Sprintf("%s%s%s = %s%s%s", ColorCyan, key, ColorReset, ColorYellow, val, ColorReset))
		}
	})

	c.RegisterCommand("del", "Delete key (del <key>)", func(parts []string, server *KVServer) {
		if len(parts) != 2 {
			printBoxed("Usage: del <key>")
			return
		}
		key := parts[1]
		if err := server.Del(key); err != nil {
			printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
		} else {
			printBoxed(fmt.Sprintf("%sDeleted%s", ColorGreen, ColorReset))
		}
	})

	c.RegisterCommand("join", "Add node to cluster (join <id> <addr>)", func(parts []string, server *KVServer) {
		if len(parts) != 3 {
			printBoxed("Usage: join <nodeID> <address>")
			return
		}
		nodeID := parts[1]
		addr := parts[2]
		if err := server.Join(nodeID, addr); err != nil {
			printBoxed(fmt.Sprintf("%sError joining node:%s %v", ColorRed, ColorReset, err))
		} else {
			printBoxed(fmt.Sprintf("%sJoin request sent for node %s%s", ColorGreen, nodeID, ColorReset))
		}
	})

	c.RegisterCommand("leave", "Remove node from cluster (leave <id>)", func(parts []string, server *KVServer) {
		if len(parts) != 2 {
			printBoxed("Usage: leave <nodeID>")
			return
		}
		nodeID := parts[1]
		if err := server.Leave(nodeID); err != nil {
			printBoxed(fmt.Sprintf("%sError leaving node:%s %v", ColorRed, ColorReset, err))
		} else {
			printBoxed(fmt.Sprintf("%sLeave request sent for node %s%s", ColorGreen, nodeID, ColorReset))
		}
	})

	c.RegisterCommand("search", "Search keys (search <pat> [lim] [off])", func(parts []string, server *KVServer) {
		if len(parts) < 2 {
			printBoxed("Usage: search <pattern> [limit] [offset]")
			return
		}
		pattern := parts[1]
		limit := 20 // Default limit
		offset := 0
		
		if len(parts) >= 3 {
			if l, err := strconv.Atoi(parts[2]); err == nil {
				limit = l
			}
		}
		if len(parts) >= 4 {
			if o, err := strconv.Atoi(parts[3]); err == nil {
				offset = o
			}
		}

		// Construct SQL for DB Engine
		sql := fmt.Sprintf("key like \"%s\" LIMIT %d OFFSET %d", pattern, limit, offset)
		results, err := server.DB.Query(sql)
		if err != nil {
			printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
			return
		}

		// Format Table to Buffer
		var buf bytes.Buffer
		w := tabwriter.NewWriter(&buf, 0, 0, 2, ' ', 0)
		fmt.Fprintln(w, "Key\tValue\tCommitIndex")
		fmt.Fprintln(w, "---\t-----\t-----------")
		for _, r := range results {
			fmt.Fprintf(w, "%s\t%s\t%d\n", r.Key, r.Value, r.CommitIndex)
		}
		w.Flush()
		
		summary := fmt.Sprintf("\n%sFound %d records (showing max %d)%s", ColorDim, len(results), limit, ColorReset)
		printBoxed(buf.String() + summary)
	})

	c.RegisterCommand("count", "Count keys (count <pat>)", func(parts []string, server *KVServer) {
		if len(parts) < 2 {
			printBoxed("Usage: count <pattern> (e.g. 'count *', 'count user.*')")
			return
		}
		pattern := parts[1]
		
		sql := ""
		if pattern == "*" {
			sql = "*"
		} else {
			sql = fmt.Sprintf("key like \"%s\"", pattern)
		}
		
		start := time.Now()
		count, err := server.DB.Count(sql)
		duration := time.Since(start)
		
		if err != nil {
			printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
		} else {
			printBoxed(fmt.Sprintf("%sCount: %d%s (took %v)", ColorGreen, count, ColorReset, duration))
		}
	})

	c.RegisterCommand("info", "Show all system stats", func(parts []string, server *KVServer) {
		// 1. Gather all data
		var m runtime.MemStats
		runtime.ReadMemStats(&m)
		
		dbSize := server.GetDBSize()
		logSize := server.GetLogSize()
		
		stats := server.GetStats()
		health := server.HealthCheck()
		dbApplied := server.DB.GetLastAppliedIndex()

		// 2. Build Output
		var sb strings.Builder
		
		// Section: Node Status
		sb.WriteString(fmt.Sprintf("%sNode Status:%s\n", ColorCyan, ColorReset))
		sb.WriteString(fmt.Sprintf("  ID:           %s%s%s\n", ColorGreen, health.NodeID, ColorReset))
		sb.WriteString(fmt.Sprintf("  State:        %s\n", health.State))
		sb.WriteString(fmt.Sprintf("  Term:         %d\n", health.Term))
		sb.WriteString(fmt.Sprintf("  Leader:       %s\n", health.LeaderID))
		sb.WriteString(fmt.Sprintf("  Healthy:      %v\n", health.IsHealthy))
		sb.WriteString("\n")

		// Section: Storage Stats
		sb.WriteString(fmt.Sprintf("%sStorage & Memory:%s\n", ColorCyan, ColorReset))
		sb.WriteString(fmt.Sprintf("  DB Size:      %s\n", formatBytes(dbSize)))
		sb.WriteString(fmt.Sprintf("  Raft Log:     %s\n", formatBytes(logSize)))
		sb.WriteString(fmt.Sprintf("  Mem Alloc:    %s\n", formatBytes(int64(m.Alloc))))
		sb.WriteString(fmt.Sprintf("  Mem Sys:      %s\n", formatBytes(int64(m.Sys))))
		sb.WriteString(fmt.Sprintf("  NumGC:        %d\n", m.NumGC))
		sb.WriteString("\n")

		// Section: Raft & DB Indices
		sb.WriteString(fmt.Sprintf("%sConsensus Indices:%s\n", ColorCyan, ColorReset))
		sb.WriteString(fmt.Sprintf("  Commit Index: %d\n", stats.CommitIndex))
		sb.WriteString(fmt.Sprintf("  Applied Index:%d\n", stats.LastApplied))
		sb.WriteString(fmt.Sprintf("  Last Log Idx: %d\n", stats.LastLogIndex))
		sb.WriteString(fmt.Sprintf("  DB Applied:   %d\n", dbApplied))
		sb.WriteString("\n")

		// Section: Cluster Members
		sb.WriteString(fmt.Sprintf("%sCluster Members (%d):%s", ColorCyan, stats.ClusterSize, ColorReset))
		
		for id, addr := range stats.ClusterNodes {
			// Determine status color
			statusColor := ColorDim
			nodeID := server.Raft.nodeID // Assuming local nodeID is accessible
			if id == nodeID {
				statusColor = ColorCyan // Self
			} else {
				// Quick check
				conn, err := net.DialTimeout("tcp", addr, 50*time.Millisecond)
				if err == nil {
					statusColor = ColorGreen
					conn.Close()
				}
			}
			
			// Marker for Leader
			marker := ""
			if id == health.LeaderID {
				marker = fmt.Sprintf(" %s(Leader)%s", ColorYellow, ColorReset)
			}
			
			sb.WriteString(fmt.Sprintf("\n  - %s%s%s: %s%s", statusColor, id, ColorReset, addr, marker))
		}
		
		printBoxed(sb.String())
	})

	c.RegisterCommand("quit", "Gracefully shutdown the node", func(parts []string, server *KVServer) {
		printBoxed("Shutting down node...")
		if err := server.Stop(); err != nil {
			printBoxed(fmt.Sprintf("%sError stopping server:%s %v", ColorRed, ColorReset, err))
		} else {
			printBoxed(fmt.Sprintf("%sNode stopped successfully.%s", ColorGreen, ColorReset))
		}
		os.Exit(0)
	})

	c.RegisterCommand("help", "Show commands", func(parts []string, server *KVServer) {
		var sb strings.Builder
		sb.WriteString(fmt.Sprintf("%sCommands:%s\n", ColorGreen, ColorReset))
		
		// Sort commands
		c.mu.RLock()
		var cmds []CLICommand
		for _, cmd := range c.commands {
			cmds = append(cmds, cmd)
		}
		c.mu.RUnlock()
		
		sort.Slice(cmds, func(i, j int) bool {
			return cmds[i].Name < cmds[j].Name
		})

		// Find max length for padding
		maxNameLen := 0
		for _, cmd := range cmds {
			if len(cmd.Name) > maxNameLen {
				maxNameLen = len(cmd.Name)
			}
		}

		for _, cmd := range cmds {
			padding := strings.Repeat(" ", maxNameLen-len(cmd.Name))
			sb.WriteString(fmt.Sprintf("  %s%s%s%s  %s%s%s\n", ColorGreen, cmd.Name, ColorReset, padding, ColorDim, cmd.Description, ColorReset))
		}
		
		printBoxed(sb.String())
	})
}