#!/usr/bin/env bash

# WireGuard Setup Script
# This script guides users through creating WireGuard keys and configuration files

set -euo pipefail  # Exit on error, undefined vars, pipe failures

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

# Configuration file for Zion settings
ZION_CONFIG_FILE="$(dirname "$0")/CURRENT_WORKING/zion.conf"

# Function to print colored output
print_status() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

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

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

print_header() {
    echo -e "${BLUE}================================${NC}"
    echo -e "${BLUE}$1${NC}"
    echo -e "${BLUE}================================${NC}"
}

# Function to check if command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Function to validate IP address with subnet
validate_ip() {
    local ip=$1
    
    # Check if input contains subnet
    if [[ ! $ip =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}/[0-9]{1,2}$ ]]; then
        return 1
    fi
    
    # Extract IP and subnet parts
    local ip_part="${ip%/*}"
    local subnet="${ip#*/}"
    
    # Validate subnet
    if [[ ! $subnet =~ ^[0-9]+$ ]] || [ "$subnet" -lt 1 ] || [ "$subnet" -gt 32 ]; then
        return 1
    fi
    
    # Validate each octet
    IFS='.' read -ra OCTETS <<< "$ip_part"
    if [ ${#OCTETS[@]} -ne 4 ]; then
        return 1
    fi
    
    for octet in "${OCTETS[@]}"; do
        if [[ ! $octet =~ ^[0-9]+$ ]] || [ "$octet" -lt 0 ] || [ "$octet" -gt 255 ]; then
            return 1
        fi
    done
    
    return 0
}

# Function to validate port number
validate_port() {
    local port=$1
    
    if [[ ! $port =~ ^[0-9]+$ ]]; then
        return 1
    fi
    
    if [ "$port" -lt 1 ] || [ "$port" -gt 65535 ]; then
        return 1
    fi
    
    # Check for common reserved ports
    if [ "$port" -le 1024 ]; then
        print_warning "Port $port is in privileged range (1-1024). You may need root access."
    fi
    
    return 0
}

# Function to validate WireGuard public key
validate_public_key() {
    local key=$1
    
    # WireGuard keys are base64 encoded and exactly 44 characters long
    if [[ ! $key =~ ^[A-Za-z0-9+/]{43}=$ ]]; then
        return 1
    fi
    
    return 0
}

# Function to detect network interface
detect_network_interface() {
    # Try to detect the main network interface
    local interface
    
    # Check for common interface names
    for name in eth0 ens33 ens160 enp0s3 eno1; do
        if ip link show "$name" >/dev/null 2>&1; then
            interface="$name"
            break
        fi
    done
    
    # Fallback to first non-loopback interface
    if [ -z "$interface" ]; then
        interface=$(ip route | grep default | awk '{print $5}' | head -1)
    fi
    
    echo "$interface"
}

# Function to load Zion configuration
load_zion_config() {
    if [[ -f "$ZION_CONFIG_FILE" ]]; then
        print_status "Found Zion configuration file: $ZION_CONFIG_FILE"
        return 0
    else
        print_warning "Zion configuration file not found: $ZION_CONFIG_FILE"
        print_warning "Using hardcoded Zion configuration"
        return 1
    fi
}

# Function to create safe filename
safe_filename() {
    local name="$1"
    # Replace invalid characters with underscores
    echo "$name" | sed 's/[^a-zA-Z0-9._-]/_/g'
}

# Main setup function
main() {
    print_header "WireGuard Configuration Setup"
    echo "This script will help you create WireGuard keys and configuration files."
    echo "Based on the CURRENT_WORKING configuration with Zion as central server."
    echo ""

    # Check if wg command exists
    if ! command_exists wg; then
        print_error "WireGuard tools not found. Please install WireGuard first:"
        echo "  Ubuntu/Debian: sudo apt install wireguard"
        echo "  CentOS/RHEL: sudo yum install wireguard-tools"
        echo "  Arch: sudo pacman -S wireguard-tools"
        echo "  Fedora: sudo dnf install wireguard-tools"
        exit 1
    fi

    # Create output directory
    OUTPUT_DIR="wireguard_configs"
    mkdir -p "$OUTPUT_DIR"
    print_status "Created output directory: $OUTPUT_DIR"

    # Step 1: Generate keys
    print_header "Step 1: Generate WireGuard Keys"
    
    echo "Generating private and public keys..."
    PRIVATE_KEY=$(wg genkey)
    PUBLIC_KEY=$(echo "$PRIVATE_KEY" | wg pubkey)
    
    print_status "Keys generated successfully!"
    echo "Private Key: $PRIVATE_KEY"
    echo "Public Key:  $PUBLIC_KEY"
    echo ""

    # Step 2: Get configuration details
    print_header "Step 2: Configuration Details"
    
    # Get node name
    local node_name
    while true; do
        read -p "Enter node name (e.g., aza, cth, galaxy): " node_name
        node_name=$(safe_filename "$node_name")
        
        if [ -n "$node_name" ]; then
            break
        else
            print_error "Node name cannot be empty"
        fi
    done
    
    # Check if configuration file already exists
    CONFIG_FILE="$OUTPUT_DIR/${node_name}.conf"
    if [ -f "$CONFIG_FILE" ]; then
        print_warning "Configuration file '$CONFIG_FILE' already exists."
        read -p "Do you want to overwrite it? (y/n): " OVERWRITE_FILE
        if [[ ! $OVERWRITE_FILE =~ ^[Yy]$ ]]; then
            print_error "Operation cancelled. Please choose a different node name or remove the existing file."
            exit 1
        fi
        print_status "Will overwrite existing configuration file."
    fi
    
    # Get IP address
    local ip_address
    while true; do
        read -p "Enter IP address with subnet (e.g., 10.8.0.2/24): " ip_address
        if validate_ip "$ip_address"; then
            break
        else
            print_error "Invalid IP address format. Use format: x.x.x.x/y"
        fi
    done
    
    # Ask if this is a server (listening) or client
    read -p "Is this a server that will listen for connections? (y/n): " IS_SERVER
    local listen_port
    if [[ $IS_SERVER =~ ^[Yy]$ ]]; then
        while true; do
            read -p "Enter listen port (1-65535): " listen_port
            if validate_port "$listen_port"; then
                break
            else
                print_error "Invalid port number. Must be between 1-65535"
            fi
        done
    fi
    
    # Step 3: Create configuration file
    print_header "Step 3: Create Configuration File"
    
    # Create the configuration file
    cat > "$CONFIG_FILE" << EOF
[Interface]
PrivateKey = $PRIVATE_KEY
Address = $ip_address
EOF

    # Add listen port if it's a server
    if [[ $IS_SERVER =~ ^[Yy]$ ]]; then
        echo "ListenPort = $listen_port" >> "$CONFIG_FILE"
    fi
    
    # Add PostUp/PostDown rules for server (Zion-like configuration)
    if [[ $IS_SERVER =~ ^[Yy]$ ]]; then
        local network_interface
        network_interface=$(detect_network_interface)
        
        if [ -n "$network_interface" ]; then
            print_status "Detected network interface: $network_interface"
        else
            print_warning "Could not detect network interface, using 'eth0'"
            network_interface="eth0"
        fi
        
        echo "" >> "$CONFIG_FILE"
        echo "# Server configuration - enable IP forwarding" >> "$CONFIG_FILE"
        echo "PostUp = iptables -A FORWARD -i wg0 -j ACCEPT; iptables -t nat -A POSTROUTING -o $network_interface -j MASQUERADE; ip route add 10.8.0.0/24 dev wg0 2>/dev/null || true" >> "$CONFIG_FILE"
        echo "PostDown = iptables -D FORWARD -i wg0 -j ACCEPT; iptables -t nat -D POSTROUTING -o $network_interface -j MASQUERADE; ip route del 10.8.0.0/24 dev wg0 2>/dev/null || true" >> "$CONFIG_FILE"
    fi
    
    print_status "Configuration file created: $CONFIG_FILE"
    
    # Step 4: Add peers
    print_header "Step 4: Add Peers"
    
    # Zion default peer information (from CURRENT_WORKING/zion.conf)
    local ZION_PUBLIC_KEY="2ztJbrN1x1NWanzPGLiKL19ZkdOhm5Y7WeKEWBT5cyg="
    local ZION_ENDPOINT="ugh.im:51820"
    local ZION_ALLOWED_IPS="10.8.0.0/24"
    
    # Try to load Zion config, fall back to hardcoded if not available
    if ! load_zion_config; then
        print_warning "Using hardcoded Zion configuration"
    fi
    
    # Ask if user wants to add Zion as default peer (unless this IS Zion)
    if [[ "$node_name" != "zion" ]]; then
        echo "Zion is the central server for this network."
        read -p "Do you want to add Zion as a peer? (y/n, default: y): " ADD_ZION
        ADD_ZION=${ADD_ZION:-y}
        
        if [[ $ADD_ZION =~ ^[Yy]$ ]]; then
            echo ""
            echo "Adding Zion as peer with default settings:"
            echo "  Public Key: $ZION_PUBLIC_KEY"
            echo "  Endpoint: $ZION_ENDPOINT"
            echo "  Allowed IPs: $ZION_ALLOWED_IPS"
            echo "  Persistent Keepalive: 25"
            echo ""
            
            # Add Zion peer to configuration
            echo "" >> "$CONFIG_FILE"
            echo "# Zion (central server)" >> "$CONFIG_FILE"
            echo "[Peer]" >> "$CONFIG_FILE"
            echo "PublicKey = $ZION_PUBLIC_KEY" >> "$CONFIG_FILE"
            echo "AllowedIPs = $ZION_ALLOWED_IPS" >> "$CONFIG_FILE"
            echo "Endpoint = $ZION_ENDPOINT" >> "$CONFIG_FILE"
            echo "PersistentKeepalive = 25" >> "$CONFIG_FILE"
            
            print_status "Added Zion as default peer"
        fi
    fi
    
    # Add additional peers
    while true; do
        read -p "Do you want to add additional peers? (y/n): " ADD_PEER
        if [[ ! $ADD_PEER =~ ^[Yy]$ ]]; then
            break
        fi
        
        echo ""
        local peer_name
        read -p "Enter peer name (e.g., aza, cth, galaxy): " peer_name
        peer_name=$(safe_filename "$peer_name")
        
        local peer_public_key
        while true; do
            read -p "Enter peer public key: " peer_public_key
            if validate_public_key "$peer_public_key"; then
                break
            else
                print_error "Invalid WireGuard public key format"
            fi
        done
        
        local peer_allowed_ips
        while true; do
            read -p "Enter allowed IPs for this peer (e.g., 10.8.0.1/32 or 10.8.0.0/24): " peer_allowed_ips
            if validate_ip "$peer_allowed_ips"; then
                break
            else
                print_error "Invalid IP address format. Use format: x.x.x.x/y"
            fi
        done
        
        # Ask if this peer has an endpoint (for clients connecting to servers)
        read -p "Does this peer have an endpoint? (y/n): " HAS_ENDPOINT
        if [[ $HAS_ENDPOINT =~ ^[Yy]$ ]]; then
            read -p "Enter endpoint (host:port): " PEER_ENDPOINT
            local peer_keepalive
            read -p "Enter persistent keepalive (seconds, default 25): " peer_keepalive
            peer_keepalive=${peer_keepalive:-25}
        fi
        
        # Add peer to configuration
        echo "" >> "$CONFIG_FILE"
        echo "# $peer_name" >> "$CONFIG_FILE"
        echo "[Peer]" >> "$CONFIG_FILE"
        echo "PublicKey = $peer_public_key" >> "$CONFIG_FILE"
        echo "AllowedIPs = $peer_allowed_ips" >> "$CONFIG_FILE"
        
        if [[ $HAS_ENDPOINT =~ ^[Yy]$ ]]; then
            echo "Endpoint = $PEER_ENDPOINT" >> "$CONFIG_FILE"
            echo "PersistentKeepalive = $peer_keepalive" >> "$CONFIG_FILE"
        fi
        
        print_status "Added peer: $peer_name"
    done
    
    # Set proper file permissions
    chmod 600 "$CONFIG_FILE"
    print_status "Set file permissions to 600 (owner read/write only)"
    
    # Step 5: Display results
    print_header "Step 5: Setup Complete"
    
    print_status "Configuration file created: $CONFIG_FILE"
    print_status "Your public key: $PUBLIC_KEY"
    echo ""
    
    # Zion update instructions
    print_header "IMPORTANT: Update Zion Server Configuration"
    echo ""
    print_warning "You MUST add this peer to Zion's configuration file:"
    echo "   /etc/wireguard/wg0.conf"
    echo ""
    echo "Add the following peer section to Zion's config:"
    echo "----------------------------------------"
    echo "# $node_name"
    echo "[Peer]"
    echo "PublicKey = $PUBLIC_KEY"
    echo "AllowedIPs = ${ip_address%/*}/32"
    echo "----------------------------------------"
    echo ""
    
    # Show Zion's current configuration structure
    if ! load_zion_config; then
        echo "Zion's current configuration structure:"
        echo "----------------------------------------"
        echo "[Interface]"
        echo "Address = 10.8.0.1/24"
        echo "ListenPort = 51820"
        echo "PrivateKey = <zion_private_key>"
        echo "PostUp = <iptables_rules>"
        echo "PostDown = <iptables_rules>"
        echo ""
        echo "#Cth"
        echo "[Peer]"
        echo "PublicKey = NBktXKy1s0n2lIlIMODvOqKNwAtYdoZH5feKt5P43i0="
        echo "AllowedIPs = 10.8.0.10/32"
        echo ""
        echo "#Aza"
        echo "[Peer]"
        echo "PublicKey = qmTKA257DLOrfhk5Zw8RyRmBSonmm6epbloT0P0ZWDc="
        echo "AllowedIPs = 10.8.0.2/32"
        echo ""
        echo "# Add your peer section here:"
        echo "# $node_name"
        echo "# [Peer]"
        echo "# PublicKey = $PUBLIC_KEY"
        echo "# AllowedIPs = ${ip_address%/*}/32"
        echo "----------------------------------------"
    fi
    
    echo ""
    print_warning "After updating Zion's config:"
    echo "1. Save the file"
    echo "2. Restart Zion's WireGuard: sudo systemctl restart wg-quick@wg0"
    echo "3. Then start this node's WireGuard: sudo wg-quick up ${node_name}"
    echo ""
    
    # Show the configuration file
    echo "Your configuration file contents:"
    echo "----------------------------------------"
    cat "$CONFIG_FILE"
    echo "----------------------------------------"
    
    # Create a summary file
    SUMMARY_FILE="$OUTPUT_DIR/${node_name}_summary.txt"
    if [ -f "$SUMMARY_FILE" ]; then
        print_warning "Summary file '$SUMMARY_FILE' already exists and will be overwritten."
    fi
    cat > "$SUMMARY_FILE" << EOF
WireGuard Configuration Summary for $node_name
===============================================

Node Name: $node_name
IP Address: $ip_address
$(if [[ $IS_SERVER =~ ^[Yy]$ ]]; then echo "Listen Port: $listen_port"; fi)

Your Public Key: $PUBLIC_KEY
Your Private Key: $PRIVATE_KEY

Configuration File: $CONFIG_FILE

Commands to use:
- Start: sudo wg-quick up ${node_name}
- Stop: sudo wg-quick down ${node_name}
- Status: sudo wg show

ZION SERVER UPDATE REQUIRED:
===========================
You MUST add this peer to Zion's configuration file (/etc/wireguard/wg0.conf):

# $node_name
[Peer]
PublicKey = $PUBLIC_KEY
AllowedIPs = ${ip_address%/*}/32

After updating Zion's config:
1. Save the file
2. Restart Zion's WireGuard: sudo systemctl restart wg-quick@wg0
3. Then start this node's WireGuard: sudo wg-quick up ${node_name}

Remember to:
1. Keep your private key secure
2. Update Zion's configuration with your public key
3. Set proper file permissions: chmod 600 $CONFIG_FILE
4. Restart both Zion and this node after configuration changes
EOF
    
    print_status "Summary saved to: $SUMMARY_FILE"
}

# Function to show usage
show_usage() {
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  -h, --help     Show this help message"
    echo "  -d, --dir DIR  Set output directory (default: wireguard_configs)"
    echo "  -c, --config   Use custom Zion config file"
    echo ""
    echo "This script will interactively guide you through creating WireGuard"
    echo "keys and configuration files."
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            show_usage
            exit 0
            ;;
        -d|--dir)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        -c|--config)
            ZION_CONFIG_FILE="$2"
            shift 2
            ;;
        *)
            print_error "Unknown option: $1"
            show_usage
            exit 1
            ;;
    esac
done

# Run main function
main "$@"