#!/bin/bash

# NextGen WireGuard Easy Setup Script (wgez-setup.sh)
# Interactive setup script for new WireGuard nodes

set -euo pipefail

# Colors for better UX
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# Globals
SCRIPT_DIR="${BASH_SOURCE[0]%/*}"
SCRIPT_DIR="$(cd "$SCRIPT_DIR" && pwd)"

# Determine if running as root and set appropriate directories
if [[ $EUID -eq 0 ]]; then
    WG_DIR="/etc/wireguard"
    BACKUP_DIR="$WG_DIR/backups"
    RUNNING_AS_ROOT=true
else
    WG_DIR="$(pwd)"
    BACKUP_DIR="$WG_DIR/backups"
    RUNNING_AS_ROOT=false
fi

# Logging
log_file="/var/log/wgez-setup.log"
if [[ "$RUNNING_AS_ROOT" == "false" ]]; then
    log_file="$WG_DIR/wgez-setup.log"
fi
exec 19> >(exec cat >&2)
exec 20> >(exec tee -a "$log_file" >&2)
BASH_XTRACEFD=20

print_status() {
    echo -e "${GREEN}[INFO]${NC} $1" | tee -a "$log_file"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1" | tee -a "$log_file"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1" | tee -a "$log_file"
}

print_header() {
    echo -e "${BLUE}================================${NC}"
    echo -e "${BLUE}  NextGen WireGuard Easy Setup  ${NC}"
    echo -e "${BLUE}================================${NC}"
    echo ""
}

cleanup() {
    local exit_code=$?
    if [[ $exit_code -ne 0 ]]; then
        print_error "Script failed with exit code $exit_code"
        if [[ -f "/tmp/${HOSTNAME:-unknown}_wg_info.json" ]]; then
            rm -f "/tmp/${HOSTNAME}_wg_info.json"
        fi
    fi
    exit $exit_code
}

