#!/usr/bin/env bash
# =============================================================================
# NetBird External Relay Server Setup Script
# Interactively generates relay.env and docker-compose.yml for one or more
# relay servers. Run this script ON each relay server, or use --dry-run to
# preview the generated files locally.
# =============================================================================

set -euo pipefail

# ── Colours ──────────────────────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; BOLD='\033[1m'; RESET='\033[0m'

info()    { echo -e "${CYAN}[INFO]${RESET}  $*"; }
success() { echo -e "${GREEN}[OK]${RESET}    $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET}  $*"; }
error()   { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
header()  { echo -e "\n${BOLD}${CYAN}══ $* ══${RESET}"; }

# ── Helpers ───────────────────────────────────────────────────────────────────
ask() {
    # ask <var_name> <prompt> [default]
    local var="$1" prompt="$2" default="${3:-}"
    local display_default=""
    [[ -n "$default" ]] && display_default=" [${default}]"
    while true; do
        read -rp "$(echo -e "${BOLD}${prompt}${display_default}: ${RESET}")" value
        value="${value:-$default}"
        if [[ -n "$value" ]]; then
            printf -v "$var" '%s' "$value"
            return
        fi
        warn "This field is required."
    done
}

ask_yn() {
    # ask_yn <prompt> <default: y|n>  → returns 0=yes 1=no
    local prompt="$1" default="${2:-y}"
    local opts="[Y/n]"; [[ "$default" == "n" ]] && opts="[y/N]"
    read -rp "$(echo -e "${BOLD}${prompt} ${opts}: ${RESET}")" reply
    reply="${reply:-$default}"
    [[ "${reply,,}" == "y" ]]
}

ask_secret() {
    local var="$1" prompt="$2"
    while true; do
        read -rsp "$(echo -e "${BOLD}${prompt}: ${RESET}")" value; echo
        if [[ -n "$value" ]]; then
            printf -v "$var" '%s' "$value"
            return
        fi
        warn "Secret cannot be empty."
    done
}

generate_secret() {
    if command -v openssl &>/dev/null; then
        openssl rand -base64 32
    else
        head -c 32 /dev/urandom | base64
    fi
}

validate_domain() {
    # Very basic domain sanity check
    [[ "$1" =~ ^[a-zA-Z0-9]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?(\.[a-zA-Z0-9]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?)*\.[a-zA-Z]{2,}$ ]]
}

# ── Arg parsing ───────────────────────────────────────────────────────────────
DRY_RUN=false
OUTPUT_DIR="."
for arg in "$@"; do
    case "$arg" in
        --dry-run)   DRY_RUN=true ;;
        --output=*)  OUTPUT_DIR="${arg#--output=}" ;;
        -h|--help)
            echo "Usage: $0 [--dry-run] [--output=<dir>]"
            echo "  --dry-run      Print generated files to stdout instead of writing them"
            echo "  --output=DIR   Write files to DIR instead of current directory"
            exit 0 ;;
    esac
done

mkdir -p "$OUTPUT_DIR"

# ── Banner ────────────────────────────────────────────────────────────────────
echo -e "${BOLD}${CYAN}"
echo "╔══════════════════════════════════════════════════════════╗"
echo "║      NetBird External Relay Server Setup Wizard          ║"
echo "╚══════════════════════════════════════════════════════════╝${RESET}"
echo ""
echo "This script generates  relay.env  and  docker-compose.yml"
echo "for one or more NetBird relay servers."
$DRY_RUN && warn "DRY-RUN mode: files will be printed, not written."
echo ""

# ═════════════════════════════════════════════════════════════════════════════
# STEP 1 – Auth secret
# ═════════════════════════════════════════════════════════════════════════════
header "Step 1 · Authentication Secret"
echo "All relay servers AND your main NetBird server must share the same secret."
echo ""

if ask_yn "Generate a new random secret automatically?" "y"; then
    AUTH_SECRET="$(generate_secret)"
    success "Generated secret: ${BOLD}${AUTH_SECRET}${RESET}"
    warn "Save this — you'll need it for every relay and your main server config."
else
    ask_secret AUTH_SECRET "Paste your existing shared secret"
    success "Using supplied secret."
fi

# ═════════════════════════════════════════════════════════════════════════════
# STEP 2 – Number of relay servers
# ═════════════════════════════════════════════════════════════════════════════
header "Step 2 · How many relay servers?"
ask RELAY_COUNT "Number of relay servers to configure" "1"
if ! [[ "$RELAY_COUNT" =~ ^[1-9][0-9]*$ ]]; then
    error "Please enter a positive integer."; exit 1
fi

# ═════════════════════════════════════════════════════════════════════════════
# STEP 3 – Per-relay configuration
# ═════════════════════════════════════════════════════════════════════════════
declare -a RELAY_DOMAINS=()
declare -a RELAY_DIRS=()
declare -a RELAY_LISTEN_PORTS=()
declare -a RELAY_STUN_PORTS_LIST=()
declare -a RELAY_TLS_MODES=()
declare -a RELAY_ENABLE_STUN=()

