#!/bin/bash
# WireGuard Client Configuration Generator
# Purpose: Generate client configs with QR codes for easy mobile import
# Usage: ./generate-client-config.sh <client-name>

set -euo pipefail

# ========================================
# Configuration
# ========================================

WG_CONFIG_DIR="/etc/wireguard"
CLIENT_CONFIG_DIR="$(dirname "$0")/../wireguard/configs"
WG_INTERFACE="wg0"
WG_SERVER_CONFIG="${WG_CONFIG_DIR}/${WG_INTERFACE}.conf"
WG_NETWORK="10.8.0.0/24"
WG_SERVER_IP="10.8.0.1"
WG_PORT="51820"

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

# ========================================
# Helper Functions
# ========================================

print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

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

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

check_root() {
    if [ "$EUID" -ne 0 ]; then
        print_error "This script must be run as root (for server config modifications)"
        exit 1
    fi
}

check_dependencies() {
    local deps=("wg" "wg-quick" "qrencode")
    local missing=()

    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            missing+=("$dep")
        fi
    done

    if [ ${#missing[@]} -ne 0 ]; then
        print_error "Missing dependencies: ${missing[*]}"
        print_info "Install with: apt install wireguard wireguard-tools qrencode"
        exit 1
    fi
}

get_next_client_ip() {
    # Find highest used IP in last octet and add 1
    if [ ! -f "$WG_SERVER_CONFIG" ]; then
        echo "10.8.0.2"
        return
    fi

    local last_octet=$(grep -oP 'AllowedIPs\s*=\s*10\.8\.0\.\K\d+' "$WG_SERVER_CONFIG" 2>/dev/null | sort -n | tail -1)

    if [ -z "$last_octet" ]; then
        echo "10.8.0.2"
    else
        echo "10.8.0.$((last_octet + 1))"
    fi
}

get_server_public_key() {
    if [ ! -f "${WG_CONFIG_DIR}/server_public.key" ]; then
        print_error "Server public key not found at ${WG_CONFIG_DIR}/server_public.key"
        print_info "Run the Ansible playbook first: ansible-playbook setup-wireguard-host.yml"
        exit 1
    fi

    cat "${WG_CONFIG_DIR}/server_public.key"
}

get_server_endpoint() {
    # Try to detect public IP
    local public_ip
    public_ip=$(curl -s -4 ifconfig.me 2>/dev/null || curl -s -4 icanhazip.com 2>/dev/null || echo "YOUR_SERVER_IP")

    echo "${public_ip}:${WG_PORT}"
}

# ========================================
# Main Script
# ========================================

main() {
    print_info "WireGuard Client Configuration Generator"
    echo ""

    # Validate input
    if [ $# -ne 1 ]; then
        print_error "Usage: $0 <client-name>"
        echo ""
        echo "Example:"
        echo "  $0 michael-laptop"
        echo "  $0 iphone"
        echo "  $0 office-desktop"
        exit 1
    fi

    local client_name="$1"

    # Validate client name (alphanumeric + dash/underscore only)
    if ! [[ "$client_name" =~ ^[a-zA-Z0-9_-]+$ ]]; then
        print_error "Client name must contain only alphanumeric characters, dashes, and underscores"
        exit 1
    fi

    # Pre-flight checks
    check_root
    check_dependencies

    # Create client config directory
    mkdir -p "$CLIENT_CONFIG_DIR"

    # Check if client already exists
    if [ -f "${CLIENT_CONFIG_DIR}/${client_name}.conf" ]; then
        print_warning "Client config for '${client_name}' already exists"
        read -p "Regenerate? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            print_info "Aborted"
            exit 0
        fi
    fi

    print_info "Generating configuration for client: ${client_name}"
    echo ""

    # Generate client keys
    print_info "Generating client keys..."
    local client_private_key=$(wg genkey)
    local client_public_key=$(echo "$client_private_key" | wg pubkey)
    local client_preshared_key=$(wg genpsk)

    # Get server information
    print_info "Reading server configuration..."
    local server_public_key=$(get_server_public_key)
    local server_endpoint=$(get_server_endpoint)

    # Assign client IP
    local client_ip=$(get_next_client_ip)
    print_info "Assigned client IP: ${client_ip}"

    # Create client config file
    print_info "Creating client configuration file..."
    cat > "${CLIENT_CONFIG_DIR}/${client_name}.conf" <<EOF
[Interface]
# Client: ${client_name}
# Generated: $(date)
PrivateKey = ${client_private_key}
Address = ${client_ip}/32

# DNS: Use Cloudflare/Quad9 (change if needed)
DNS = 1.1.1.1, 9.9.9.9

[Peer]
# Server
PublicKey = ${server_public_key}
PresharedKey = ${client_preshared_key}
Endpoint = ${server_endpoint}

# Route only VPN network through tunnel (split-tunnel)
AllowedIPs = ${WG_NETWORK}

# Keep connection alive (NAT traversal)
PersistentKeepalive = 25
EOF

    chmod 600 "${CLIENT_CONFIG_DIR}/${client_name}.conf"

    # Add peer to server configuration
    print_info "Adding peer to server configuration..."

    # Check if peer already exists in server config
    if grep -q "# ${client_name}" "$WG_SERVER_CONFIG" 2>/dev/null; then
        print_info "Removing old peer entry..."
        # Remove old peer entry (from comment to next empty line or end of file)
        sed -i "/# ${client_name}/,/^$/d" "$WG_SERVER_CONFIG"
    fi

    # Append new peer
    cat >> "$WG_SERVER_CONFIG" <<EOF

[Peer]
# ${client_name}
PublicKey = ${client_public_key}
PresharedKey = ${client_preshared_key}
AllowedIPs = ${client_ip}/32
PersistentKeepalive = 25
EOF

    # Reload WireGuard to apply changes
    print_info "Reloading WireGuard configuration..."
    systemctl reload wg-quick@${WG_INTERFACE}

    # Verify peer is active
    sleep 1
    if wg show ${WG_INTERFACE} | grep -q "$client_public_key"; then
        print_success "Peer successfully added to server"
    else
        print_warning "Peer added to config but not yet active (will activate when client connects)"
    fi

    # Generate QR code for mobile devices
    print_info "Generating QR code for mobile import..."
    qrencode -t ansiutf8 < "${CLIENT_CONFIG_DIR}/${client_name}.conf" > "${CLIENT_CONFIG_DIR}/${client_name}.qr.txt"
    qrencode -t png -o "${CLIENT_CONFIG_DIR}/${client_name}.qr.png" < "${CLIENT_CONFIG_DIR}/${client_name}.conf"

    # Display success summary
    echo ""
    print_success "=========================================="
    print_success "Client Configuration Created!"
    print_success "=========================================="
    echo ""
    echo "Client Name:      ${client_name}"
    echo "Client IP:        ${client_ip}"
    echo "Config File:      ${CLIENT_CONFIG_DIR}/${client_name}.conf"
    echo "QR Code (text):   ${CLIENT_CONFIG_DIR}/${client_name}.qr.txt"
    echo "QR Code (PNG):    ${CLIENT_CONFIG_DIR}/${client_name}.qr.png"
    echo ""
    echo "Server Endpoint:  ${server_endpoint}"
    echo "VPN Network:      ${WG_NETWORK}"
    echo ""
    print_info "=========================================="
    print_info "Import Instructions:"
    print_info "=========================================="
    echo ""
    echo "Desktop (Linux/macOS):"
    echo "  sudo cp ${CLIENT_CONFIG_DIR}/${client_name}.conf /etc/wireguard/"
    echo "  sudo wg-quick up ${client_name}"
    echo ""
    echo "Desktop (Windows):"
    echo "  1. Open WireGuard GUI"
    echo "  2. Click 'Import tunnel(s) from file'"
    echo "  3. Select: ${CLIENT_CONFIG_DIR}/${client_name}.conf"
    echo ""
    echo "Mobile (iOS/Android):"
    echo "  1. Open WireGuard app"
    echo "  2. Tap '+' > 'Create from QR code'"
    echo "  3. Scan QR code below or from: ${CLIENT_CONFIG_DIR}/${client_name}.qr.png"
    echo ""
    print_info "QR Code (scan with phone):"
    echo ""
    cat "${CLIENT_CONFIG_DIR}/${client_name}.qr.txt"
    echo ""
    print_info "=========================================="
    print_info "Verify Connection:"
    print_info "=========================================="
    echo ""
    echo "After connecting:"
    echo "  ping ${WG_SERVER_IP}"
    echo "  curl -k https://${WG_SERVER_IP}:8080  # Traefik Dashboard"
    echo ""
    print_success "Configuration complete! Client is ready to connect."
}

# Run main function
main "$@"