package main

import (
	"bufio"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"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"
)

// Configuration structs
type WGConfig 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"`
	Generated   string `json:"generated"`
	ScriptVer   string `json:"script_version"`
	RunningRoot bool   `json:"running_as_root"`
}

type StaticServer struct {
	IP   string
	Port string
}

// Global variables
var (
	staticServers = map[string]string{
		"10.8.0.1":  "51820",
		"10.8.0.10": "53535",
		"10.8.0.99": "54382",
	}

	peers = map[string]Peer{
		"zion": {
			Name:      "Zion peer (central server) - for access to entire network",
			PublicKey: "2ztJbrN1x1NWanzPGLiKL19ZkdOhm5Y7WeKEWBT5cyg=",
			Endpoint:  "ugh.im:51820",
			Keepalive: 25,
		},
		"cthulhu": {
			Name:       "Cthulhu (optional if port 53 is also forwarded to bypass firewalls)",
			PublicKey:  "NBktXKy1s0n2lIlIMODvOqKNwAtYdoZH5feKt5P43i0=",
			AllowedIPs: "10.8.0.10/32",
			Endpoint:   "aw2cd67.glddns.com:53535",
			Keepalive:  25,
		},
		"galaxy": {
			Name:       "Galaxy (located in Europe, NL)",
			PublicKey:  "QBNt00VSedxPlq3ZvsdYaqIcbudCAyxv9TG65aPVZzM=",
			AllowedIPs: "10.8.0.99/32",
			Endpoint:   "galaxyspin.space:54382",
			Keepalive:  25,
		},
	}
)

type Peer struct {
	Name       string
	PublicKey  string
	AllowedIPs string
	Endpoint   string
	Keepalive  int
}

// Utility functions
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() {
	fmt.Printf("%s================================%s\n", Blue, Reset)
	fmt.Printf("%s  NextGen WireGuard Easy Setup  %s\n", Blue, 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) > 63 {
		return fmt.Errorf("hostname too long (max 63 characters)")
	}
	// Basic alphanumeric and hyphen validation
	for _, char := range hostname {
		if !((char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') ||
			(char >= '0' && char <= '9') || char == '-') {
			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")
	}

	parts := strings.Split(ip, ".")
	if len(parts) != 4 {
		return fmt.Errorf("invalid IP format")
	}

	for _, part := range parts {
		num, err := strconv.Atoi(part)
		if err != nil || num < 0 || num > 255 {
			return fmt.Errorf("invalid IP octet: %s", part)
		}
	}

	// Check if it's in the expected range
	if !strings.HasPrefix(ip, "10.8.0.") && !strings.HasPrefix(ip, "10.0.0.") {
		printWarning("IP should be in 10.8.0.x or 10.0.0.x range for NextGen network")
	}

	return nil
}

func validateInterfaceName(iface string) error {
	if iface == "" {
		return fmt.Errorf("interface name cannot be empty")
	}
	if len(iface) > 15 {
		return fmt.Errorf("interface name too long (max 15 characters)")
	}

	// Check if starts with letter and contains only alphanumeric
	if len(iface) == 0 || !((iface[0] >= 'a' && iface[0] <= 'z') || (iface[0] >= 'A' && iface[0] <= 'Z')) {
		return fmt.Errorf("interface name must start with a letter")
	}

	for _, char := range iface {
		if !((char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') || (char >= '0' && char <= '9')) {
			return fmt.Errorf("interface name can only contain letters and numbers")
		}
	}

	return nil
}

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

		if err := validator(input); err != nil {
			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) 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) {
			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
	}
}

// WireGuard key generation
func generateWireGuardKeys() (string, string, error) {
	// Generate private key
	privateKeyBytes := make([]byte, 32)
	if _, err := rand.Read(privateKeyBytes); err != nil {
		return "", "", err
	}

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

	// Generate public key
	var publicKeyBytes [32]byte
	curve25519.ScalarBaseMult(&publicKeyBytes, (*[32]byte)(privateKeyBytes))

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

	return privateKey, publicKey, nil
}

// Configuration generation
func generateConfig(hostname, ipAddress, privateKey, routingMode 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 := staticServers[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")
	}

	// Zion peer (always enabled)
	config.WriteString("\n#Zion peer (central server) - for access to entire network\n")
	config.WriteString("[Peer]\n")
	config.WriteString(fmt.Sprintf("PublicKey = %s\n", peers["zion"].PublicKey))

	// Set AllowedIPs based on routing mode
	if routingMode == "full_tunnel" {
		config.WriteString("AllowedIPs = 0.0.0.0/0, ::/0\n")
	} else {
		config.WriteString("AllowedIPs = 10.8.0.0/24\n")
	}

	config.WriteString(fmt.Sprintf("Endpoint = %s\n", peers["zion"].Endpoint))
	config.WriteString(fmt.Sprintf("PersistentKeepalive = %d\n", peers["zion"].Keepalive))

	// Optional peers (commented out)
	config.WriteString("\n#Cthulhu (optional if port 53 is also forwarded to bypass firewalls)\n")
	config.WriteString("#[Peer]\n")
	config.WriteString(fmt.Sprintf("#PublicKey = %s\n", peers["cthulhu"].PublicKey))
	config.WriteString(fmt.Sprintf("#AllowedIPs = %s\n", peers["cthulhu"].AllowedIPs))
	config.WriteString(fmt.Sprintf("#Endpoint = %s\n", peers["cthulhu"].Endpoint))
	config.WriteString(fmt.Sprintf("#PersistentKeepalive = %d\n", peers["cthulhu"].Keepalive))

	config.WriteString("\n#Galaxy (located in Europe, NL)\n")
	config.WriteString("#[Peer]\n")
	config.WriteString(fmt.Sprintf("#PublicKey = %s\n", peers["galaxy"].PublicKey))
	config.WriteString(fmt.Sprintf("#AllowedIPs = %s\n", peers["galaxy"].AllowedIPs))
	config.WriteString(fmt.Sprintf("#Endpoint = %s\n", peers["galaxy"].Endpoint))
	config.WriteString(fmt.Sprintf("#PersistentKeepalive = %d\n", peers["galaxy"].Keepalive))

	return config.String()
}

// Check dependencies
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
}

