package main

import (
	"bufio"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"golang.org/x/crypto/curve25519"
)

// Colors for terminal output
const (
	Red    = "\033[0;31m"
	Green  = "\033[0;32m"
	Yellow = "\033[1;33m"
	Blue   = "\033[0;34m"
	Reset  = "\033[0m"
)

// Script constants
const (
	ScriptVersion = "3.0.0"
	DefaultPort   = "51820"
	MaxHostname   = 63
	MaxInterface  = 15
)

// WireGuard peer configuration
type WireGuardPeer struct {
	Name                string
	PublicKey           string
	AllowedIPs          string
	Endpoint            string
	PersistentKeepalive int
	Description         string
}

// Node configuration
type NodeConfig struct {
	Hostname    string `json:"hostname"`
	IPAddress   string `json:"ip_address"`
	PrivateKey  string `json:"private_key"`
	PublicKey   string `json:"public_key"`
	RoutingMode string `json:"routing_mode"`
	Interface   string `json:"interface"`
	Generated   string `json:"generated"`
	ScriptVer   string `json:"script_version"`
	RunningRoot bool   `json:"running_as_root"`
}

// Application state
type AppState struct {
	ForceMode     bool
	RunningAsRoot bool
	WGDirectory   string
	StepCounter   int
}

// Default peers - automatically included in every config
var defaultPeers = []WireGuardPeer{
	{
		Name:                "Zion",
		PublicKey:           "2ztJbrN1x1NWanzPGLiKL19ZkdOhm5Y7WeKEWBT5cyg=",
		AllowedIPs:          "10.8.0.0/24",
		Endpoint:            "ugh.im:51820",
		PersistentKeepalive: 25,
		Description:         "Central server (always included)",
	},
	{
		Name:                "CTH",
		PublicKey:           "NBktXKy1s0n2lIlIMODvOqKNwAtYdoZH5feKt5P43i0=",
		AllowedIPs:          "10.8.0.10/32",
		Endpoint:            "aw2cd67.glddns.com:53535",
		PersistentKeepalive: 25,
		Description:         "Secondary server (always included)",
	},
}

// Reserved IP addresses
var reservedIPs = map[string]string{
	"10.8.0.1":  "Zion",
	"10.8.0.10": "CTH",
	"10.8.0.2":  "Aza",
	"10.8.0.20": "Nyar",
	"10.8.0.99": "Galaxy",
	"10.8.0.7":  "nanocube",
	"10.8.0.42": "jupiter",
	"10.8.0.8":  "HASS",
	"10.8.0.40": "framebot",
}

// Static server ports
var staticServerPorts = map[string]string{
	"10.8.0.1":  "51820", // Zion
	"10.8.0.10": "53535", // CTH
	"10.8.0.99": "54382", // Galaxy
}

// Validation regexes
var (
	ipRegex        = regexp.MustCompile(`^10\.8\.0\.\d{1,3}$`)
	hostnameRegex  = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9-]*[a-zA-Z0-9]$|^[a-zA-Z0-9]$`)
	interfaceRegex = regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9]*$`)
	keyRegex       = regexp.MustCompile(`^[A-Za-z0-9+/]{43}=$`)
)

// Utility functions
func step(n int, msg string) {
	fmt.Printf("%sStep %d:%s %s\n", Blue, n, Reset, msg)
}

func printStatus(message string) {
	fmt.Printf("%s[INFO]%s %s\n", Green, Reset, message)
}

func printWarning(message string) {
	fmt.Printf("%s[WARNING]%s %s\n", Yellow, Reset, message)
}

func printError(message string) {
	fmt.Printf("%s[ERROR]%s %s\n", Red, Reset, message)
}

func printHeader(title string) {
	fmt.Printf("%s================================%s\n", Blue, Reset)
	fmt.Printf("%s%s%s\n", Blue, title, Reset)
	fmt.Printf("%s================================%s\n", Blue, Reset)
	fmt.Println()
}

