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

	c.RegisterCommand("login", "Login to system (login <user> <pass> [code])", func(parts []string, server *KVServer) {
		if len(parts) < 3 {
			printBoxed("Usage: login <username> <password> [mfa_code]")
			return
		}
		username := parts[1]
		password := parts[2]
		code := ""
		if len(parts) > 3 {
			code = parts[3]
		}

		token, 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 = token
			c.mu.Unlock()
			printBoxed(fmt.Sprintf("%sLogin Success! Token saved.%s", ColorGreen, ColorReset))
		}
	})

	c.RegisterCommand("logout", "Logout from system", func(parts []string, server *KVServer) {
		c.mu.Lock()
		c.token = ""
		c.mu.Unlock()
		printBoxed("Logged out.")
	})

	c.RegisterCommand("auth-init", "Initialize Auth System (auth-init <root_pass>)", func(parts []string, server *KVServer) {
		if len(parts) != 2 {
			printBoxed("Usage: auth-init <root_password>")
			return
		}
		
		// 1. Create Root User
		if err := server.AuthManager.CreateRootUser(parts[1]); err != nil {
			printBoxed(fmt.Sprintf("Failed to create root user: %v", err))
			return
		}
		
		// 2. Enable Auth
		if err := server.AuthManager.EnableAuth(); err != nil {
			printBoxed(fmt.Sprintf("Failed to enable auth: %v", err))
			return
		}
		
		printBoxed("Auth Initialized. Please login as 'root'.")
	})
	
	c.RegisterCommand("whoami", "Show current user info", func(parts []string, server *KVServer) {
		c.mu.RLock()
		token := c.token
		c.mu.RUnlock()
		
		if token == "" {
			printBoxed("Not logged in.")
			return
		}
		
		// Decode token or check session
		// Accessing AuthManager internal session map is protected.
		// We can add a helper in AuthManager to get session info or just show token.
		printBoxed(fmt.Sprintf("Logged in (Token: %s...)", token[:8]))
	})

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