#!/bin/bash
# setup-wireguard.sh — Set up WireGuard VPN to HQ on a new Mac
# Usage: sudo ./scripts/setup-wireguard.sh <PRIVATE_KEY> <IP>
#
# Example:
#   sudo ./scripts/setup-wireguard.sh "YOUR_PRIVATE_KEY" "192.168.8.2"
#
# Peers (assigned on UDM Pro SE VPN server, port 51822):
#   192.168.8.2 — Robert Laptop
#
# Prerequisites: brew install wireguard-tools
# Note: Requires bash 4+ (macOS ships with bash 3; Homebrew's bash works)

set -euo pipefail

PRIVATE_KEY="${1:?Usage: sudo $0 <PRIVATE_KEY> <IP>}"
ADDRESS="${2:?Usage: sudo $0 <PRIVATE_KEY> <IP>}"
WG_PUBKEY="BAs+URnA37pXXUNZs7VEDAGxxEGfmnIuFQ0doVx5tWQ="
WG_ENDPOINT="147.160.62.196:51822"
USERNAME="$(logname)"
WG_BIN="$(dirname "$(command -v wg-quick 2>/dev/null || echo /opt/homebrew/bin/wg-quick)")"
BASH_BIN="$WG_BIN/bash"

echo "=== WireGuard HQ VPN Setup ==="
echo "Peer IP: $ADDRESS"
echo "User: $USERNAME"
echo "WG binaries: $WG_BIN"
echo ""

# Step 1: Install wireguard-tools if missing
if ! command -v wg &>/dev/null; then
    echo "[1/6] Installing wireguard-tools..."
    sudo -u "$USERNAME" brew install wireguard-tools
    WG_BIN="$(dirname "$(command -v wg-quick)")"
    BASH_BIN="$WG_BIN/bash"
else
    echo "[1/6] wireguard-tools already installed"
fi

WG_ETC="$(dirname "$WG_BIN")/etc/wireguard"

# Step 2: Write WireGuard config
echo "[2/6] Writing $WG_ETC/hq.conf..."
mkdir -p "$WG_ETC"
cat > "$WG_ETC/hq.conf" << EOF
[Interface]
PrivateKey = $PRIVATE_KEY
Address = $ADDRESS/32

[Peer]
PublicKey = $WG_PUBKEY
Endpoint = $WG_ENDPOINT
AllowedIPs = 192.168.0.0/22, 192.168.8.0/24
PersistentKeepalive = 25
EOF
chmod 600 "$WG_ETC/hq.conf"

# Step 3: Sudoers for passwordless wg commands
echo "[3/6] Configuring sudoers..."
cat > /etc/sudoers.d/wireguard << EOF
$USERNAME ALL=(ALL) NOPASSWD: $WG_BIN/wg, $WG_BIN/wg-quick, $BASH_BIN $WG_BIN/wg-quick *
EOF
chmod 440 /etc/sudoers.d/wireguard

# Step 4: Install wg-auto script
echo "[4/6] Installing /usr/local/bin/wg-auto..."
cat > /usr/local/bin/wg-auto << 'SCRIPT'
#!/bin/bash
WG_BIN="__WG_BIN__"
BASH_BIN="__BASH_BIN__"
WG_PEER="__WG_PEER__"
WG_ENDPOINT="__WG_ENDPOINT__"
WG_IFACE="hq"
OFFICE_GW="192.168.3.1"
LOG="/tmp/wg-auto.log"

log() { echo "$(date '+%Y-%m-%d %H:%M:%S') $1" >> "$LOG"; }

GW=$(route -n get default 2>/dev/null | awk '/gateway:/{print $2}')
SSID=$(networksetup -getairportnetwork en0 2>/dev/null | sed 's/^Current Wi-Fi Network: //')
[[ "$SSID" == *"not associated"* || "$SSID" == *"Error"* ]] && SSID=""
NET_LABEL="${SSID:-ethernet}"

WG_UP=0
sudo "$WG_BIN/wg" show "$WG_IFACE" >/dev/null 2>&1 && WG_UP=1

if [[ -z "$GW" ]]; then
    log "No network detected (no gateway) — no action"
elif [[ "$GW" == "$OFFICE_GW" ]]; then
    if [[ $WG_UP -eq 1 ]]; then
        log "At office ($NET_LABEL, gw=$GW) — tearing down WireGuard"
        sudo "$BASH_BIN" "$WG_BIN/wg-quick" down "$WG_IFACE" 2>/dev/null
    else
        log "At office ($NET_LABEL, gw=$GW) — WireGuard already off"
    fi