for (( i=1; i<=RELAY_COUNT; i++ )); do

    header "Step 3.$i · Relay Server #${i}"

    # ── Domain ────────────────────────────────────────────────────────────────
    while true; do
        ask DOMAIN "  Domain name (e.g. relay-us.example.com)" ""
        if validate_domain "$DOMAIN"; then break
        else warn "  That doesn't look like a valid domain. Try again."; fi
    done
    RELAY_DOMAINS+=("$DOMAIN")

    # ── Listen & exposed ports ────────────────────────────────────────────────
    ask LISTEN_PORT  "  HTTPS listen port" "443"
    ask LOG_LEVEL    "  Log level (debug/info/warn/error)" "info"

    # ── STUN ──────────────────────────────────────────────────────────────────
    echo ""
    if ask_yn "  Enable embedded STUN server?" "y"; then
        ENABLE_STUN=true
        ask STUN_PORTS "  STUN port(s) — comma-separated for multiple (e.g. 3478,3479)" "3478"
    else
        ENABLE_STUN=false
        STUN_PORTS=""
    fi

    # ── TLS mode ──────────────────────────────────────────────────────────────
    echo ""
    echo "  TLS mode:"
    echo "    1) Let's Encrypt (automatic — server needs port 80/tcp open)"
    echo "    2) Existing certificates (wildcard / own CA)"
    ask TLS_MODE "  Choose [1/2]" "1"

    if [[ "$TLS_MODE" == "1" ]]; then
        ask LE_EMAIL    "  Let's Encrypt email" ""
        LE_DATA_DIR="/data/letsencrypt"
        CERT_HOST_PATH=""; CERT_FILE=""; KEY_FILE=""
    else
        ask CERT_HOST_PATH "  Host path to certs directory" "/opt/1panel/www/sites/${DOMAIN}/ssl"
        ask CERT_FILE      "  Cert file path inside container" "/certs/fullchain.pem"
        ask KEY_FILE       "  Key  file path inside container" "/certs/privkey.pem"
        LE_EMAIL=""; LE_DATA_DIR=""
    fi

    # ── Output directory ──────────────────────────────────────────────────────
    ask RELAY_DIR "  Output directory for this relay's files" "/opt/netbird-relay"
    RELAY_DIRS+=("$RELAY_DIR")
    RELAY_LISTEN_PORTS+=("$LISTEN_PORT")
    RELAY_STUN_PORTS_LIST+=("$STUN_PORTS")
    RELAY_TLS_MODES+=("$TLS_MODE")
    RELAY_ENABLE_STUN+=("$ENABLE_STUN")

    # ── Build relay.env ───────────────────────────────────────────────────────
    ENV_FILE="${RELAY_DIR}/relay.env"
    COMPOSE_FILE="${RELAY_DIR}/docker-compose.yml"

    ENV_CONTENT="# NetBird Relay — ${DOMAIN}
# Generated by setup-relay.sh on $(date -u '+%Y-%m-%dT%H:%M:%SZ')
# ---------------------------------------------------------------

NB_LOG_LEVEL=${LOG_LEVEL}
NB_LISTEN_ADDRESS=:${LISTEN_PORT}
NB_EXPOSED_ADDRESS=rels://${DOMAIN}:${LISTEN_PORT}
NB_AUTH_SECRET=${AUTH_SECRET}
"

    if [[ "$TLS_MODE" == "1" ]]; then
        ENV_CONTENT+="
# TLS — Let's Encrypt (automatic certificate provisioning)
NB_LETSENCRYPT_DOMAINS=${DOMAIN}
NB_LETSENCRYPT_EMAIL=${LE_EMAIL}
NB_LETSENCRYPT_DATA_DIR=${LE_DATA_DIR}
"
    else
        ENV_CONTENT+="
# TLS — Existing certificates
NB_TLS_CERT_FILE=${CERT_FILE}
NB_TLS_KEY_FILE=${KEY_FILE}
"
    fi

    if $ENABLE_STUN; then
        ENV_CONTENT+="
# Embedded STUN
NB_ENABLE_STUN=true
NB_STUN_PORTS=${STUN_PORTS}
"
    else
        ENV_CONTENT+="
# Embedded STUN disabled
NB_ENABLE_STUN=false
"
    fi

    # ── Build STUN port mappings ───────────────────────────────────────────────
    STUN_PORT_LINES=""
    if $ENABLE_STUN && [[ -n "$STUN_PORTS" ]]; then
        IFS=',' read -ra SPORT_ARRAY <<< "$STUN_PORTS"
        for SP in "${SPORT_ARRAY[@]}"; do
            SP="${SP// /}"
            STUN_PORT_LINES+="      - '${SP}:${SP}/udp'"$'\n'
        done
    fi

    # ── Build volume section ───────────────────────────────────────────────────
    if [[ "$TLS_MODE" == "1" ]]; then
        VOLUME_LINES="      - relay_data:/data"
    else
        VOLUME_LINES="      - ${CERT_HOST_PATH}:$(dirname "${CERT_FILE}"):ro
      - relay_data:/data"
    fi

    # ── Build docker-compose.yml ──────────────────────────────────────────────
    if [[ "$TLS_MODE" == "1" ]]; then
        LE_PORT_LINE="      - '80:80'"$'\n'
    else
        LE_PORT_LINE=""
    fi

    COMPOSE_CONTENT="# NetBird Relay — ${DOMAIN}