// Check if running as root
func isRunningAsRoot() bool {
	return os.Geteuid() == 0
}

// Main function
func main() {
	printHeader()

	// Check if running as root
	runningAsRoot := isRunningAsRoot()

	var wgDir string
	if runningAsRoot {
		wgDir = "/etc/wireguard"
		printStatus("Running as root - using system directories")
		printStatus(fmt.Sprintf("WireGuard directory: %s", wgDir))
	} else {
		wgDir = "."
		printWarning("Not running as root - using current directory")
		printStatus(fmt.Sprintf("WireGuard directory: %s", wgDir))
		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 !runningAsRoot {
		fmt.Printf("%sStep 1: Directory Selection%s\n", Blue, Reset)
		fmt.Println()
		printStatus("Choose where to save WireGuard files:")
		fmt.Printf("  - Current directory: %s\n", wgDir)
		fmt.Printf("  - Home directory: %s\n", os.Getenv("HOME"))
		fmt.Println("  - Custom directory")
		fmt.Println()
		wgDir = getDirectoryInput("Enter directory path for WireGuard files", wgDir)
		fmt.Println()
	}

	// Get hostname
	stepNum := "2"
	if !runningAsRoot {
		stepNum = "2"
	} else {
		stepNum = "1"
	}
	fmt.Printf("%sStep %s: Node Information%s\n", Blue, stepNum, Reset)
	fmt.Println()

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

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

	// 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): ", validateInterfaceName)
	if interfaceName == "" {
		interfaceName = "wg0"
	}

	// Configuration options
	fmt.Println()
	if !runningAsRoot {
		stepNum = "3"
	} else {
		stepNum = "2"
	}
	fmt.Printf("%sStep %s: Configuration Options%s\n", Blue, stepNum, Reset)
	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()

	reader := bufio.NewReader(os.Stdin)
	var configChoice string
	for {
		fmt.Print("Enter your choice (1 or 2): ")
		configChoice, _ = reader.ReadString('\n')
		configChoice = strings.TrimSpace(configChoice)
		if configChoice == "1" || configChoice == "2" {
			break
		}
		printError("Invalid choice. Please enter 1 or 2.")
	}

	// 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 string
		for {
			fmt.Print("Enter your choice (1 or 2): ")
			routingChoice, _ = reader.ReadString('\n')
			routingChoice = strings.TrimSpace(routingChoice)
			if routingChoice == "1" || routingChoice == "2" {
				break
			}
			printError("Invalid choice. Please enter 1 or 2.")
		}

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

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

	// 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(wgDir, fmt.Sprintf("%s_private.key", hostname))
	publicKeyFile := filepath.Join(wgDir, fmt.Sprintf("%s_public.key", hostname))

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

	if err := ioutil.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
	if !runningAsRoot {
		stepNum = "4"
	} else {
		stepNum = "3"
	}
	fmt.Printf("%sStep %s: Node Information%s\n", Blue, stepNum, Reset)
	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.Println("==========================================")
	fmt.Println()

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

	wgConfig := WGConfig{
		Hostname:    hostname,
		IPAddress:   ipAddress,
		PrivateKey:  privateKey,
		PublicKey:   publicKey,
		RoutingMode: routingMode,
		Generated:   time.Now().Format(time.RFC3339),
		ScriptVer:   "2.2",
		RunningRoot: runningAsRoot,
	}

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

	if err := ioutil.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" {
		configFile := filepath.Join(wgDir, fmt.Sprintf("%s.conf", interfaceName))

		printStatus(fmt.Sprintf("Generating complete %s.conf...", interfaceName))
		configContent := generateConfig(hostname, ipAddress, privateKey, routingMode)

		if err := ioutil.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 runningAsRoot {
			printStatus("Permissions set to 600")
		}
		fmt.Println()

		// Next steps
		if !runningAsRoot {
			stepNum = "5"
		} else {
			stepNum = "4"
		}
		fmt.Printf("%sStep %s: Next Steps%s\n", Blue, stepNum, Reset)
		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", configFile)
			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()

		// Config preview
		fmt.Printf("%sConfig Preview:%s\n", Blue, Reset)
		fmt.Println("----------------------------------------")
		lines := strings.Split(configContent, "\n")
		for i, line := range lines {
			if i >= 5 {
				break
			}
			fmt.Println(line)
		}
		fmt.Printf("   [... and %d total lines]\n", len(lines))
		fmt.Println("----------------------------------------")
		fmt.Println()

		// Full tunnel instructions
		if routingMode == "full_tunnel" {
			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()
		}

		// 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)
		}
	} else {
		// Manual config generation path
		if !runningAsRoot {
			stepNum = "5"
		} else {
			stepNum = "4"
		}
		fmt.Printf("%sStep %s: Next Steps%s\n", Blue, stepNum, Reset)
		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(wgDir, fmt.Sprintf("%s.conf", interfaceName)))
		if runningAsRoot {
			fmt.Printf("   - Set permissions: chmod 600 %s\n", filepath.Join(wgDir, 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(wgDir, 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 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()
}