// Validation functions
func validateHostname(hostname string) error {
	if hostname == "" {
		return fmt.Errorf("hostname cannot be empty")
	}
	if len(hostname) > MaxHostname {
		return fmt.Errorf("hostname too long (max %d characters)", MaxHostname)
	}
	if !hostnameRegex.MatchString(hostname) {
		return fmt.Errorf("invalid hostname format. Use alphanumeric characters and hyphens only")
	}
	return nil
}

func validateIP(ip string) error {
	if ip == "" {
		return fmt.Errorf("IP address cannot be empty")
	}

	if !ipRegex.MatchString(ip) {
		return fmt.Errorf("IP should be in 10.8.0.x range for NextGen network")
	}

	if parsedIP := net.ParseIP(ip); parsedIP == nil {
		return fmt.Errorf("invalid IP address format")
	}

	if peerName, exists := reservedIPs[ip]; exists {
		return fmt.Errorf("IP %s is already reserved for %s", ip, peerName)
	}

	return nil
}

func validateInterface(name string) error {
	if name == "" {
		return nil // Allow empty for default
	}
	if len(name) > MaxInterface {
		return fmt.Errorf("interface name too long (max %d characters)", MaxInterface)
	}
	if !interfaceRegex.MatchString(name) {
		return fmt.Errorf("interface name must start with a letter and contain only letters and numbers")
	}
	return nil
}

func validateWireGuardKey(key string) error {
	if key == "" {
		return fmt.Errorf("key cannot be empty")
	}
	if !keyRegex.MatchString(key) {
		return fmt.Errorf("invalid WireGuard key format")
	}
	return nil
}

// User input functions
func getUserInput(prompt string, validator func(string) error, forceMode bool) string {
	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print(prompt)
		input, _ := reader.ReadString('\n')
		input = strings.TrimSpace(input)

		if err := validator(input); err != nil {
			if forceMode {
				printWarning(fmt.Sprintf("Validation failed: %s, but continuing due to force mode", err.Error()))
				return input
			}
			printError(err.Error())
			fmt.Print("Continue anyway? (y/N): ")
			response, _ := reader.ReadString('\n')
			response = strings.TrimSpace(strings.ToLower(response))
			if response == "y" || response == "yes" {
				return input
			}
			continue
		}
		return input
	}
}

func getDirectoryInput(prompt, defaultDir string, forceMode bool) string {
	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Printf("%s [%s]: ", prompt, defaultDir)
		input, _ := reader.ReadString('\n')
		input = strings.TrimSpace(input)

		if input == "" {
			input = defaultDir
		}

		// Check if directory exists
		if _, err := os.Stat(input); os.IsNotExist(err) {
			if forceMode {
				if err := os.MkdirAll(input, 0755); err != nil {
					printError(fmt.Sprintf("Failed to create directory '%s': %v", input, err))
					continue
				}
			} else {
				fmt.Printf("Directory '%s' doesn't exist. Create it? (Y/n): ", input)
				response, _ := reader.ReadString('\n')
				response = strings.TrimSpace(strings.ToLower(response))
				if response == "n" || response == "no" {
					continue
				}

				if err := os.MkdirAll(input, 0755); err != nil {
					printError(fmt.Sprintf("Failed to create directory '%s': %v", input, err))
					continue
				}
			}
		}

		// Check if directory is writable
		if info, err := os.Stat(input); err == nil {
			if info.Mode()&0200 == 0 {
				printError(fmt.Sprintf("Directory '%s' is not writable", input))
				continue
			}
		}

		return input
	}
}

func getYesNoInput(prompt string, defaultYes bool, forceMode bool) bool {
	if forceMode {
		return defaultYes
	}

	reader := bufio.NewReader(os.Stdin)
	for {
		if defaultYes {
			fmt.Printf("%s (Y/n): ", prompt)
		} else {
			fmt.Printf("%s (y/N): ", prompt)
		}
		input, _ := reader.ReadString('\n')
		input = strings.TrimSpace(strings.ToLower(input))

		if input == "" {
			return defaultYes
		}
		if input == "y" || input == "yes" {
			return true
		}
		if input == "n" || input == "no" {
			return false
		}
		printError("Please enter 'y' or 'n'")
	}
}

