raft/cli.go

package raft

import (
	"bufio"
	"bytes"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"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
	token       string
}

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

// scanCRLF is a split function for a Scanner that returns each line of
// text, stripped of any trailing end-of-line marker. The returned line may
// be empty. The end-of-line marker is one optional carriage return followed
// by one mandatory newline, OR just a carriage return.
func scanCRLF(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	if i := bytes.IndexByte(data, '\n'); i >= 0 {
		// We have a full newline-terminated line.
		return i + 1, dropCR(data[0:i]), nil
	}
	if i := bytes.IndexByte(data, '\r'); i >= 0 {
		// We have a CR-terminated line.
		return i + 1, data[0:i], nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), dropCR(data), nil
	}
	// Request more data.
	return 0, nil, nil
}

// dropCR drops a terminal \r from the data.
func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

// Start starts the command line interface
func (c *CLI) Start() {
	// Attempt to fix terminal settings if they are messed up (e.g. missing ICRNL)
	// This helps when Enter sends \r but the terminal driver expects \n
	cmd := exec.Command("stty", "sane")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	_ = cmd.Run()

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

	// Initialize history handling
	historyPath := filepath.Join(os.TempDir(), fmt.Sprintf(".raft_history_%s", c.server.Raft.nodeID))
	historyFile, err := os.OpenFile(historyPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		// Just warn, don't fail
		fmt.Printf("%sWarning: Could not open history file: %v%s\n", ColorYellow, err, ColorReset)
	}
	defer func() {
		if historyFile != nil {
			historyFile.Close()
		}
	}()

	// 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)
	scanner.Split(scanCRLF)
	fmt.Print("> ")
	for scanner.Scan() {
		text := strings.TrimSpace(scanner.Text())
		if text == "" {
			fmt.Print("> ")
			continue
		}

		// Save command to history
		if historyFile != nil {
			if _, err := historyFile.WriteString(text + "\n"); err != nil {
				// Ignore write error
			}
		}

		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()
			case "history":
				c.printHistory(historyPath)
			default:
				printBoxed("Unknown command. Type 'help'.")
			}
		}
		fmt.Print("> ")
	}
}

// printHistory prints the command history
func (c *CLI) printHistory(path string) {
	content, err := os.ReadFile(path)
	if err != nil {
		printBoxed("No history found.")
		return
	}
	
	lines := strings.Split(string(content), "\n")
	// Show last 20 lines
	start := 0
	if len(lines) > 20 {
		start = len(lines) - 20
	}
	
	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("%sCommand History (Last 20):%s\n", ColorCyan, ColorReset))
	for i := start; i < len(lines); i++ {
		line := strings.TrimSpace(lines[i])
		if line != "" {
			sb.WriteString(fmt.Sprintf("%d: %s\n", i+1, line))
		}
	}
	printBoxed(sb.String())
}