# Enhanced validation functions
validate_hostname() {
    local hostname="$1"
    if [[ -z "$hostname" ]]; then
        print_error "Hostname cannot be empty"
        return 1
    fi
    if [[ ! "$hostname" =~ ^[a-zA-Z0-9][a-zA-Z0-9-]*[a-zA-Z0-9]$|^[a-zA-Z0-9]$ ]]; then
        print_error "Invalid hostname format. Use alphanumeric characters and hyphens only."
        return 1
    fi
    if [[ ${#hostname} -gt 63 ]]; then
        print_error "Hostname too long (max 63 characters)"
        return 1
    fi
    return 0
}

validate_ip() {
    local ip="$1"
    local octets
    
    if [[ -z "$ip" ]]; then
        print_error "IP address cannot be empty"
        return 1
    fi
    
    IFS='.' read -ra octets <<< "$ip"
    if [[ ${#octets[@]} -ne 4 ]]; then
        print_error "Invalid IP format"
        return 1
    fi
    
    for octet in "${octets[@]}"; do
        if ! [[ "$octet" =~ ^[0-9]+$ ]] || [[ $octet -gt 255 ]] || [[ $octet -lt 0 ]]; then
            print_error "Invalid IP octet: $octet"
            return 1
        fi
    done
    
    if [[ ! "$ip" =~ ^10\.(8|0)\.0\.[0-9]{1,3}$ ]]; then
        print_warning "IP should be in 10.8.0.x or 10.0.0.x range for NextGen network"
        return 2
    fi
    
    return 0
}

validate_interface_name() {
    local interface="$1"
    
    if [[ -z "$interface" ]]; then
        print_error "Interface name cannot be empty"
        return 1
    fi
    
    if [[ ! "$interface" =~ ^[a-zA-Z][a-zA-Z0-9]*$ ]]; then
        print_error "Invalid interface name. Use letters and numbers only, starting with a letter."
        return 1
    fi
    
    if [[ ${#interface} -gt 15 ]]; then
        print_error "Interface name too long (max 15 characters)"
        return 1
    fi
    
    # Check if interface already exists (only if running as root)
    if [[ "$RUNNING_AS_ROOT" == "true" ]] && ip link show "$interface" &>/dev/null; then
        print_warning "Interface '$interface' already exists"
        return 2
    fi
    
    # Check if config file already exists
    if [[ -f "$WG_DIR/${interface}.conf" ]]; then
        print_warning "Config file '$WG_DIR/${interface}.conf' already exists"
        return 2
    fi
    
    return 0
}

check_ip_conflict() {
    local ip="$1"
    local config_file="$2"
    
    # Check if IP is already in use on network
    if ping -c1 -W1 "$ip" &>/dev/null; then
        print_warning "IP $ip appears to be in use on the network"
        return 1
    fi
    
    # Check existing WireGuard configs
    if [[ -f "$config_file" ]] && grep -q "Address = $ip" "$config_file"; then
        print_warning "IP $ip already configured in existing WireGuard config"
        return 1
    fi
    
    return 0
}

check_dependencies() {
    local deps=("wg" "wg-quick")
    local missing=()
    
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &>/dev/null; then
            missing+=("$dep")
        fi
    done
    
    if [[ ${#missing[@]} -gt 0 ]]; then
        print_error "Missing dependencies: ${missing[*]}"
        print_status "Install with: apt install wireguard-tools"
        return 1
    fi
    
    # Only check systemctl if running as root
    if [[ "$RUNNING_AS_ROOT" == "true" ]] && ! command -v "systemctl" &>/dev/null; then
        print_error "Missing dependency: systemctl"
        print_status "Install with: apt install systemd"
        return 1
    fi
    
    return 0
}

backup_existing_config() {
    local config_file="$1"
    if [[ -f "$config_file" ]]; then
        mkdir -p "$BACKUP_DIR"
        local config_name=$(basename "$config_file")
        local backup_file="$BACKUP_DIR/${config_name}.backup.$(date +%Y%m%d_%H%M%S)"
        cp "$config_file" "$backup_file"
        print_status "Existing config backed up to: $backup_file"
    fi
}

get_user_input() {
    local prompt="$1"
    local var_name="$2"
    local validator="$3"
    local value
    
    while true; do
        read -p "$prompt" value
        if $validator "$value"; then
            eval "$var_name='$value'"
            break
        elif [[ $? -eq 2 ]]; then
            read -p "Continue anyway? (y/N): " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                eval "$var_name='$value'"
                break
            fi
        fi
    done
}

get_directory_input() {
    local prompt="$1"
    local var_name="$2"
    local default_dir="$3"
    local directory
    
    while true; do
        read -p "$prompt [$default_dir]: " directory
        if [[ -z "$directory" ]]; then
            directory="$default_dir"
        fi
        
        # Create directory if it doesn't exist
        if [[ ! -d "$directory" ]]; then
            read -p "Directory '$directory' doesn't exist. Create it? (Y/n): " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Nn]$ ]]; then
                continue
            fi
            if ! mkdir -p "$directory" 2>/dev/null; then
                print_error "Failed to create directory '$directory'"
                continue
            fi
        fi
        
        # Check if directory is writable
        if [[ ! -w "$directory" ]]; then
            print_error "Directory '$directory' is not writable"
            continue
        fi
        
        eval "$var_name='$directory'"
        break
    done
}

generate_config() {
    local hostname="$1"
    local ip_address="$2"
    local private_key="$3"
    local routing_mode="${4:-wg_only}"
    
    # Define static servers
    local -A static_servers=(
        ["10.8.0.1"]="51820"
        ["10.8.0.10"]="53535"
        ["10.8.0.99"]="54382"
    )
    
    local config_content="[Interface]
Address = $ip_address/24
PrivateKey = $private_key"
    
    # Add ListenPort for static servers
    local is_static_server=false
    if [[ -n "${static_servers[$ip_address]:-}" ]]; then
        is_static_server=true
        config_content="$config_content
ListenPort = ${static_servers[$ip_address]}"
    fi
    
    # Add DNS for full tunnel mode
    if [[ "$routing_mode" == "full_tunnel" ]]; then
        config_content="$config_content
DNS = 1.1.1.1, 8.8.8.8"
    fi
    
    config_content="$config_content

#Zion peer (central server) - for access to entire network
[Peer]
PublicKey = 2ztJbrN1x1NWanzPGLiKL19ZkdOhm5Y7WeKEWBT5cyg="
    
    # Set AllowedIPs based on routing mode
    if [[ "$routing_mode" == "full_tunnel" ]]; then
        config_content="$config_content
AllowedIPs = 0.0.0.0/0, ::/0"
    else
        config_content="$config_content
AllowedIPs = 10.8.0.0/24"
    fi
    
    config_content="$config_content
Endpoint = ugh.im:51820"
    
    # Set keepalive value (25 for all peers as per new configuration)
    config_content="$config_content
PersistentKeepalive = 25"
    
    config_content="$config_content

#Cthulhu (optional if port 53 is also forwarded to bypass firewalls)
#[Peer]
#PublicKey = NBktXKy1s0n2lIlIMODvOqKNwAtYdoZH5feKt5P43i0=
#AllowedIPs = 10.8.0.10/32
#Endpoint = aw2cd67.glddns.com:53535"
    
    config_content="$config_content
#PersistentKeepalive = 25"

    config_content="$config_content

#Galaxy (located in Europe, NL)
#[Peer]
#PublicKey = QBNt00VSedxPlq3ZvsdYaqIcbudCAyxv9TG65aPVZzM=
#AllowedIPs = 10.8.0.99/32
#Endpoint = galaxyspin.space:54382"
    
    config_content="$config_content
#PersistentKeepalive = 25"
    
    echo "$config_content"
}

print_progress() {
    local step="$1"
    local total="$2"
    echo -e "${BLUE}[$step/$total]${NC} $3"
}

validate_config() {
    local config_file="$1"
    if ! wg-quick strip "$config_file" >/dev/null 2>&1; then
        print_error "Generated config has syntax errors"
        return 1
    fi
    return 0
}

check_network_connectivity() {
    if ! ping -c1 -W3 8.8.8.8 &>/dev/null; then
        print_warning "No internet connectivity detected"
        return 1
    fi
    return 0
}

main() {
    trap cleanup EXIT
    
    print_header
    
    # Check if running as root and inform user
    if [[ "$RUNNING_AS_ROOT" == "true" ]]; then
        print_status "Running as root - using system directories"
        print_status "WireGuard directory: $WG_DIR"
    else
        print_warning "Not running as root - using current directory"
        print_status "WireGuard directory: $WG_DIR"
        print_warning "You'll need to manually copy config files to /etc/wireguard/ later"
    fi
    echo ""
    
    # Check dependencies
    if ! check_dependencies; then
        exit 1
    fi
    
    # Get directory for non-root users
    if [[ "$RUNNING_AS_ROOT" == "false" ]]; then
        echo -e "${BLUE}Step 1: Directory Selection${NC}"
        echo ""
        print_status "Choose where to save WireGuard files:"
        echo "  - Current directory: $(pwd)"
        echo "  - Home directory: $HOME"
        echo "  - Custom directory"
        echo ""
        get_directory_input "Enter directory path for WireGuard files" WG_DIR "$(pwd)"
        BACKUP_DIR="$WG_DIR/backups"
        echo ""
    fi
    
    # Get hostname
    echo -e "${BLUE}Step $([[ "$RUNNING_AS_ROOT" == "false" ]] && echo "2" || echo "1"): Node Information${NC}"
    echo ""
    
    get_user_input "Enter hostname for this node: " HOSTNAME validate_hostname
    
    # Get IP address
    echo ""
    echo "Available IP ranges:"
    echo "  - 10.8.0.x (recommended for NextGen network)"
    echo "  - 10.0.0.x (alternative range)"
    echo ""
    
    get_user_input "Enter IP address for this node (e.g., 10.8.0.30): " IP_ADDRESS validate_ip
    
    # Get interface name
    echo ""
    echo "Interface name options:"
    echo "  - wg0 (default, most common)"
    echo "  - wg1, wg2, etc. (if wg0 is already in use)"
    echo "  - Custom name (e.g., nextgen, vpn, etc.)"
    echo ""
    
    get_user_input "Enter interface name (default: wg0): " INTERFACE_NAME validate_interface_name
    if [[ -z "$INTERFACE_NAME" ]]; then
        INTERFACE_NAME="wg0"
    fi
    
    # Set config file path
    CONFIG_FILE="$WG_DIR/${INTERFACE_NAME}.conf"
    
    # Check for IP conflicts
    if ! check_ip_conflict "$IP_ADDRESS" "$CONFIG_FILE"; then
        read -p "Continue anyway? (y/N): " -n 1 -r
        echo
        [[ ! $REPLY =~ ^[Yy]$ ]] && exit 1
    fi
    
    # Check if keys already exist
    if [[ -f "$WG_DIR/${HOSTNAME}_private.key" ]]; then
        print_warning "Keys for $HOSTNAME already exist!"
        read -p "Overwrite? (y/N): " -n 1 -r
        echo
        [[ ! $REPLY =~ ^[Yy]$ ]] && exit 1
    fi
    
    # Configuration options
    echo ""
    echo -e "${BLUE}Step $([[ "$RUNNING_AS_ROOT" == "false" ]] && echo "3" || echo "2"): Configuration Options${NC}"
    echo ""
    echo "Choose an option:"
    echo "1. Generate keys only (manual config creation)"
    echo "2. Generate keys + complete ${INTERFACE_NAME}.conf (recommended)"
    echo ""
    
    while true; do
        read -p "Enter your choice (1 or 2): " CONFIG_CHOICE
        if [[ "$CONFIG_CHOICE" == "1" || "$CONFIG_CHOICE" == "2" ]]; then
            break
        fi
        print_error "Invalid choice. Please enter 1 or 2."
    done
    
    # Traffic routing options (only if generating complete config)
    if [[ "$CONFIG_CHOICE" == "2" ]]; then
        echo ""
        echo "Traffic routing options:"
        echo "1. WireGuard traffic only (10.8.0.x network only)"
        echo "2. All traffic through VPN (full tunnel)"
        echo ""
        echo "Note: Full tunnel routes ALL internet traffic through the VPN."
        echo "      WireGuard-only keeps your regular internet traffic separate."
        echo ""
        
        while true; do
            read -p "Enter your choice (1 or 2): " ROUTING_CHOICE
            if [[ "$ROUTING_CHOICE" == "1" || "$ROUTING_CHOICE" == "2" ]]; then
                break
            fi
            print_error "Invalid choice. Please enter 1 or 2."
        done
        
        if [[ "$ROUTING_CHOICE" == "1" ]]; then
            ROUTING_MODE="wg_only"
            print_status "Selected: WireGuard traffic only"
        else
            ROUTING_MODE="full_tunnel"
            print_status "Selected: All traffic through VPN"
        fi
    fi
    
    print_status "Starting setup for $HOSTNAME ($IP_ADDRESS)..."
    echo ""
    
    # Create directories
    mkdir -p "$WG_DIR" "$BACKUP_DIR"
    cd "$WG_DIR"
    
    # Set secure permissions (only if running as root)
    if [[ "$RUNNING_AS_ROOT" == "true" ]]; then
        umask 077
    fi
    
    print_status "Generating WireGuard keys..."
    PRIVATE_KEY=$(wg genkey)
    PUBLIC_KEY=$(echo "$PRIVATE_KEY" | wg pubkey)
    
    # Save keys
    printf "%s\n%s\n" "$PRIVATE_KEY" "$PUBLIC_KEY" > "${HOSTNAME}_keys.tmp"
    mv "${HOSTNAME}_keys.tmp" "${HOSTNAME}_private.key"
    echo "$PUBLIC_KEY" > "${HOSTNAME}_public.key"
    
    # Set permissions (only if running as root)
    if [[ "$RUNNING_AS_ROOT" == "true" ]]; then
        chmod 600 "${HOSTNAME}_private.key" "${HOSTNAME}_public.key"
    fi
    
    print_status "Keys generated successfully!"
    echo "   Private key: $WG_DIR/${HOSTNAME}_private.key"
    echo "   Public key: $WG_DIR/${HOSTNAME}_public.key"
    echo ""
    
    echo -e "${BLUE}Step $([[ "$RUNNING_AS_ROOT" == "false" ]] && echo "4" || echo "3"): Node Information${NC}"
    echo "=========================================="
    echo "HOSTNAME: $HOSTNAME"
    echo "IP ADDRESS: $IP_ADDRESS"
    echo "PRIVATE KEY: $PRIVATE_KEY"
    echo "PUBLIC KEY: $PUBLIC_KEY"
    echo "=========================================="
    echo ""
    
    # Save structured info
    INFO_FILE="/tmp/${HOSTNAME}_wg_info.json"
    if [[ "$RUNNING_AS_ROOT" == "false" ]]; then
        INFO_FILE="$WG_DIR/${HOSTNAME}_wg_info.json"
    fi
    cat > "$INFO_FILE" << EOF
{
  "hostname": "$HOSTNAME",
  "ip_address": "$IP_ADDRESS",
  "private_key": "$PRIVATE_KEY",
  "public_key": "$PUBLIC_KEY",
  "routing_mode": "${ROUTING_MODE:-wg_only}",
  "generated": "$(date -Iseconds)",
  "script_version": "2.2",
  "running_as_root": $RUNNING_AS_ROOT
}
EOF
    
    print_status "Information saved to: $INFO_FILE"
    echo ""
    
    # Generate complete config if requested
    if [[ "$CONFIG_CHOICE" == "2" ]]; then
        backup_existing_config "$CONFIG_FILE"
        
        print_status "Generating complete ${INTERFACE_NAME}.conf..."
        generate_config "$HOSTNAME" "$IP_ADDRESS" "$PRIVATE_KEY" "$ROUTING_MODE" > "$CONFIG_FILE"
        
        # Set permissions (only if running as root)
        if [[ "$RUNNING_AS_ROOT" == "true" ]]; then
            chmod 600 "$CONFIG_FILE"
        fi
        
        print_status "Config written to: $CONFIG_FILE"
        if [[ "$RUNNING_AS_ROOT" == "true" ]]; then
            print_status "Permissions set to 600"
        fi
        echo ""
        
        echo -e "${BLUE}Step $([[ "$RUNNING_AS_ROOT" == "false" ]] && echo "5" || echo "4"): Next Steps${NC}"
        echo ""
        
        if [[ "$RUNNING_AS_ROOT" == "true" ]]; then
            print_status "Ready to start WireGuard:"
            echo "   systemctl enable --now wg-quick@${INTERFACE_NAME}"
        else
            print_warning "To enable WireGuard (requires root):"
            echo "   sudo cp $CONFIG_FILE /etc/wireguard/"
            echo "   sudo chmod 600 /etc/wireguard/${INTERFACE_NAME}.conf"
            echo "   sudo systemctl enable --now wg-quick@${INTERFACE_NAME}"
        fi
        echo ""
        print_warning "IMPORTANT: Update other nodes with this peer info:"
        echo "   PublicKey = $PUBLIC_KEY"
        echo "   AllowedIPs = $IP_ADDRESS/32"
        echo ""
        
        echo -e "${BLUE}Config Preview:${NC}"
        echo "----------------------------------------"
        head -5 "$CONFIG_FILE"
        echo "   [... and $(wc -l < "$CONFIG_FILE" | tr -d ' ') total lines]"
        echo "----------------------------------------"
        
        echo ""
        if [[ "$RUNNING_AS_ROOT" == "true" ]]; then
            print_status "Configuration complete! To start WireGuard:"
            echo "   systemctl enable --now wg-quick@${INTERFACE_NAME}"
            echo ""
            print_status "To check status:"
            echo "   wg show ${INTERFACE_NAME}"
            echo "   systemctl status wg-quick@${INTERFACE_NAME}"
            echo ""
            print_status "To view logs:"
            echo "   journalctl -u wg-quick@${INTERFACE_NAME} -f"
        else
            print_status "Configuration complete! To enable WireGuard:"
            echo "   sudo cp $CONFIG_FILE /etc/wireguard/"
            echo "   sudo chmod 600 /etc/wireguard/${INTERFACE_NAME}.conf"
            echo "   sudo systemctl enable --now wg-quick@${INTERFACE_NAME}"
            echo ""
            print_status "To check status (requires root):"
            echo "   sudo wg show ${INTERFACE_NAME}"
            echo "   sudo systemctl status wg-quick@${INTERFACE_NAME}"
        fi
        echo ""
        echo -e "${RED}========================================${NC}"
        echo -e "${RED}  !!! NOW UPDATE ZION SERVER !!!  ${NC}"
        echo -e "${RED}========================================${NC}"
        echo ""
        print_warning "You MUST add this peer to Zion's config (/etc/wireguard/wg0.conf):"
        echo ""
        echo -e "${YELLOW}#$HOSTNAME${NC}"
        echo -e "${YELLOW}[Peer]${NC}"
        echo -e "${YELLOW}PublicKey = $PUBLIC_KEY${NC}"
        echo -e "${YELLOW}AllowedIPs = $IP_ADDRESS/32${NC}"
        echo ""
        print_warning "After updating Zion, restart its WireGuard:"
        echo "   systemctl restart wg-quick@wg0"
        echo ""
        
        # Add endpoint-specific instructions for full tunnel mode
        if [[ "$ROUTING_MODE" == "full_tunnel" ]]; then
            echo ""
            print_warning "FULL TUNNEL MODE DETECTED - Endpoint Changes Required:"
            echo ""
            echo "Since this node will route ALL traffic through the VPN, you need to:"
            echo ""
            echo "1. Update Zion's config (/etc/wireguard/wg0.conf) to allow this peer:"
            echo "   - Add the peer section as shown above"
            echo "   - Ensure Zion has proper iptables rules for NAT/masquerading"
            echo ""
            echo "2. Check Zion's iptables rules (run on Zion server):"
            echo "   sudo iptables -t nat -L POSTROUTING"
            echo "   sudo iptables -L FORWARD"
            echo ""
            echo "3. If Zion doesn't have proper NAT rules, add them:"
            echo "   sudo iptables -t nat -A POSTROUTING -s 10.8.0.0/24 -o eth0 -j MASQUERADE"
            echo "   sudo iptables -A FORWARD -i wg0 -j ACCEPT"
            echo "   sudo iptables -A FORWARD -o wg0 -j ACCEPT"
            echo ""
            echo "4. Enable IP forwarding on Zion (if not already enabled):"
            echo "   echo 'net.ipv4.ip_forward=1' | sudo tee -a /etc/sysctl.conf"
            echo "   sudo sysctl -p"
            echo ""
        fi
        
        if [[ "$RUNNING_AS_ROOT" == "true" ]]; then
            print_warning "Then restart this node's WireGuard:"
            echo "   systemctl restart wg-quick@${INTERFACE_NAME}"
        else
            print_warning "Then restart this node's WireGuard:"
            echo "   sudo systemctl restart wg-quick@${INTERFACE_NAME}"
        fi
    else
        echo -e "${BLUE}Step $([[ "$RUNNING_AS_ROOT" == "false" ]] && echo "5" || echo "4"): Next Steps${NC}"
        echo ""
        print_status "Keys generated successfully!"
        echo "   Private key: $WG_DIR/${HOSTNAME}_private.key"
        echo "   Public key: $WG_DIR/${HOSTNAME}_public.key"
        echo ""
        print_warning "Next steps:"
        echo "   - Create ${INTERFACE_NAME}.conf manually using the keys above"
        echo "   - Copy config to $CONFIG_FILE"
        if [[ "$RUNNING_AS_ROOT" == "true" ]]; then
            echo "   - Set permissions: chmod 600 $CONFIG_FILE"
            echo "   - Enable/start: systemctl enable --now wg-quick@${INTERFACE_NAME}"
        else
            echo "   - Copy to system: sudo cp $CONFIG_FILE /etc/wireguard/"
            echo "   - Set permissions: sudo chmod 600 /etc/wireguard/${INTERFACE_NAME}.conf"
            echo "   - Enable/start: sudo systemctl enable --now wg-quick@${INTERFACE_NAME}"
        fi
        echo ""
        echo -e "${RED}========================================${NC}"
        echo -e "${RED}  !!! NOW UPDATE ZION SERVER !!!  ${NC}"
        echo -e "${RED}========================================${NC}"
        echo ""
        print_warning "You MUST add this peer to Zion's config (/etc/wireguard/wg0.conf):"
        echo ""
        echo -e "${YELLOW}#$HOSTNAME${NC}"
        echo -e "${YELLOW}[Peer]${NC}"
        echo -e "${YELLOW}PublicKey = $PUBLIC_KEY${NC}"
        echo -e "${YELLOW}AllowedIPs = $IP_ADDRESS/32${NC}"
        echo ""
        print_warning "After updating Zion, restart its WireGuard:"
        echo "   systemctl restart wg-quick@wg0"
        echo ""
        
        # Add endpoint-specific instructions for full tunnel mode (manual path)
        if [[ "${ROUTING_MODE:-wg_only}" == "full_tunnel" ]]; then
            echo ""
            print_warning "FULL TUNNEL MODE DETECTED - Endpoint Changes Required:"
            echo ""
            echo "Since this node will route ALL traffic through the VPN, you need to:"
            echo ""
            echo "1. Update Zion's config (/etc/wireguard/wg0.conf) to allow this peer:"
            echo "   - Add the peer section as shown above"
            echo "   - Ensure Zion has proper iptables rules for NAT/masquerading"
            echo ""
            echo "2. Check Zion's iptables rules (run on Zion server):"
            echo "   sudo iptables -t nat -L POSTROUTING"
            echo "   sudo iptables -L FORWARD"
            echo ""
            echo "3. If Zion doesn't have proper NAT rules, add them:"
            echo "   sudo iptables -t nat -A POSTROUTING -s 10.8.0.0/24 -o eth0 -j MASQUERADE"
            echo "   sudo iptables -A FORWARD -i wg0 -j ACCEPT"
            echo "   sudo iptables -A FORWARD -o wg0 -j ACCEPT"
            echo ""
            echo "4. Enable IP forwarding on Zion (if not already enabled):"
            echo "   echo 'net.ipv4.ip_forward=1' | sudo tee -a /etc/sysctl.conf"
            echo "   sudo sysctl -p"
            echo ""
        fi
        
        if [[ "$RUNNING_AS_ROOT" == "true" ]]; then
            print_warning "Then restart this node's WireGuard:"
            echo "   systemctl restart wg-quick@${INTERFACE_NAME}"
        else
            print_warning "Then restart this node's WireGuard:"
            echo "   sudo systemctl restart wg-quick@${INTERFACE_NAME}"
        fi
    fi
    
    echo ""
    print_status "Setup complete for $HOSTNAME!"
    print_status "Log file available at: $log_file"
    echo ""
}

main "$@"