func getChoice(prompt string, choices []string, forceMode bool) int {
	if forceMode && len(choices) > 0 {
		return 1 // Default to first choice in force mode
	}

	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Println(prompt)
		for i, choice := range choices {
			fmt.Printf("%d. %s\n", i+1, choice)
		}
		fmt.Print("Enter your choice: ")
		input, _ := reader.ReadString('\n')
		input = strings.TrimSpace(input)

		if choice, err := strconv.Atoi(input); err == nil && choice >= 1 && choice <= len(choices) {
			return choice
		}
		printError(fmt.Sprintf("Invalid choice. Please enter a number between 1 and %d.", len(choices)))
	}
}

// WireGuard key generation (safe version)
func generateWireGuardKeys() (string, string, error) {
	// Generate private key
	privateKeyBytes := make([]byte, 32)
	if _, err := rand.Read(privateKeyBytes); err != nil {
		return "", "", fmt.Errorf("failed to generate random bytes: %w", err)
	}

	// Ensure the key is valid for curve25519
	privateKeyBytes[0] &= 248
	privateKeyBytes[31] &= 127
	privateKeyBytes[31] |= 64

	// Generate public key (safe conversion)
	var privateKeyArray [32]byte
	copy(privateKeyArray[:], privateKeyBytes)

	var publicKeyBytes [32]byte
	curve25519.ScalarBaseMult(&publicKeyBytes, &privateKeyArray)

	privateKey := base64.StdEncoding.EncodeToString(privateKeyBytes[:])
	publicKey := base64.StdEncoding.EncodeToString(publicKeyBytes[:])

	return privateKey, publicKey, nil
}

// Configuration generation
func generateConfig(hostname, ipAddress, privateKey, routingMode, interfaceName string) string {
	var config strings.Builder

	// Interface section
	config.WriteString("[Interface]\n")
	config.WriteString(fmt.Sprintf("Address = %s/24\n", ipAddress))
	config.WriteString(fmt.Sprintf("PrivateKey = %s\n", privateKey))

	// Add ListenPort for static servers
	if port, exists := staticServerPorts[ipAddress]; exists {
		config.WriteString(fmt.Sprintf("ListenPort = %s\n", port))
	}

	// Add DNS for full tunnel mode
	if routingMode == "full_tunnel" {
		config.WriteString("DNS = 1.1.1.1, 8.8.8.8\n")
	}

	// Add default peers
	for _, peer := range defaultPeers {
		config.WriteString(fmt.Sprintf("\n# %s (%s)\n", peer.Name, peer.Description))
		config.WriteString("[Peer]\n")
		config.WriteString(fmt.Sprintf("PublicKey = %s\n", peer.PublicKey))

		// Set AllowedIPs based on routing mode
		if routingMode == "full_tunnel" {
			config.WriteString("AllowedIPs = 0.0.0.0/0, ::/0\n")
		} else {
			config.WriteString(fmt.Sprintf("AllowedIPs = %s\n", peer.AllowedIPs))
		}

		config.WriteString(fmt.Sprintf("Endpoint = %s\n", peer.Endpoint))
		config.WriteString(fmt.Sprintf("PersistentKeepalive = %d\n", peer.PersistentKeepalive))
	}

	return config.String()
}

// Generate peer configuration for Zion
func generateZionPeerConfig(hostname, publicKey, ipAddress string) string {
	var config strings.Builder
	config.WriteString(fmt.Sprintf("# %s\n", hostname))
	config.WriteString("[Peer]\n")
	config.WriteString(fmt.Sprintf("PublicKey = %s\n", publicKey))
	config.WriteString(fmt.Sprintf("AllowedIPs = %s/32\n", ipAddress))
	return config.String()
}

// System checks
func checkDependencies() error {
	deps := []string{"wg", "wg-quick"}

	for _, dep := range deps {
		if _, err := exec.LookPath(dep); err != nil {
			return fmt.Errorf("missing dependency: %s", dep)
		}
	}
	return nil
}

