#!/usr/bin/env bash
# wireguard-status.sh -- JSON status for Waybar
# Outputs: {"text":"...","class":"connected|disconnected","tooltip":"..."}
#
# Waybar config example (in ~/.config/waybar/config.jsonc):
#   "custom/wireguard": {
#       "exec": "$HOME/.config/waybar/scripts/wireguard-status.sh",
#       "interval": 10,
#       "return-type": "json"
#   }

set -euo pipefail

# --- Config -------------------------------------------------------
ICON_LOCKED="\uf023"       # nf-fa-lock (FontAwesome)
ICON_UNLOCKED="\uf09e"     # nf-fa-unlock
# Override with environment variables if set:
ICON_LOCKED="${WG_ICON_LOCKED:-$ICON_LOCKED}"
ICON_UNLOCKED="${WG_ICON_UNLOCKED:-$ICON_UNLOCKED}"

# --- Helpers ------------------------------------------------------
die() { echo "$*" >&2; exit 1; }

# Find the wg binary — works on NixOS, standard Linux, etc.
find_wg() {
    if command -v wg &>/dev/null; then
        echo "wg"
    elif [ -x /run/current-system/sw/bin/wg ]; then
        echo "/run/current-system/sw/bin/wg"
    elif [ -x /usr/bin/wg ]; then
        echo "/usr/bin/wg"
    else
        return 1
    fi
}

# Find the ip binary
find_ip() {
    if command -v ip &>/dev/null; then
        echo "ip"
    elif [ -x /run/current-system/sw/bin/ip ]; then
        echo "/run/current-system/sw/bin/ip"
    elif [ -x /usr/sbin/ip ]; then
        echo "/usr/sbin/ip"
    elif [ -x /sbin/ip ]; then
        echo "/sbin/ip"
    else
        return 1
    fi
}

format_bytes() {
    local bytes=$1
    if [ "$bytes" -lt 1024 ]; then
        echo "${bytes}B"
    elif [ "$bytes" -lt $((1024 * 1024)) ]; then
        echo "$(awk "BEGIN { printf \"%.1f KiB\", $bytes / 1024 }")"
    elif [ "$bytes" -lt $((1024 * 1024 * 1024)) ]; then
        echo "$(awk "BEGIN { printf \"%.1f MiB\", $bytes / (1024 * 1024) }")"
    else
        echo "$(awk "BEGIN { printf \"%.2f GiB\", $bytes / (1024 * 1024 * 1024) }")"
    fi
}

# --- Main ---------------------------------------------------------

WG_BIN=$(find_wg 2>/dev/null) || {
    # wg not installed — report as disconnected
    printf '{"text":"%s","class":"disconnected","tooltip":"WireGuard not installed"}\n' "$ICON_UNLOCKED"
    exit 0
}

IP_BIN=$(find_ip 2>/dev/null) || true

# Find all WireGuard interfaces
INTERFACES=$("$WG_BIN" show interfaces 2>/dev/null) || {
    # No WireGuard interfaces exist (wg module not loaded?)
    printf '{"text":"%s","class":"disconnected","tooltip":"No WireGuard interfaces"}\n' "$ICON_UNLOCKED"
    exit 0
}

# Parse into array, handling possible single-line space-separated output
IFS=' ' read -ra ifaces <<< "$INTERFACES"
if [ ${#ifaces[@]} -eq 0 ]; then
    printf '{"text":"%s","class":"disconnected","tooltip":"No WireGuard interfaces"}\n' "$ICON_UNLOCKED"
    exit 0
fi

# Just use the first interface (most setups have one)
IFACE="${ifaces[0]}"

# Get transfer stats
TRANSFER=$("$WG_BIN" show "$IFACE" transfer 2>/dev/null) || {
    printf '{"text":"%s","class":"disconnected","tooltip":"Cannot read stats for %s"}\n' "$ICON_UNLOCKED" "$IFACE"
    exit 0
}

# Parse: "rx_bytes tx_bytes"
RX_BYTES=$(echo "$TRANSFER" | awk '{print $1}')
TX_BYTES=$(echo "$TRANSFER" | awk '{print $2}')

RX_HUMAN=$(format_bytes "${RX_BYTES:-0}")
TX_HUMAN=$(format_bytes "${TX_BYTES:-0}")

# Get endpoint info for tooltip
ENDPOINT=$("$WG_BIN" show "$IFACE" endpoints 2>/dev/null | head -1 | awk '{print $2}') || ENDPOINT=""
PEER=$("$WG_BIN" show "$IFACE" peers 2>/dev/null | head -1) || PEER=""

# Build status line
STATUS_TEXT="${ICON_LOCKED} \u2193${RX_HUMAN} \u2191${TX_HUMAN}"

# Build tooltip
TOOLTIP="${IFACE}"
if [ -n "$PEER" ]; then
    TOOLTIP="${TOOLTIP}\nPeer: ${PEER}"
fi
if [ -n "$ENDPOINT" ]; then
    TOOLTIP="${TOOLTIP}\nEndpoint: ${ENDPOINT}"
fi
# Add IP info if ip is available
if [ -n "$IP_BIN" ]; then
    IFACE_IP=$("$IP_BIN" -4 addr show dev "$IFACE" 2>/dev/null | awk '/inet / {print $2}') || IFACE_IP=""
    if [ -n "$IFACE_IP" ]; then
        TOOLTIP="${TOOLTIP}\nIP: ${IFACE_IP}"
    fi
fi
TOOLTIP="${TOOLTIP}\nRX: ${RX_HUMAN}\nTX: ${TX_HUMAN}"

printf '{"text":"%s","class":"connected","tooltip":"%s"}\n' "$STATUS_TEXT" "$TOOLTIP"