#!/usr/bin/env bash
# WireGuard key generation and setup helper
# Run this on BOTH the home server and the droplet to generate keys.
# Then copy the public keys into the appropriate config files.

set -euo pipefail

echo "=== Vigilar WireGuard Setup ==="
echo ""

# Check if WireGuard is installed
if ! command -v wg &>/dev/null; then
    echo "Installing WireGuard..."
    if command -v apt &>/dev/null; then
        sudo apt update && sudo apt install -y wireguard
    elif command -v pacman &>/dev/null; then
        sudo pacman -S --noconfirm wireguard-tools
    else
        echo "ERROR: Install WireGuard manually for your OS"
        exit 1
    fi
fi

echo "Generating WireGuard keys..."
PRIV_KEY=$(wg genkey)
PUB_KEY=$(echo "$PRIV_KEY" | wg pubkey)

echo ""
echo "Private Key: $PRIV_KEY"
echo "Public Key:  $PUB_KEY"
echo ""
echo "Save the private key in /etc/wireguard/ and share the PUBLIC key"
echo "with the other end of the tunnel."
echo ""

# Detect if this is the home server or droplet
read -p "Is this the (h)ome server or (d)roplet? [h/d]: " ROLE

if [[ "$ROLE" == "d" ]]; then
    echo ""
    echo "=== DROPLET SETUP ==="
    echo ""

    # Save keys
    sudo mkdir -p /etc/wireguard
    echo "$PRIV_KEY" | sudo tee /etc/wireguard/droplet_private.key > /dev/null
    echo "$PUB_KEY" | sudo tee /etc/wireguard/droplet_public.key > /dev/null
    sudo chmod 600 /etc/wireguard/droplet_private.key

    read -p "Home server's PUBLIC key: " HOME_PUB_KEY

    # Generate config
    sudo tee /etc/wireguard/wg0.conf > /dev/null <<EOF
[Interface]
Address = 10.99.0.1/32
ListenPort = 51820
PrivateKey = $PRIV_KEY

[Peer]
PublicKey = $HOME_PUB_KEY
AllowedIPs = 10.99.0.2/32
EOF

    sudo chmod 600 /etc/wireguard/wg0.conf

    # Open firewall
    if command -v ufw &>/dev/null; then
        sudo ufw allow 51820/udp
        sudo ufw allow 443/tcp
        sudo ufw allow 80/tcp
        echo "Firewall rules added (51820/udp, 80/tcp, 443/tcp)"
    fi

    # Enable and start
    sudo systemctl enable --now wg-quick@wg0
    echo ""
    echo "WireGuard started on droplet."
    echo "Droplet tunnel IP: 10.99.0.1"
    echo ""
    echo "Share this public key with your home server: $PUB_KEY"

elif [[ "$ROLE" == "h" ]]; then
    echo ""
    echo "=== HOME SERVER SETUP ==="
    echo ""

    # Save keys
    sudo mkdir -p /etc/wireguard
    echo "$PRIV_KEY" | sudo tee /etc/wireguard/home_private.key > /dev/null
    echo "$PUB_KEY" | sudo tee /etc/wireguard/home_public.key > /dev/null
    sudo chmod 600 /etc/wireguard/home_private.key

    read -p "Droplet's PUBLIC key: " DROPLET_PUB_KEY
    read -p "Droplet's public IP address: " DROPLET_IP

    # Generate config
    sudo tee /etc/wireguard/wg0.conf > /dev/null <<EOF
[Interface]
Address = 10.99.0.2/32
PrivateKey = $PRIV_KEY

[Peer]
PublicKey = $DROPLET_PUB_KEY
AllowedIPs = 10.99.0.1/32
Endpoint = ${DROPLET_IP}:51820
PersistentKeepalive = 25
EOF

    sudo chmod 600 /etc/wireguard/wg0.conf

    # Enable and start
    sudo systemctl enable --now wg-quick@wg0
    echo ""
    echo "WireGuard started on home server."
    echo "Home tunnel IP: 10.99.0.2"
    echo ""
    echo "Share this public key with your droplet: $PUB_KEY"
fi

echo ""
echo "Test connectivity with: ping 10.99.0.1 (from home) or ping 10.99.0.2 (from droplet)"