func isRunningAsRoot() bool {
	return os.Geteuid() == 0
}

func checkFileExists(filepath string, forceMode bool) bool {
	if _, err := os.Stat(filepath); err == nil {
		if forceMode {
			printWarning(fmt.Sprintf("File '%s' already exists, overwriting due to force mode", filepath))
			return true
		}
		printWarning(fmt.Sprintf("File '%s' already exists", filepath))
		return getYesNoInput("Do you want to overwrite it?", false, false)
	}
	return true
}

// Display instructions
func displayZionInstructions(hostname, publicKey, ipAddress string) {
	printHeader("IMPORTANT: Update Zion Server Configuration")
	fmt.Println()
	printWarning("You MUST add this peer to Zion's configuration file:")
	fmt.Println("   /etc/wireguard/wg0.conf")
	fmt.Println()

	fmt.Println("Add the following peer section to Zion's config:")
	fmt.Println("----------------------------------------")
	fmt.Println(generateZionPeerConfig(hostname, publicKey, ipAddress))
	fmt.Println("----------------------------------------")
	fmt.Println()

	printWarning("After updating Zion's config:")
	fmt.Println("1. Save the file")
	fmt.Println("2. Restart Zion's WireGuard: sudo systemctl restart wg-quick@wg0")
	fmt.Println("3. Then start this node's WireGuard: sudo wg-quick up " + hostname)
	fmt.Println()

	// Show Zion's current configuration structure
	fmt.Println("Zion's current configuration structure:")
	fmt.Println("----------------------------------------")
	fmt.Println("[Interface]")
	fmt.Println("Address = 10.8.0.1/24")
	fmt.Println("ListenPort = 51820")
	fmt.Println("PrivateKey = <zion_private_key>")
	fmt.Println("PostUp = <iptables_rules>")
	fmt.Println("PostDown = <iptables_rules>")
	fmt.Println()

	// Show existing peers
	for name, pubKey := range map[string]string{
		"Cth":      "NBktXKy1s0n2lIlIMODvOqKNwAtYdoZH5feKt5P43i0=",
		"Aza":      "qmTKA257DLOrfhk5Zw8RyRmBSonmm6epbloT0P0ZWDc=",
		"Nyar":     "2BA7L1oJP1tK6dIUNHMgcZmOmYmlyPRe2RaBqfUsEWo=",
		"Galaxy":   "QBNt00VSedxPlq3ZvsdYaqIcbudCAyxv9TG65aPVZzM=",
		"nanocube": "/ZImoATDIS0e0N08CD7mqWbhtGlSnynpPuY04Ed4Zyc=",
		"jupiter":  "YIFQ43ULk/YoCgOv3SBU6+MOrbxd+mlvaw9rT8uoNmw=",
		"HASS":     "C+Poz/7DaXCxe4HZiL6D5cld4jMt5o1gBq3iPiBzrg0=",
		"framebot": "loS3yZapqmt6lP53Q+s4EvUzw6FmwgZC8jzgLluJ1Es=",
	} {
		fmt.Printf("#%s\n", name)
		fmt.Println("[Peer]")
		fmt.Printf("PublicKey = %s\n", pubKey)
		fmt.Printf("AllowedIPs = 10.8.0.x/32\n")
		fmt.Println()
	}

	fmt.Println("# Add your peer here:")
	fmt.Printf("# %s\n", hostname)
	fmt.Println("# [Peer]")
	fmt.Printf("# PublicKey = %s\n", publicKey)
	fmt.Printf("# AllowedIPs = %s/32\n", ipAddress)
	fmt.Println("----------------------------------------")
	fmt.Println()
}