else
    if [[ $WG_UP -eq 0 ]]; then
        log "Remote ($NET_LABEL, gw=$GW) — bringing up WireGuard"
        sudo "$BASH_BIN" "$WG_BIN/wg-quick" up "$WG_IFACE" 2>/dev/null
        IFACE=$(sudo "$WG_BIN/wg" show interfaces 2>/dev/null)
        sudo "$WG_BIN/wg" set "$IFACE" peer "$WG_PEER" endpoint "$WG_ENDPOINT"
        log "WireGuard up, endpoint forced to $WG_ENDPOINT"
    else
        # Re-force endpoint in case it drifted (e.g. after network change)
        IFACE=$(sudo "$WG_BIN/wg" show interfaces 2>/dev/null)
        CURRENT_EP=$(sudo "$WG_BIN/wg" show "$IFACE" endpoints 2>/dev/null | awk '{print $2}')
        if [[ "$CURRENT_EP" != "$WG_ENDPOINT" ]]; then
            log "Remote ($NET_LABEL, gw=$GW) — endpoint drifted to $CURRENT_EP, re-forcing to $WG_ENDPOINT"
            sudo "$BASH_BIN" "$WG_BIN/wg-quick" down "$WG_IFACE" 2>/dev/null
            sudo "$BASH_BIN" "$WG_BIN/wg-quick" up "$WG_IFACE" 2>/dev/null
            IFACE=$(sudo "$WG_BIN/wg" show interfaces 2>/dev/null)
            sudo "$WG_BIN/wg" set "$IFACE" peer "$WG_PEER" endpoint "$WG_ENDPOINT"
            log "WireGuard bounced, endpoint forced to $WG_ENDPOINT"
        else
            log "Remote ($NET_LABEL, gw=$GW) — WireGuard already up"
        fi
    fi
fi
SCRIPT
sed -i '' "s|__WG_BIN__|$WG_BIN|g" /usr/local/bin/wg-auto
sed -i '' "s|__BASH_BIN__|$BASH_BIN|g" /usr/local/bin/wg-auto
sed -i '' "s|__WG_PEER__|$WG_PUBKEY|g" /usr/local/bin/wg-auto
sed -i '' "s|__WG_ENDPOINT__|$WG_ENDPOINT|g" /usr/local/bin/wg-auto
chmod +x /usr/local/bin/wg-auto

# Step 5: Install shell aliases
echo "[5/6] Installing shell aliases..."
ALIAS_FILE="/Users/$USERNAME/.wg-aliases"
cat > "$ALIAS_FILE" << 'ALIASES'
alias wg-up='sudo /opt/homebrew/bin/bash /opt/homebrew/bin/wg-quick up hq'
alias wg-down='sudo /opt/homebrew/bin/bash /opt/homebrew/bin/wg-quick down hq'
alias wg-status='sudo /opt/homebrew/bin/wg show'
alias wg-bounce='wg-down 2>/dev/null; wg-up'
ALIASES
chown "$USERNAME" "$ALIAS_FILE"

# Add source to shell profiles if not already present
for RC in "/Users/$USERNAME/.bashrc" "/Users/$USERNAME/.zshrc"; do
    if [ -f "$RC" ] && ! grep -q '.wg-aliases' "$RC" 2>/dev/null; then
        echo '[ -f ~/.wg-aliases ] && source ~/.wg-aliases' >> "$RC"
    elif [ ! -f "$RC" ]; then
        echo '[ -f ~/.wg-aliases ] && source ~/.wg-aliases' > "$RC"
        chown "$USERNAME" "$RC"
    fi
done

# Step 6: Install LaunchAgent for auto-toggle on network change
echo "[6/6] Installing LaunchAgent..."
AGENT_DIR="/Users/$USERNAME/Library/LaunchAgents"
mkdir -p "$AGENT_DIR"
cat > "$AGENT_DIR/com.tts.wg-auto.plist" << EOF
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
    <key>Label</key>
    <string>com.tts.wg-auto</string>
    <key>ProgramArguments</key>
    <array>
        <string>/usr/local/bin/wg-auto</string>
    </array>
    <key>WatchPaths</key>
    <array>
        <string>/Library/Preferences/SystemConfiguration/com.apple.airport.preferences.plist</string>
        <string>/Library/Preferences/SystemConfiguration/NetworkInterfaces.plist</string>
        <string>/Library/Preferences/SystemConfiguration/com.apple.network.identification.plist</string>
    </array>
    <key>ThrottleInterval</key>
    <integer>5</integer>
    <key>StandardOutPath</key>
    <string>/tmp/wg-auto-stdout.log</string>
    <key>StandardErrorPath</key>
    <string>/tmp/wg-auto-stderr.log</string>
</dict>
</plist>
EOF
chown "$USERNAME" "$AGENT_DIR/com.tts.wg-auto.plist"

# Load the agent
sudo -u "$USERNAME" launchctl unload "$AGENT_DIR/com.tts.wg-auto.plist" 2>/dev/null || true
sudo -u "$USERNAME" launchctl load "$AGENT_DIR/com.tts.wg-auto.plist" 2>/dev/null || true

echo ""
echo "=== Setup complete ==="
echo "Commands: wg-up, wg-down, wg-status, wg-bounce"
echo "Auto-toggle: LaunchAgent watches network changes"
echo "Log: cat /tmp/wg-auto.log"
echo "Verify: ping -c 3 192.168.0.1"