#!/bin/bash
# Helper script to add a new WireGuard client
# Usage: ./add_client.sh <client_name> <ip_suffix>
# Example: ./add_client.sh mylaptop 4

set -e

CLIENT_NAME="$1"
IP_SUFFIX="$2"

if [ -z "$CLIENT_NAME" ] || [ -z "$IP_SUFFIX" ]; then
    echo "Usage: $0 <client_name> <ip_suffix>"
    echo "Example: $0 mylaptop 4"
    echo ""
    echo "This will:"
    echo "  - Generate WireGuard keys for the client"
    echo "  - Assign IP 10.8.0.$IP_SUFFIX to the client"
    echo "  - Add peer to WireGuard server config"
    echo "  - Generate client config file and QR codes"
    exit 1
fi

WG_DIR="/etc/wireguard"
KEYS_DIR="$WG_DIR/keys"
CLIENTS_DIR="$WG_DIR/clients"
NETDEV_FILE="/etc/systemd/network/99-wg0.netdev"

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    echo "Please run as root"
    exit 1
fi

# Check if client already exists
if [ -f "$KEYS_DIR/${CLIENT_NAME}.key" ]; then
    echo "Error: Client '$CLIENT_NAME' already exists!"
    exit 1
fi

# Check if IP is already in use
if grep -q "10.8.0.${IP_SUFFIX}/32" "$NETDEV_FILE"; then
    echo "Error: IP 10.8.0.${IP_SUFFIX} is already assigned!"
    echo "Choose a different IP suffix."
    exit 1
fi

echo "=========================================="
echo "Adding WireGuard Client: $CLIENT_NAME"
echo "=========================================="
echo "Client IP: 10.8.0.${IP_SUFFIX}"
echo ""

# Generate keys
echo "Generating keys..."
cd "$KEYS_DIR"
"$WG_DIR/generate_key.sh" "$CLIENT_NAME"

CLIENT_PRIVATE_KEY=$(cat "${KEYS_DIR}/${CLIENT_NAME}.key")
CLIENT_PUBLIC_KEY=$(cat "${KEYS_DIR}/${CLIENT_NAME}.pub")
SERVER_PUBLIC_KEY=$(cat "${KEYS_DIR}/server.pub")

echo "Keys generated successfully!"
echo "Public key: $CLIENT_PUBLIC_KEY"
echo ""

# Add peer to netdev file
echo "Adding peer to server configuration..."
cat >> "$NETDEV_FILE" << EOF

# $CLIENT_NAME
[WireGuardPeer]
PublicKey=$CLIENT_PUBLIC_KEY
AllowedIPs=10.8.0.${IP_SUFFIX}/32
PersistentKeepalive=25
EOF

echo "Peer added to $NETDEV_FILE"
echo ""

# Get server endpoint (try to detect public IP)
SERVER_ENDPOINT=$(curl -s ifconfig.me 2>/dev/null || echo "YOUR_PUBLIC_IP")
if [ "$SERVER_ENDPOINT" = "YOUR_PUBLIC_IP" ]; then
    # Fallback to ansible_host if available
    SERVER_ENDPOINT=$(hostname -I | awk '{print $1}')
fi

# Create client config
echo "Creating client configuration..."
cat > "$CLIENTS_DIR/${CLIENT_NAME}.conf" << EOF
[Interface]
Address = 10.8.0.${IP_SUFFIX}/24
PrivateKey = $CLIENT_PRIVATE_KEY
DNS = 10.8.0.1

[Peer]
PublicKey = $SERVER_PUBLIC_KEY
Endpoint = ${SERVER_ENDPOINT}:51820
# Route home network traffic through WireGuard
# To route ALL traffic (full VPN), change to: 0.0.0.0/0
AllowedIPs = 192.168.178.0/24, 10.8.0.0/24
PersistentKeepalive = 25
EOF

chmod 600 "$CLIENTS_DIR/${CLIENT_NAME}.conf"
echo "Client config created at $CLIENTS_DIR/${CLIENT_NAME}.conf"
echo ""

# Generate QR codes
echo "Generating QR codes..."
qrencode -t ansiutf8 -r "$CLIENTS_DIR/${CLIENT_NAME}.conf" > "$CLIENTS_DIR/${CLIENT_NAME}.qr.txt"
qrencode -t png -r "$CLIENTS_DIR/${CLIENT_NAME}.conf" -o "$CLIENTS_DIR/${CLIENT_NAME}.qr.png"
echo "QR codes generated!"
echo ""

# Reload systemd-networkd
echo "Reloading systemd-networkd..."
networkctl reload
sleep 2
echo ""

# Display summary
echo "=========================================="
echo "Client Added Successfully!"
echo "=========================================="
echo "Client name: $CLIENT_NAME"
echo "Client IP: 10.8.0.${IP_SUFFIX}"
echo "Public key: $CLIENT_PUBLIC_KEY"
echo ""
echo "Configuration files:"
echo "  - $CLIENTS_DIR/${CLIENT_NAME}.conf"
echo "  - $CLIENTS_DIR/${CLIENT_NAME}.qr.txt (terminal QR)"
echo "  - $CLIENTS_DIR/${CLIENT_NAME}.qr.png (image QR)"
echo ""
echo "To view QR code in terminal:"
echo "  cat $CLIENTS_DIR/${CLIENT_NAME}.qr.txt"
echo ""
echo "Next steps:"
echo "  1. Copy the client config to your device"
echo "  2. Import into WireGuard client app"
echo "  3. Connect and test with: ping 10.8.0.1"
echo "=========================================="