func displayFullTunnelInstructions() {
	fmt.Println()
	printWarning("FULL TUNNEL MODE DETECTED - Endpoint Changes Required:")
	fmt.Println()
	fmt.Println("Since this node will route ALL traffic through the VPN, you need to:")
	fmt.Println()
	fmt.Println("1. Update Zion's config (/etc/wireguard/wg0.conf) to allow this peer:")
	fmt.Println("   - Add the peer section as shown above")
	fmt.Println("   - Ensure Zion has proper iptables rules for NAT/masquerading")
	fmt.Println()
	fmt.Println("2. Check Zion's iptables rules (run on Zion server):")
	fmt.Println("   sudo iptables -t nat -L POSTROUTING")
	fmt.Println("   sudo iptables -L FORWARD")
	fmt.Println()
	fmt.Println("3. If Zion doesn't have proper NAT rules, add them:")
	fmt.Println("   sudo iptables -t nat -A POSTROUTING -s 10.8.0.0/24 -o eth0 -j MASQUERADE")
	fmt.Println("   sudo iptables -A FORWARD -i wg0 -j ACCEPT")
	fmt.Println("   sudo iptables -A FORWARD -o wg0 -j ACCEPT")
	fmt.Println()
	fmt.Println("4. Enable IP forwarding on Zion (if not already enabled):")
	fmt.Println("   echo 'net.ipv4.ip_forward=1' | sudo tee -a /etc/sysctl.conf")
	fmt.Println("   sudo sysctl -p")
	fmt.Println()
}

func displayNextSteps(hostname, interfaceName, publicKey, ipAddress, routingMode string, runningAsRoot bool) {
	step(4, "Next Steps")
	fmt.Println()

	if runningAsRoot {
		printStatus("Ready to start WireGuard:")
		fmt.Printf("   systemctl enable --now wg-quick@%s\n", interfaceName)
	} else {
		printWarning("To enable WireGuard (requires root):")
		fmt.Printf("   sudo cp %s /etc/wireguard/\n", filepath.Join("wireguard_configs", fmt.Sprintf("%s.conf", interfaceName)))
		fmt.Printf("   sudo chmod 600 /etc/wireguard/%s.conf\n", interfaceName)
		fmt.Printf("   sudo systemctl enable --now wg-quick@%s\n", interfaceName)
	}
	fmt.Println()
	printWarning("IMPORTANT: Update other nodes with this peer info:")
	fmt.Printf("   PublicKey = %s\n", publicKey)
	fmt.Printf("   AllowedIPs = %s/32\n", ipAddress)
	fmt.Println()

	if routingMode == "full_tunnel" {
		displayFullTunnelInstructions()
	}

	// Zion update instructions
	fmt.Printf("%s========================================%s\n", Red, Reset)
	fmt.Printf("%s  !!! NOW UPDATE ZION SERVER !!!  %s\n", Red, Reset)
	fmt.Printf("%s========================================%s\n", Red, Reset)
	fmt.Println()
	printWarning("You MUST add this peer to Zion's config (/etc/wireguard/wg0.conf):")
	fmt.Println()
	fmt.Printf("%s#%s%s\n", Yellow, hostname, Reset)
	fmt.Printf("%s[Peer]%s\n", Yellow, Reset)
	fmt.Printf("%sPublicKey = %s%s\n", Yellow, publicKey, Reset)
	fmt.Printf("%sAllowedIPs = %s/32%s\n", Yellow, ipAddress, Reset)
	fmt.Println()
	printWarning("After updating Zion, restart its WireGuard:")
	fmt.Println("   systemctl restart wg-quick@wg0")
	fmt.Println()

	if runningAsRoot {
		printWarning("Then restart this node's WireGuard:")
		fmt.Printf("   systemctl restart wg-quick@%s\n", interfaceName)
	} else {
		printWarning("Then restart this node's WireGuard:")
		fmt.Printf("   sudo systemctl restart wg-quick@%s\n", interfaceName)
	}
}

