kalzu-value-store/main.go

package main

import (
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"path/filepath"
	"strings"

	"kvs/config"
	"kvs/daemon"
	"kvs/server"
)

func main() {
	if len(os.Args) < 2 {
		// No arguments - run in foreground with default config
		runServer("./config.yaml", false)
		return
	}

	// Check if this is a daemon spawn
	if os.Args[1] == "--daemon" {
		if len(os.Args) < 3 {
			fmt.Fprintf(os.Stderr, "Error: --daemon flag requires config path\n")
			os.Exit(1)
		}
		runServer(os.Args[2], true)
		return
	}

	// Parse subcommand
	command := os.Args[1]

	switch command {
	case "start":
		if len(os.Args) < 3 {
			fmt.Fprintf(os.Stderr, "Usage: kvs start <config>\n")
			os.Exit(1)
		}
		cmdStart(normalizeConfigPath(os.Args[2]))

	case "stop":
		if len(os.Args) < 3 {
			fmt.Fprintf(os.Stderr, "Usage: kvs stop <config>\n")
			os.Exit(1)
		}
		cmdStop(normalizeConfigPath(os.Args[2]))

	case "restart":
		if len(os.Args) < 3 {
			fmt.Fprintf(os.Stderr, "Usage: kvs restart <config>\n")
			os.Exit(1)
		}
		cmdRestart(normalizeConfigPath(os.Args[2]))

	case "status":
		if len(os.Args) > 2 {
			cmdStatusSingle(normalizeConfigPath(os.Args[2]))
		} else {
			cmdStatusAll()
		}

	case "help", "--help", "-h":
		printHelp()

	default:
		// Backward compatibility: assume it's a config file path
		runServer(command, false)
	}
}

func runServer(configPath string, isDaemon bool) {
	cfg, err := config.Load(configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load configuration: %v\n", err)
		os.Exit(1)
	}

	// Write PID file if running as daemon
	if isDaemon {
		if err := daemon.WritePID(configPath); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to write PID file: %v\n", err)
			os.Exit(1)
		}
		defer daemon.RemovePID(configPath)
	}

	kvServer, err := server.NewServer(cfg)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create server: %v\n", err)
		os.Exit(1)
	}

	// Handle graceful shutdown
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-sigCh
		kvServer.Stop()
	}()

	if err := kvServer.Start(); err != nil && err != http.ErrServerClosed {
		fmt.Fprintf(os.Stderr, "Server error: %v\n", err)
		os.Exit(1)
	}
}

func cmdStart(configPath string) {
	if err := daemon.Daemonize(configPath); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to start: %v\n", err)
		os.Exit(1)
	}
}

func cmdStop(configPath string) {
	pid, running, err := daemon.ReadPID(configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read PID: %v\n", err)
		os.Exit(1)
	}

	if !running {
		fmt.Printf("Instance '%s' is not running\n", configPath)
		// Clean up stale PID file
		daemon.RemovePID(configPath)
		return
	}

	fmt.Printf("Stopping instance '%s' (PID %d)...\n", configPath, pid)
	if err := daemon.StopProcess(pid); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to stop process: %v\n", err)
		os.Exit(1)
	}

	// Wait a bit and verify it stopped
	time.Sleep(1 * time.Second)
	_, stillRunning, _ := daemon.ReadPID(configPath)
	if stillRunning {
		fmt.Printf("Warning: Process may still be running\n")
	} else {
		daemon.RemovePID(configPath)
		fmt.Printf("Stopped successfully\n")
	}
}

func cmdRestart(configPath string) {
	// Check if running
	_, running, err := daemon.ReadPID(configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to check status: %v\n", err)
		os.Exit(1)
	}

	if running {
		cmdStop(configPath)
		// Wait a bit for clean shutdown
		time.Sleep(2 * time.Second)
	}

	cmdStart(configPath)
}

func cmdStatusSingle(configPath string) {
	pid, running, err := daemon.ReadPID(configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read PID: %v\n", err)
		os.Exit(1)
	}

	if running {
		fmt.Printf("Instance '%s': RUNNING (PID %d)\n", configPath, pid)
	} else if pid > 0 {
		fmt.Printf("Instance '%s': STOPPED (stale PID %d)\n", configPath, pid)
	} else {
		fmt.Printf("Instance '%s': STOPPED\n", configPath)
	}
}

func cmdStatusAll() {
	instances, err := daemon.ListRunningInstances()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to list instances: %v\n", err)
		os.Exit(1)
	}

	if len(instances) == 0 {
		fmt.Println("No KVS instances found")
		return
	}

	fmt.Println("KVS Instances:")
	for _, inst := range instances {
		status := "STOPPED"
		if inst.Running {
			status = "RUNNING"
		}
		fmt.Printf("  %-20s %s (PID %d)\n", inst.Name, status, inst.PID)
	}
}

// normalizeConfigPath ensures config path has .yaml extension if not specified
func normalizeConfigPath(path string) string {
	// If path doesn't have an extension, add .yaml
	if filepath.Ext(path) == "" {
		return path + ".yaml"
	}
	return path
}

// getConfigIdentifier returns the identifier for a config (basename without extension)
// This is used for PID files and status display
func getConfigIdentifier(path string) string {
	basename := filepath.Base(path)
	return strings.TrimSuffix(basename, filepath.Ext(basename))
}

func printHelp() {
	help := `KVS - Distributed Key-Value Store

Usage:
  kvs [config.yaml]              Run in foreground (default: ./config.yaml)
  kvs start <config>             Start as daemon (.yaml extension optional)
  kvs stop <config>              Stop daemon (.yaml extension optional)
  kvs restart <config>           Restart daemon (.yaml extension optional)
  kvs status [config]            Show status (all instances if no config given)
  kvs help                       Show this help

Examples:
  kvs                            # Run with ./config.yaml in foreground
  kvs node1.yaml                 # Run with node1.yaml in foreground
  kvs start node1                # Start node1.yaml as daemon
  kvs start node1.yaml           # Same as above
  kvs stop node1                 # Stop node1 daemon
  kvs status                     # Show all running instances
  kvs status node1               # Show status of node1
`
	fmt.Print(help)
}