#!/bin/bash

set -e

WG_DIR="/etc/wireguard"
WG_CONF="$WG_DIR/wg0.conf"
INTERFACE="wg0"
SUBNET_PREFIX="10.0.0"
PORT=51820
DNS_SERVER="1.1.1.1"

function get_next_ip() {
    USED_IPS=$(grep AllowedIPs "$WG_CONF" | grep -oE "$SUBNET_PREFIX\.[0-9]+" | sort -n -t. -k4)
    NEXT_IP=2
    for ip in $USED_IPS; do
        last_octet=$(echo $ip | cut -d. -f4)
        if [[ $last_octet -ge $NEXT_IP ]]; then
            ((NEXT_IP=last_octet+1))
        fi
    done
    echo "$SUBNET_PREFIX.$NEXT_IP"
}

function add_client() {
    CLIENT_NAME="$1"
    CLIENT_DIR="$WG_DIR/$CLIENT_NAME"
    mkdir -p "$CLIENT_DIR"
    cd "$CLIENT_DIR"

    echo "[+] Generating keys for $CLIENT_NAME..."
    wg genkey | tee private | wg pubkey > public
    PRIVATE_KEY=$(cat private)
    PUBLIC_KEY=$(cat public)

    SERVER_PUBLIC_KEY=$(cat "$WG_DIR/self/public")
    SERVER_IP=$(curl -s ifconfig.me)
    CLIENT_IP=$(get_next_ip)

    echo "[+] Creating config for $CLIENT_NAME ($CLIENT_IP)..."
    cat > "$CLIENT_DIR/$CLIENT_NAME.conf" <<EOF
[Interface]
PrivateKey = $PRIVATE_KEY
Address = $CLIENT_IP/24
DNS = $DNS_SERVER

[Peer]
PublicKey = $SERVER_PUBLIC_KEY
Endpoint = $SERVER_IP:$PORT
AllowedIPs = 0.0.0.0/0, ::/0
PersistentKeepalive = 25
EOF

    echo "[+] Adding peer to running interface..."
    wg set $INTERFACE peer $PUBLIC_KEY allowed-ips $CLIENT_IP/32

    echo "[+] Persisting peer to $WG_CONF..."
    cat >> "$WG_CONF" <<EOF

[Peer]
PublicKey = $PUBLIC_KEY
AllowedIPs = $CLIENT_IP/32
EOF

    echo "[+] Generating QR code..."
    qrencode -o "$CLIENT_DIR/$CLIENT_NAME.png" < "$CLIENT_DIR/$CLIENT_NAME.conf"
    qrencode -t ansiutf8 < "$CLIENT_DIR/$CLIENT_NAME.conf"

    echo
    echo "✅ Added client: $CLIENT_NAME"
    echo "📄 Config: $CLIENT_DIR/$CLIENT_NAME.conf"
    echo "📸 QR Code: $CLIENT_DIR/$CLIENT_NAME.png"
    echo "📶 Assigned IP: $CLIENT_IP"
}

function remove_client() {
    CLIENT_NAME="$1"
    CLIENT_DIR="$WG_DIR/$CLIENT_NAME"

    if [ ! -d "$CLIENT_DIR" ]; then
        echo "❌ No such client: $CLIENT_NAME"
        exit 1
    fi

    echo "[+] Getting public key of $CLIENT_NAME..."
    PUBLIC_KEY=$(cat "$CLIENT_DIR/public")

    echo "[+] Removing peer from live interface..."
    wg set $INTERFACE peer $PUBLIC_KEY remove || true

    echo "[+] Removing peer from $WG_CONF..."
    awk -v pub="$PUBLIC_KEY" '
        BEGIN {skip=0}
        $0 ~ "\\[Peer\\]" {skip=0}
        $0 ~ "PublicKey = " pub {skip=1; next}
        skip && /^AllowedIPs/ {next}
        {print}
    ' "$WG_CONF" > "$WG_CONF.tmp" && mv "$WG_CONF.tmp" "$WG_CONF"

    echo "[+] Removing client files..."
    rm -rf "$CLIENT_DIR"

    echo "✅ Removed client: $CLIENT_NAME"
}

case "$1" in
  add)
    if [ -z "$2" ]; then echo "Usage: $0 add <client_name>"; exit 1; fi
    add_client "$2"
    ;;
  remove)
    if [ -z "$2" ]; then echo "Usage: $0 remove <client_name>"; exit 1; fi
    remove_client "$2"
    ;;
  *)
    echo "Usage: $0 {add|remove} <client_name>"
    exit 1
    ;;
esac

# bash wg_config.sh add mobile_2
# bash wg_config.sh remove mobile_2