// Main function
func main() {
	// Parse command line flags
	var forceMode bool
	flag.BoolVar(&forceMode, "force", false, "Skip confirmations and use defaults")
	flag.Parse()

	// Initialize application state
	state := &AppState{
		ForceMode:     forceMode,
		RunningAsRoot: isRunningAsRoot(),
		StepCounter:   1,
	}

	printHeader("WireGuard Configuration Setup v" + ScriptVersion)
	fmt.Println("This script will help you create WireGuard keys and configuration files.")
	fmt.Println("Based on the CURRENT_WORKING configuration with Zion as central server.")
	fmt.Println("Default peers (Zion & CTH) will be automatically included.")
	fmt.Println()

	if forceMode {
		printWarning("Force mode enabled - skipping confirmations and using defaults")
		fmt.Println()
	}

	// Determine WireGuard directory
	if state.RunningAsRoot {
		state.WGDirectory = "/etc/wireguard"
		printStatus("Running as root - using system directories")
		printStatus(fmt.Sprintf("WireGuard directory: %s", state.WGDirectory))
	} else {
		state.WGDirectory = "wireguard_configs"
		printWarning("Not running as root - using current directory")
		printStatus(fmt.Sprintf("WireGuard directory: %s", state.WGDirectory))
		printWarning("You'll need to manually copy config files to /etc/wireguard/ later")
	}
	fmt.Println()

	// Check dependencies
	if err := checkDependencies(); err != nil {
		printError(err.Error())
		printStatus("Install with: apt install wireguard-tools")
		os.Exit(1)
	}

	// Get directory for non-root users
	if !state.RunningAsRoot {
		step(state.StepCounter, "Directory Selection")
		state.StepCounter++
		printStatus("Choose where to save WireGuard files:")
		fmt.Printf("  - Current directory: %s\n", state.WGDirectory)
		fmt.Printf("  - Home directory: %s\n", os.Getenv("HOME"))
		fmt.Println("  - Custom directory")
		fmt.Println()
		state.WGDirectory = getDirectoryInput("Enter directory path for WireGuard files", state.WGDirectory, forceMode)
		fmt.Println()
	}

	// Create directory
	if err := os.MkdirAll(state.WGDirectory, 0755); err != nil {
		printError(fmt.Sprintf("Failed to create directory: %v", err))
		os.Exit(1)
	}

	// Get node information
	step(state.StepCounter, "Node Information")
	state.StepCounter++
	fmt.Println()

	hostname := getUserInput("Enter hostname for this node: ", validateHostname, forceMode)

	// Get IP address
	fmt.Println()
	fmt.Println("Available IP ranges:")
	fmt.Println("  - 10.8.0.x (recommended for NextGen network)")
	fmt.Println("  - 10.0.0.x (alternative range)")
	fmt.Println()

	ipAddress := getUserInput("Enter IP address for this node (e.g., 10.8.0.30): ", validateIP, forceMode)

	// Get interface name
	fmt.Println()
	fmt.Println("Interface name options:")
	fmt.Println("  - wg0 (default, most common)")
	fmt.Println("  - wg1, wg2, etc. (if wg0 is already in use)")
	fmt.Println("  - Custom name (e.g., nextgen, vpn, etc.)")
	fmt.Println()

	interfaceName := getUserInput("Enter interface name (default: wg0): ", validateInterface, forceMode)
	if interfaceName == "" {
		interfaceName = "wg0"
	}

	// Check if configuration file already exists
	configFile := filepath.Join(state.WGDirectory, fmt.Sprintf("%s.conf", interfaceName))
	if !checkFileExists(configFile, forceMode) {
		printError("Operation cancelled. Please choose a different interface name or remove the existing file.")
		os.Exit(1)
	}

	// Configuration options
	fmt.Println()
	step(state.StepCounter, "Configuration Options")
	state.StepCounter++
	fmt.Println()
	fmt.Println("Choose an option:")
	fmt.Println("1. Generate keys only (manual config creation)")
	fmt.Println("2. Generate keys + complete config (recommended)")
	fmt.Println()

	var configChoice int
	if forceMode {
		configChoice = 2 // Default to complete config in force mode
		printStatus("Force mode: Using complete config generation")
	} else {
		configChoice = getChoice("Enter your choice:", []string{"Generate keys only", "Generate keys + complete config"}, false)
	}

	// Traffic routing options
	var routingMode string
	if configChoice == 2 {
		fmt.Println()
		fmt.Println("Traffic routing options:")
		fmt.Println("1. WireGuard traffic only (10.8.0.x network only)")
		fmt.Println("2. All traffic through VPN (full tunnel)")
		fmt.Println()
		fmt.Println("Note: Full tunnel routes ALL internet traffic through the VPN.")
		fmt.Println("      WireGuard-only keeps your regular internet traffic separate.")
		fmt.Println()

		var routingChoice int
		if forceMode {
			routingChoice = 1 // Default to WireGuard-only in force mode
			printStatus("Force mode: Using WireGuard-only routing")
		} else {
			routingChoice = getChoice("Enter your choice:", []string{"WireGuard traffic only", "All traffic through VPN"}, false)
		}

		if routingChoice == 1 {
			routingMode = "wg_only"
			printStatus("Selected: WireGuard traffic only")
		} else {
			routingMode = "full_tunnel"
			printStatus("Selected: All traffic through VPN")
		}
	}

	printStatus(fmt.Sprintf("Starting setup for %s (%s)...", hostname, ipAddress))
	fmt.Println()

	// Generate keys
	printStatus("Generating WireGuard keys...")
	privateKey, publicKey, err := generateWireGuardKeys()
	if err != nil {
		printError(fmt.Sprintf("Failed to generate keys: %v", err))
		os.Exit(1)
	}

	// Save keys
	privateKeyFile := filepath.Join(state.WGDirectory, fmt.Sprintf("%s_private.key", hostname))
	publicKeyFile := filepath.Join(state.WGDirectory, fmt.Sprintf("%s_public.key", hostname))

	if err := os.WriteFile(privateKeyFile, []byte(privateKey), 0600); err != nil {
		printError(fmt.Sprintf("Failed to save private key: %v", err))
		os.Exit(1)
	}

	if err := os.WriteFile(publicKeyFile, []byte(publicKey), 0600); err != nil {
		printError(fmt.Sprintf("Failed to save public key: %v", err))
		os.Exit(1)
	}

	printStatus("Keys generated successfully!")
	fmt.Printf("   Private key: %s\n", privateKeyFile)
	fmt.Printf("   Public key: %s\n", publicKeyFile)
	fmt.Println()

	// Display node information
	step(state.StepCounter, "Node Information")
	state.StepCounter++
	fmt.Println("==========================================")
	fmt.Printf("HOSTNAME: %s\n", hostname)
	fmt.Printf("IP ADDRESS: %s\n", ipAddress)
	fmt.Printf("PRIVATE KEY: %s\n", privateKey)
	fmt.Printf("PUBLIC KEY: %s\n", publicKey)
	fmt.Printf("INTERFACE: %s\n", interfaceName)
	fmt.Printf("ROUTING MODE: %s\n", routingMode)
	fmt.Println("==========================================")
	fmt.Println()

	// Save structured info
	infoFile := filepath.Join(state.WGDirectory, fmt.Sprintf("%s_wg_info.json", hostname))
	if state.RunningAsRoot {
		infoFile = filepath.Join("/tmp", fmt.Sprintf("%s_wg_info.json", hostname))
	}

	nodeConfig := NodeConfig{
		Hostname:    hostname,
		IPAddress:   ipAddress,
		PrivateKey:  privateKey,
		PublicKey:   publicKey,
		RoutingMode: routingMode,
		Interface:   interfaceName,
		Generated:   time.Now().Format(time.RFC3339),
		ScriptVer:   ScriptVersion,
		RunningRoot: state.RunningAsRoot,
	}

	infoData, err := json.MarshalIndent(nodeConfig, "", "  ")
	if err != nil {
		printError(fmt.Sprintf("Failed to marshal config: %v", err))
		os.Exit(1)
	}

	if err := os.WriteFile(infoFile, infoData, 0600); err != nil {
		printError(fmt.Sprintf("Failed to save info file: %v", err))
		os.Exit(1)
	}

	printStatus(fmt.Sprintf("Information saved to: %s", infoFile))
	fmt.Println()

	// Generate complete config if requested
	if configChoice == 2 {
		printStatus(fmt.Sprintf("Generating complete %s.conf...", interfaceName))
		configContent := generateConfig(hostname, ipAddress, privateKey, routingMode, interfaceName)

		if err := os.WriteFile(configFile, []byte(configContent), 0600); err != nil {
			printError(fmt.Sprintf("Failed to write config file: %v", err))
			os.Exit(1)
		}

		printStatus(fmt.Sprintf("Config written to: %s", configFile))
		if state.RunningAsRoot {
			printStatus("Permissions set to 600")
		}
		fmt.Println()

		// Display Zion integration instructions
		displayZionInstructions(hostname, publicKey, ipAddress)

		// Display next steps
		displayNextSteps(hostname, interfaceName, publicKey, ipAddress, routingMode, state.RunningAsRoot)

		// Config preview
		fmt.Printf("%sConfig Preview:%s\n", Blue, Reset)
		fmt.Println("----------------------------------------")
		lines := strings.Split(configContent, "\n")
		for i, line := range lines {
			if i >= 8 { // Show more lines to include default peers
				break
			}
			fmt.Println(line)
		}
		fmt.Printf("   [... and %d total lines]\n", len(lines))
		fmt.Println("----------------------------------------")
		fmt.Println()
	} else {
		// Manual config generation path
		step(state.StepCounter, "Next Steps")
		state.StepCounter++
		fmt.Println()
		printStatus("Keys generated successfully!")
		fmt.Printf("   Private key: %s\n", privateKeyFile)
		fmt.Printf("   Public key: %s\n", publicKeyFile)
		fmt.Println()
		printWarning("Next steps:")
		fmt.Printf("   - Create %s.conf manually using the keys above\n", interfaceName)
		fmt.Printf("   - Copy config to %s\n", filepath.Join(state.WGDirectory, fmt.Sprintf("%s.conf", interfaceName)))
		if state.RunningAsRoot {
			fmt.Printf("   - Set permissions: chmod 600 %s\n", filepath.Join(state.WGDirectory, fmt.Sprintf("%s.conf", interfaceName)))
			fmt.Printf("   - Enable/start: systemctl enable --now wg-quick@%s\n", interfaceName)
		} else {
			fmt.Printf("   - Copy to system: sudo cp %s /etc/wireguard/\n", filepath.Join(state.WGDirectory, fmt.Sprintf("%s.conf", interfaceName)))
			fmt.Printf("   - Set permissions: sudo chmod 600 /etc/wireguard/%s.conf\n", interfaceName)
			fmt.Printf("   - Enable/start: sudo systemctl enable --now wg-quick@%s\n", interfaceName)
		}
		fmt.Println()
		fmt.Printf("%s========================================%s\n", Red, Reset)
		fmt.Printf("%s  !!! NOW UPDATE ZION SERVER !!!  %s\n", Red, Reset)
		fmt.Printf("%s========================================%s\n", Red, Reset)
		fmt.Println()
		printWarning("You MUST add this peer to Zion's config (/etc/wireguard/wg0.conf):")
		fmt.Println()
		fmt.Printf("%s#%s%s\n", Yellow, hostname, Reset)
		fmt.Printf("%s[Peer]%s\n", Yellow, Reset)
		fmt.Printf("%sPublicKey = %s%s\n", Yellow, publicKey, Reset)
		fmt.Printf("%sAllowedIPs = %s/32%s\n", Yellow, ipAddress, Reset)
		fmt.Println()
		printWarning("After updating Zion, restart its WireGuard:")
		fmt.Println("   systemctl restart wg-quick@wg0")
		fmt.Println()

		if state.RunningAsRoot {
			printWarning("Then restart this node's WireGuard:")
			fmt.Printf("   systemctl restart wg-quick@%s\n", interfaceName)
		} else {
			printWarning("Then restart this node's WireGuard:")
			fmt.Printf("   sudo systemctl restart wg-quick@%s\n", interfaceName)
		}
	}

	fmt.Println()
	printStatus(fmt.Sprintf("Setup complete for %s!", hostname))
	fmt.Println()
}