// 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]
		
		// Use Authenticated Set
		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()

		if err := server.SetAuthenticated(key, val, token); 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]
		
		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()

		if val, ok, err := server.GetLinearAuthenticated(key, token); 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]
		
		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()

		if err := server.DelAuthenticated(key, token); err != nil {
			printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
		} else {
			printBoxed(fmt.Sprintf("%sDeleted%s", ColorGreen, ColorReset))
		}
	})

	// User Command Group
	c.RegisterCommand("user", "User management commands (type 'user help' for more)", func(parts []string, server *KVServer) {
		if len(parts) < 2 {
			printBoxed("Usage: user <subcommand> [args...]\nType 'user help' for list of subcommands.")
			return
		}
		
		subCmd := strings.ToLower(parts[1])
		args := parts[2:]
		
		// Common vars
		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()

		switch subCmd {
		case "init": // Was auth-init
			if server.AuthManager.IsEnabled() && !server.IsRoot(token) {
				printBoxed(fmt.Sprintf("%sPermission Denied: Auth already enabled. Login as root to re-init.%s", ColorRed, ColorReset))
				return
			}
			fmt.Print("Enter root password: ")
			password, err := readPassword()
			fmt.Println()
			if err != nil || password == "" {
				printBoxed("Error reading password")
				return
			}
			if err := server.AuthManager.CreateRootUser(password); err != nil {
				printBoxed(fmt.Sprintf("Failed to create root user: %v", err))
				return
			}
			if err := server.AuthManager.EnableAuth(); err != nil {
				printBoxed(fmt.Sprintf("Failed to enable auth: %v", err))
				return
			}
			printBoxed("Auth Initialized. Please login as 'root'.")

		case "login": // Was login
			if len(args) < 1 {
				printBoxed("Usage: user login <username> [code]")
				return
			}
			username := args[0]
			code := ""
			if len(args) > 1 {
				code = args[1]
			}
			fmt.Print("Password: ")
			password, err := readPassword()
			fmt.Println()
			if err != nil {
				printBoxed("Error reading password")
				return
			}
			t, err := server.AuthManager.Login(username, password, code, "cli")
			if err != nil {
				printBoxed(fmt.Sprintf("%sLogin Failed:%s %v", ColorRed, ColorReset, err))
			} else {
				c.mu.Lock()
				c.token = t
				c.mu.Unlock()
				printBoxed(fmt.Sprintf("%sLogin Success!%s", ColorGreen, ColorReset))
			}

		case "logout": // Was logout
			if token != "" {
				server.Logout(token)
				c.mu.Lock()
				c.token = ""
				c.mu.Unlock()
				printBoxed("Logged out.")
			} else {
				printBoxed("Not logged in.")
			}

		case "who": // Was whoami
			if token == "" {
				printBoxed("Not logged in.")
				return
			}
			session, err := server.GetSessionInfo(token)
			if err != nil {
				printBoxed(fmt.Sprintf("Error: %v", err))
				return
			}
			printBoxed(fmt.Sprintf("Logged in as: %s%s%s", ColorCyan, session.Username, ColorReset))

		case "add": // Was user-add
			if len(args) < 1 {
				printBoxed("Usage: user add <username> [role1,role2...]")
				return
			}
			username := args[0]
			var roles []string
			if len(args) > 1 {
				roles = strings.Split(args[1], ",")
			}
			fmt.Printf("Enter password for %s: ", username)
			password, err := readPassword()
			fmt.Println()
			if err != nil || password == "" {
				printBoxed("Error reading password")
				return
			}
			if err := server.CreateUser(username, password, roles, token); err != nil {
				printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
			} else {
				printBoxed(fmt.Sprintf("%sUser %s created%s", ColorGreen, username, ColorReset))
			}

		case "del": // Was user-del
			if len(args) < 1 {
				printBoxed("Usage: user del <username>")
				return
			}
			username := args[0]
			fmt.Printf("Delete user '%s'? (y/N): ", username)
			var resp string
			fmt.Scanln(&resp)
			if strings.ToLower(resp) != "y" {
				return
			}
			if err := server.DeleteUser(username, token); err != nil {
				printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
			} else {
				printBoxed(fmt.Sprintf("%sUser deleted%s", ColorGreen, ColorReset))
			}

		case "list":
			users, err := server.ListUsers(token)
			if err != nil {
				printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
				return
			}
			var buf bytes.Buffer
			w := tabwriter.NewWriter(&buf, 0, 0, 2, ' ', 0)
			fmt.Fprintln(w, "Username\tRoles\tStatus")
			fmt.Fprintln(w, "--------\t-----\t------")
			for _, u := range users {
				status := "Offline"
				if u.IsOnline { status = "Online" }
				fmt.Fprintf(w, "%s\t%s\t%s\n", u.Username, strings.Join(u.Roles, ","), status)
			}
			w.Flush()
			printBoxed(buf.String())

		case "update":
			// user update <username> [roles=r1,r2]
			// Currently only password update via prompt, maybe roles later
			// The prompt implies a generic update but also mentions "user update"
			// Let's implement password update here or allow changing roles.
			// Ideally arguments: user update <username> roles=admin,dev
			// If no args (other than username), prompt for password?
			if len(args) < 1 {
				printBoxed("Usage: user update <username> [roles=r1,r2...]")
				return
			}
			username := args[0]
			
			// If args present, update roles
			if len(args) > 1 {
				// Parse options
				updated := false
				// user, err := server.AuthManager.GetUser(username) // Use internal GetUser helper or logic?
				// We need authenticated way. server.ListUsers gives us access if admin.
				// But we need the User object to modify.
				// Let's rely on server.UpdateUser taking a struct.
				// We first need to GET the user to avoid overwriting fields.
				// Since we don't have GetUser exposed in server (only List), let's add one or iterate List.
				// Iterating List is inefficient but works for now.
				users, _ := server.ListUsers(token)
				var targetUser *User
				for _, u := range users {
					if u.Username == username {
						targetUser = u
						break
					}
				}
				if targetUser == nil {
					printBoxed("User not found or permission denied")
					return
				}
				
				// Apply updates
				for _, arg := range args[1:] {
					if strings.HasPrefix(arg, "roles=") {
						roleStr := strings.TrimPrefix(arg, "roles=")
						targetUser.Roles = strings.Split(roleStr, ",")
						updated = true
					}
				}
				
				if updated {
					if err := server.UpdateUser(*targetUser, token); err != nil {
						printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
					} else {
						printBoxed(fmt.Sprintf("%sUser updated%s", ColorGreen, ColorReset))
					}
				}
			} else {
				// No args -> Change Password
				fmt.Printf("Enter new password for %s: ", username)
				pass, err := readPassword()
				fmt.Println()
				if err != nil || pass == "" {
					printBoxed("Error reading password")
					return
				}
				if err := server.ChangeUserPassword(username, pass, token); err != nil {
					printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
				} else {
					printBoxed(fmt.Sprintf("%sPassword updated%s", ColorGreen, ColorReset))
				}
			}

		case "passwd": // Keep alias or subcommand
			if len(args) < 1 {
				printBoxed("Usage: user passwd <username>")
				return
			}
			// Reuse logic
			username := args[0]
			fmt.Printf("Enter new password for %s: ", username)
			pass, err := readPassword()
			fmt.Println()
			if err != nil || pass == "" {
				printBoxed("Error reading password")
				return
			}
			if err := server.ChangeUserPassword(username, pass, token); err != nil {
				printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
			} else {
				printBoxed(fmt.Sprintf("%sPassword updated%s", ColorGreen, ColorReset))
			}

		case "role":
			if len(args) < 1 {
				printBoxed("Usage: user role <list|add|delete|update> ...")
				return
			}
			roleCmd := strings.ToLower(args[0])
			roleArgs := args[1:]
			
			switch roleCmd {
			case "list":
				roles, err := server.ListRoles(token)
				if err != nil {
					printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
					return
				}
				var buf bytes.Buffer
				w := tabwriter.NewWriter(&buf, 0, 0, 2, ' ', 0)
				fmt.Fprintln(w, "Role\tParents\tPerms")
				fmt.Fprintln(w, "----\t-------\t-----")
				for _, r := range roles {
					parents := strings.Join(r.ParentRoles, ",")
					permCount := len(r.Permissions)
					fmt.Fprintf(w, "%s\t%s\t%d rules\n", r.Name, parents, permCount)
				}
				w.Flush()
				printBoxed(buf.String())
				
			case "add":
				if len(roleArgs) < 1 {
					printBoxed("Usage: user role add <name>")
					return
				}
				if err := server.CreateRole(roleArgs[0], token); err != nil {
					printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
				} else {
					printBoxed(fmt.Sprintf("%sRole %s created%s", ColorGreen, roleArgs[0], ColorReset))
				}
				
			case "delete":
				if len(roleArgs) < 1 {
					printBoxed("Usage: user role delete <name>")
					return
				}
				if err := server.DeleteRole(roleArgs[0], token); err != nil {
					printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
				} else {
					printBoxed(fmt.Sprintf("%sRole %s deleted%s", ColorGreen, roleArgs[0], ColorReset))
				}
				
			case "update":
				// Simplistic update: user role update <name> parents=p1,p2
				// Complex permissions editing is hard in CLI without JSON
				if len(roleArgs) < 1 {
					printBoxed("Usage: user role update <name> [parents=p1,p2]")
					return
				}
				name := roleArgs[0]
				roles, _ := server.ListRoles(token)
				var targetRole *Role
				for _, r := range roles {
					if r.Name == name {
						targetRole = r
						break
					}
				}
				if targetRole == nil {
					printBoxed("Role not found")
					return
				}
				
				updated := false
				for _, arg := range roleArgs[1:] {
					if strings.HasPrefix(arg, "parents=") {
						pStr := strings.TrimPrefix(arg, "parents=")
						targetRole.ParentRoles = strings.Split(pStr, ",")
						updated = true
					}
				}
				
				if updated {
					if err := server.UpdateRole(*targetRole, token); err != nil {
						printBoxed(fmt.Sprintf("%sError:%s %v", ColorRed, ColorReset, err))
					} else {
						printBoxed(fmt.Sprintf("%sRole updated%s", ColorGreen, ColorReset))
					}
				} else {
					printBoxed("No changes specified.")
				}
				
			default:
				printBoxed("Unknown sub-command. Options: list, add, delete, update")
			}

		case "help":
			var sb strings.Builder
			sb.WriteString(fmt.Sprintf("%sUser Management Commands:%s\n", ColorCyan, ColorReset))
			
			// Define subcommands for formatting
			type subCmd struct {
				Name string
				Desc string
			}
			cmds := []subCmd{
				{"user init", "Initialize auth system (Root)"},
				{"user login <user> [code]", "Login"},
				{"user logout", "Logout"},
				{"user who", "Current user info"},
				{"user add <user> [roles]", "Create user (Root)"},
				{"user del <user>", "Delete user (Root)"},
				{"user update <user> [opts]", "Update user/pass (Root)"},
				{"user list", "List users (Root)"},
				{"user role list", "List roles (Root)"},
				{"user role add <name>", "Create role (Root)"},
				{"user role delete <name>", "Delete role (Root)"},
				{"user role update <name> ...", "Update role (Root)"},
			}
			
			// Find max length
			maxLen := 0
			for _, c := range cmds {
				if len(c.Name) > maxLen {
					maxLen = len(c.Name)
				}
			}
			
			for _, c := range cmds {
				padding := strings.Repeat(" ", maxLen-len(c.Name))
				sb.WriteString(fmt.Sprintf("  %s%s%s%s  %s%s%s\n", ColorGreen, c.Name, ColorReset, padding, ColorDim, c.Desc, ColorReset))
			}
			
			printBoxed(sb.String())

		default:
			printBoxed(fmt.Sprintf("Unknown subcommand '%s'. Type 'user help'.", subCmd))
		}
	})

	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
		}
		
		// Admin only
		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()
		if server.AuthManager.IsEnabled() && !server.IsRoot(token) {
			printBoxed(fmt.Sprintf("%sPermission Denied: Root access required%s", ColorRed, ColorReset))
			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
		}

		// Admin only
		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()
		if server.AuthManager.IsEnabled() && !server.IsRoot(token) {
			printBoxed(fmt.Sprintf("%sPermission Denied: Root access required%s", ColorRed, ColorReset))
			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
			}
		}

		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()

		// Use Authenticated Search
		results, err := server.SearchAuthenticated(pattern, limit, offset, token)
		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]
		
		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()
		
		start := time.Now()
		count, err := server.CountAuthenticated(pattern, token)
		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) {
		// Admin only
		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()
		if server.AuthManager.IsEnabled() && !server.IsRoot(token) {
			printBoxed(fmt.Sprintf("%sPermission Denied: Root access required%s", ColorRed, ColorReset))
			return
		}

		// 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) {
		// Admin only
		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()
		if server.AuthManager.IsEnabled() && !server.IsRoot(token) {
			printBoxed(fmt.Sprintf("%sPermission Denied: Root access required%s", ColorRed, ColorReset))
			return
		}

		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("clear", "Clear screen", func(parts []string, server *KVServer) {
		fmt.Print("\033[H\033[2J")
	})

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

// readPassword securely reads a password from stdin without echoing.
func readPassword() (string, error) {
	// 1. Disable echo
	cmd := exec.Command("stty", "-echo")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("failed to disable echo: %v", err)
	}

	defer func() {
		// 3. Re-enable echo
		cmd := exec.Command("stty", "echo")
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		_ = cmd.Run()
	}()

	// 2. Read input
	reader := bufio.NewReader(os.Stdin)
	pass, err := reader.ReadString('\n')
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(pass), nil
}