# Generated by setup-relay.sh on $(date -u '+%Y-%m-%dT%H:%M:%SZ')
# ---------------------------------------------------------------
services:
  relay:
    image: netbirdio/relay:latest
    container_name: netbird-relay
    restart: unless-stopped
    ports:
      - '${LISTEN_PORT}:${LISTEN_PORT}'
${LE_PORT_LINE}${STUN_PORT_LINES}    env_file:
      - relay.env
    volumes:
${VOLUME_LINES}
    logging:
      driver: \"json-file\"
      options:
        max-size: \"500m\"
        max-file: \"2\"

volumes:
  relay_data:
"

    # ── Write or print ────────────────────────────────────────────────────────
    if $DRY_RUN; then
        echo ""
        echo -e "${BOLD}▶ ${ENV_FILE}${RESET}"
        echo "────────────────────────────────────────"
        echo "$ENV_CONTENT"
        echo ""
        echo -e "${BOLD}▶ ${COMPOSE_FILE}${RESET}"
        echo "────────────────────────────────────────"
        echo "$COMPOSE_CONTENT"
    else
        mkdir -p "$RELAY_DIR"
        printf '%s' "$ENV_CONTENT"   > "$ENV_FILE"
        printf '%s' "$COMPOSE_CONTENT" > "$COMPOSE_FILE"
        chmod 600 "$ENV_FILE"   # protect the secret
        success "  Written: ${ENV_FILE}"
        success "  Written: ${COMPOSE_FILE}"
    fi

done   # end per-relay loop

# ═════════════════════════════════════════════════════════════════════════════
# STEP 4 – Main server snippet
# ═════════════════════════════════════════════════════════════════════════════
header "Step 4 · Main Server Configuration Snippet"

echo "Add the following relay entries to your main NetBird server's"
echo "management.json  (or equivalent config), replacing any existing ones:"
echo ""
echo -e "${BOLD}Relay URLs:${RESET}"
for i in "${!RELAY_DOMAINS[@]}"; do
    echo "  rels://${RELAY_DOMAINS[$i]}:${RELAY_LISTEN_PORTS[$i]}"
done
echo ""
echo -e "${BOLD}STUN URLs:${RESET}"
for i in "${!RELAY_DOMAINS[@]}"; do
    if [[ "${RELAY_ENABLE_STUN[$i]}" == "true" ]] && [[ -n "${RELAY_STUN_PORTS_LIST[$i]}" ]]; then
        IFS=',' read -ra _SP <<< "${RELAY_STUN_PORTS_LIST[$i]}"
        for _P in "${_SP[@]}"; do
            _P="${_P// /}"
            echo "  stun:${RELAY_DOMAINS[$i]}:${_P}"
        done
    fi
done
echo ""
echo -e "${BOLD}Shared secret:${RESET} ${AUTH_SECRET}"
echo ""

# ═════════════════════════════════════════════════════════════════════════════
# STEP 5 – Next steps
# ═════════════════════════════════════════════════════════════════════════════
header "Step 5 · Next Steps"

for i in "${!RELAY_DOMAINS[@]}"; do
    D="${RELAY_DOMAINS[$i]}"
    DIR="${RELAY_DIRS[$i]}"
    LP="${RELAY_LISTEN_PORTS[$i]}"
    TM="${RELAY_TLS_MODES[$i]}"
    ES="${RELAY_ENABLE_STUN[$i]}"
    SP="${RELAY_STUN_PORTS_LIST[$i]}"

    echo -e "${BOLD}Relay: ${D}${RESET}"
    echo "  1. Copy ${DIR}/ to the relay server"
    echo "  2. On the relay server:"
    echo "       cd ${DIR}"
    echo "       docker compose up -d"
    echo "       docker compose logs -f"
    if [[ "$TM" == "1" ]]; then
        echo "  3. Trigger TLS cert (Let's Encrypt):"
        echo "       curl -v https://${D}/"
        echo "     Expect: 404 page not found  +  valid LE cert"
    fi
    echo ""
    echo -e "  ${BOLD}Firewall ports to open:${RESET}"
    [[ "$TM" == "1" ]] && echo "    80/tcp   — Let's Encrypt HTTP challenge"
    echo "    ${LP}/tcp  — Relay (HTTPS)"
    if [[ "$ES" == "true" ]] && [[ -n "$SP" ]]; then
        IFS=',' read -ra _SP <<< "$SP"
        for _P in "${_SP[@]}"; do
            _P="${_P// /}"
            echo "    ${_P}/udp  — STUN"
        done
    fi
    echo ""
done
success "Setup complete! 🎉"