#!/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
}

parse_existing_env() {
    # parse_existing_env <relay.env path>
    # Sets EX_* variables from an existing relay.env file.
    local envfile="$1"
    [[ -f "$envfile" ]] || return 1

    _get() { grep -m1 "^${1}=" "$envfile" 2>/dev/null | cut -d= -f2- || true; }

    EX_LOG_LEVEL="$(_get NB_LOG_LEVEL)"
    EX_AUTH_SECRET="$(_get NB_AUTH_SECRET)"
    EX_LE_EMAIL="$(_get NB_LETSENCRYPT_EMAIL)"
    EX_CERT_FILE="$(_get NB_TLS_CERT_FILE)"
    EX_KEY_FILE="$(_get NB_TLS_KEY_FILE)"
    EX_ENABLE_STUN="$(_get NB_ENABLE_STUN)"
    EX_STUN_PORTS="$(_get NB_STUN_PORTS)"

    # Parse domain and port from NB_EXPOSED_ADDRESS=rels://domain:port
    local exposed; exposed="$(_get NB_EXPOSED_ADDRESS)"
    EX_DOMAIN="${exposed#rels://}"; EX_DOMAIN="${EX_DOMAIN%%:*}"
    EX_LISTEN_PORT="${exposed##*:}"

    # Detect TLS mode
    local le_domains; le_domains="$(_get NB_LETSENCRYPT_DOMAINS)"
    if [[ -n "$le_domains" ]]; then
        EX_TLS_MODE="1"
    elif [[ -n "$EX_CERT_FILE" ]]; then
        EX_TLS_MODE="2"   # could be self-signed, but treat as existing cert
    else
        EX_TLS_MODE="1"   # fallback
    fi
}

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 ""

# Try to pre-read an existing secret from /opt/netbird-relay/relay.env
_PREREAD_SECRET=""
if [[ -f "/opt/netbird-relay/relay.env" ]]; then
    _PREREAD_SECRET="$(grep -m1 '^NB_AUTH_SECRET=' /opt/netbird-relay/relay.env 2>/dev/null | cut -d= -f2- || true)"
fi

if [[ -n "$_PREREAD_SECRET" ]]; then
    info "Found existing secret in /opt/netbird-relay/relay.env"
    echo -e "  Current secret: ${BOLD}${_PREREAD_SECRET}${RESET}"
    echo "  1) Keep existing secret"
    echo "  2) Generate a new random secret"
    echo "  3) Enter a different secret manually"
    ask _SECRET_CHOICE "  Choose [1/2/3]" "1"
    case "$_SECRET_CHOICE" in
        2)
            AUTH_SECRET="$(generate_secret)"
            success "Generated new secret: ${BOLD}${AUTH_SECRET}${RESET}"
            warn "Update NB_AUTH_SECRET on ALL relay servers and your main server config."
            ;;
        3)
            ask AUTH_SECRET "Paste your shared secret" ""
            success "Using supplied secret."
            ;;
        *)
            AUTH_SECRET="$_PREREAD_SECRET"
            success "Keeping existing secret."
            ;;
    esac
elif 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 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}"

    # ── Output directory (ask first so we can read existing config) ───────────
    ask RELAY_DIR "  Output directory for this relay's files" "/opt/netbird-relay"

    # ── Load existing config as defaults if relay.env is present ─────────────
    EX_DOMAIN=""; EX_LISTEN_PORT="443"; EX_LOG_LEVEL="info"
    EX_AUTH_SECRET=""; EX_LE_EMAIL=""; EX_CERT_FILE=""; EX_KEY_FILE=""
    EX_ENABLE_STUN="true"; EX_STUN_PORTS="3478"; EX_TLS_MODE="1"
    _EXISTING_ENV="${RELAY_DIR}/relay.env"
    if parse_existing_env "$_EXISTING_ENV"; then
        info "  Found existing config in ${_EXISTING_ENV} — using as defaults."
    fi

    # ── Domain ────────────────────────────────────────────────────────────────
    while true; do
        ask DOMAIN "  Domain name (e.g. relay-us.example.com)" "${EX_DOMAIN}"
        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" "${EX_LISTEN_PORT:-443}"
    ask LOG_LEVEL    "  Log level (debug/info/warn/error)" "${EX_LOG_LEVEL:-info}"

    # ── STUN ──────────────────────────────────────────────────────────────────
    echo ""
    _STUN_DEFAULT="y"; [[ "${EX_ENABLE_STUN}" == "false" ]] && _STUN_DEFAULT="n"
    if ask_yn "  Enable embedded STUN server?" "$_STUN_DEFAULT"; then
        ENABLE_STUN=true
        ask STUN_PORTS "  STUN port(s) — comma-separated for multiple (e.g. 3478,3479)" "${EX_STUN_PORTS:-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)"
    echo "    3) Self-signed certificate (generated by this script)"
    ask TLS_MODE "  Choose [1/2/3]" "${EX_TLS_MODE:-1}"

    if [[ "$TLS_MODE" == "1" ]]; then
        ask LE_EMAIL    "  Let's Encrypt email" "${EX_LE_EMAIL}"
        LE_DATA_DIR="/data/letsencrypt"
        CERT_HOST_PATH=""; CERT_FILE=""; KEY_FILE=""
    elif [[ "$TLS_MODE" == "2" ]]; then
        # Derive existing host path from cert file path stored in env (best-effort)
        _EX_CERT_HOST=""
        [[ -n "$EX_CERT_FILE" ]] && _EX_CERT_HOST="$(
            grep -m1 'volumes:' -A5 "${RELAY_DIR}/docker-compose.yml" 2>/dev/null \
            | grep -m1 ':/certs:ro\|:/certs ' \
            | awk -F: '{print $1}' | sed 's/^ *- *//' || true
        )"
        ask CERT_HOST_PATH "  Host path to certs directory" "${_EX_CERT_HOST:-/opt/1panel/www/sites/${DOMAIN}/ssl}"
        ask CERT_FILE      "  Cert file path inside container" "${EX_CERT_FILE:-/certs/fullchain.pem}"
        ask KEY_FILE       "  Key  file path inside container" "${EX_KEY_FILE:-/certs/privkey.pem}"
        LE_EMAIL=""; LE_DATA_DIR=""
    else
        echo ""
        echo "  ── Self-signed certificate details ──"

        ask SS_CN   "  Common Name (CN)" "${DOMAIN}"
        ask SS_O    "  Organization (O)" ""
        ask SS_OU   "  Organizational Unit (OU)" ""
        ask SS_C    "  Country (C, 2-letter ISO)" ""
        ask SS_ST   "  State / Province (ST)" ""
        ask SS_L    "  Locality / City (L)" ""

        SS_SANS="DNS:${DOMAIN}"
        echo ""
        echo "  Subject Alternative Names (SANs):"
        echo "  DNS:${DOMAIN} is included automatically."
        echo "  Add extra SANs one by one (IP:x.x.x.x or DNS:other.example.com)."
        echo "  Press Enter on an empty line when done."
        while true; do
            read -rp "$(echo -e "${BOLD}  Extra SAN (or Enter to finish): ${RESET}")" _san
            [[ -z "$_san" ]] && break
            SS_SANS+=",${_san}"
        done

        ask SS_DAYS "  Certificate validity (days)" "3650"

        echo ""
        echo "  Key type:"
        echo "    1) RSA"
        echo "    2) ECC (ECDSA)"
        ask SS_KEY_TYPE "  Choose [1/2]" "1"

        if [[ "$SS_KEY_TYPE" == "1" ]]; then
            echo "  RSA key size:"
            echo "    1) 2048 bit"
            echo "    2) 4096 bit"
            ask SS_KEY_SIZE "  Choose [1/2]" "2"
            [[ "$SS_KEY_SIZE" == "1" ]] && SS_NEWKEY="rsa:2048" || SS_NEWKEY="rsa:4096"
        else
            echo "  ECC curve:"
            echo "    1) P-256 (prime256v1)"
            echo "    2) P-384 (secp384r1)"
            echo "    3) P-521 (secp521r1)"
            ask SS_CURVE "  Choose [1/2/3]" "1"
            case "$SS_CURVE" in
                1) SS_NEWKEY="ec" ; SS_PKEYOPT="ec_paramgen_curve:P-256" ;;
                2) SS_NEWKEY="ec" ; SS_PKEYOPT="ec_paramgen_curve:P-384" ;;
                3) SS_NEWKEY="ec" ; SS_PKEYOPT="ec_paramgen_curve:P-521" ;;
            esac
        fi

        SS_SUBJ=""
        [[ -n "$SS_C"  ]] && SS_SUBJ+="/C=${SS_C}"
        [[ -n "$SS_ST" ]] && SS_SUBJ+="/ST=${SS_ST}"
        [[ -n "$SS_L"  ]] && SS_SUBJ+="/L=${SS_L}"
        [[ -n "$SS_O"  ]] && SS_SUBJ+="/O=${SS_O}"
        [[ -n "$SS_OU" ]] && SS_SUBJ+="/OU=${SS_OU}"
        SS_SUBJ+="/CN=${SS_CN}"

        CERT_HOST_PATH=""
        CERT_FILE="/certs/fullchain.pem"
        KEY_FILE="/certs/privkey.pem"
        LE_EMAIL=""; LE_DATA_DIR=""
    fi
    # For self-signed mode, certs live inside the relay dir
    [[ "$TLS_MODE" == "3" ]] && CERT_HOST_PATH="${RELAY_DIR}/certs"
    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 — $([ "$TLS_MODE" == "3" ] && echo "Self-signed certificate" || echo "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"
        if [[ "$TLS_MODE" == "3" ]]; then
            echo ""
            echo -e "${BOLD}▶ Self-signed cert (would run):${RESET}"
            echo "────────────────────────────────────────"
            echo "mkdir -p ${CERT_HOST_PATH}"
            if [[ "$SS_KEY_TYPE" == "1" ]]; then
                echo "openssl req -x509 -newkey ${SS_NEWKEY} -sha256 -days ${SS_DAYS} -nodes \\"
            else
                echo "openssl req -x509 -newkey ${SS_NEWKEY} -pkeyopt ${SS_PKEYOPT} -sha256 -days ${SS_DAYS} -nodes \\"
            fi
            echo "  -keyout ${CERT_HOST_PATH}/privkey.pem \\"
            echo "  -out    ${CERT_HOST_PATH}/fullchain.pem \\"
            echo "  -subj   \"${SS_SUBJ}\" \\"
            echo "  -addext \"subjectAltName=${SS_SANS}\""
        fi
    else
        mkdir -p "$RELAY_DIR"
        printf '%s' "$ENV_CONTENT"     > "$ENV_FILE"
        printf '%s' "$COMPOSE_CONTENT" > "$COMPOSE_FILE"
        chmod 600 "$ENV_FILE"
        success "  Written: ${ENV_FILE}"
        success "  Written: ${COMPOSE_FILE}"

        if [[ "$TLS_MODE" == "3" ]]; then
            if ! command -v openssl &>/dev/null; then
                error "  openssl not found — cannot generate self-signed certificate."
                error "  Install openssl and re-run, or run the following manually:"
                if [[ "$SS_KEY_TYPE" == "1" ]]; then
                    error "    openssl req -x509 -newkey ${SS_NEWKEY} -sha256 -days ${SS_DAYS} -nodes \\"
                else
                    error "    openssl req -x509 -newkey ${SS_NEWKEY} -pkeyopt ${SS_PKEYOPT} -sha256 -days ${SS_DAYS} -nodes \\"
                fi
                error "      -keyout ${CERT_HOST_PATH}/privkey.pem \\"
                error "      -out    ${CERT_HOST_PATH}/fullchain.pem \\"
                error "      -subj   \"${SS_SUBJ}\" \\"
                error "      -addext \"subjectAltName=${SS_SANS}\""
            else
                info "  Generating self-signed certificate ..."
                mkdir -p "${CERT_HOST_PATH}"
                if [[ "$SS_KEY_TYPE" == "1" ]]; then
                    openssl req -x509 \
                        -newkey "${SS_NEWKEY}" \
                        -sha256 -days "${SS_DAYS}" -nodes \
                        -keyout "${CERT_HOST_PATH}/privkey.pem" \
                        -out    "${CERT_HOST_PATH}/fullchain.pem" \
                        -subj   "${SS_SUBJ}" \
                        -addext "subjectAltName=${SS_SANS}" \
                        2>/dev/null
                else
                    openssl req -x509 \
                        -newkey "${SS_NEWKEY}" -pkeyopt "${SS_PKEYOPT}" \
                        -sha256 -days "${SS_DAYS}" -nodes \
                        -keyout "${CERT_HOST_PATH}/privkey.pem" \
                        -out    "${CERT_HOST_PATH}/fullchain.pem" \
                        -subj   "${SS_SUBJ}" \
                        -addext "subjectAltName=${SS_SANS}" \
                        2>/dev/null
                fi
                chmod 600 "${CERT_HOST_PATH}/privkey.pem"
                success "  Certificate : ${CERT_HOST_PATH}/fullchain.pem"
                success "  Private key : ${CERT_HOST_PATH}/privkey.pem"
                info   "  Subject     : ${SS_SUBJ}"
                info   "  SANs        : ${SS_SANS}"
                info   "  Valid for   : ${SS_DAYS} days"
            fi
        fi
    fi

done   # end per-relay loop

# ═════════════════════════════════════════════════════════════════════════════
# STEP 4 – Main server config.yaml snippet
# ═════════════════════════════════════════════════════════════════════════════
header "Step 4 · Main Server config.yaml Snippet"

echo "Reference: Set Up External Relay Servers — NetBird Docs"
echo "  https://docs.netbird.io/selfhosted/splitting-self-hosted-deployment/set-up-external-relay-servers"
echo ""
echo "On your main server:"
echo "  cd ~/netbird   # or wherever your deployment lives"
echo "  nano config.yaml"
echo ""
echo "Remove 'authSecret' from the 'server' section (disables embedded relay),"
echo "and add/replace the 'relays' and 'stuns' sections as shown below."
echo "(The presence of 'relays' also disables the embedded STUN server,"
echo " so 'stuns' is required.)"
echo ""

# ── Build stuns block ─────────────────────────────────────────────────────────
_STUNS_YAML=""
_STUNS_YAML_PLACEHOLDER=""
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// /}"
            _STUNS_YAML+="    - uri: \"stun:${RELAY_DOMAINS[$i]}:${_P}\""$'\n'
            _STUNS_YAML+="      proto: \"udp\""$'\n'
            _STUNS_YAML_PLACEHOLDER+="    - uri: \"stun:<relay-$((i+1))-domain>:${_P}\""$'\n'
            _STUNS_YAML_PLACEHOLDER+="      proto: \"udp\""$'\n'
        done
    fi
done

# ── Build relays.addresses block ──────────────────────────────────────────────
_RELAY_ADDRS_YAML=""
_RELAY_ADDRS_YAML_PLACEHOLDER=""
for i in "${!RELAY_DOMAINS[@]}"; do
    _RELAY_ADDRS_YAML+="      - \"rels://${RELAY_DOMAINS[$i]}:${RELAY_LISTEN_PORTS[$i]}\""$'\n'
    _RELAY_ADDRS_YAML_PLACEHOLDER+="      - \"rels://<relay-$((i+1))-domain>:${RELAY_LISTEN_PORTS[$i]}\""$'\n'
done

# ── Print the snippet ─────────────────────────────────────────────────────────
echo -e "${BOLD}┌─ config.yaml (relevant section) ───────────────────────────────────┐${RESET}"
cat <<YAML
server:
  # listenAddress / exposedAddress / metricsPort / etc. — keep as-is

  # Remove or comment out the embedded relay secret:
  # authSecret: ...
  # Remove or comment out embedded STUN ports:
  # stunPorts:
  #   - 3478

  # External STUN servers (your relay servers)
  stuns:
${_STUNS_YAML_PLACEHOLDER}
  # External relay servers
  relays:
    addresses:
${_RELAY_ADDRS_YAML_PLACEHOLDER}    secret: "<your-shared-secret>"
    credentialsTTL: "24h"

  # auth: ... (keep your existing auth config below)
YAML
echo -e "${BOLD}└────────────────────────────────────────────────────────────────────┘${RESET}"
echo ""
warn "The 'secret' under relays MUST match NB_AUTH_SECRET on all relay servers."
warn "Mismatched secrets cause relay connections to fail silently."
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"
    elif [[ "$TM" == "3" ]]; then
        echo "  Note: Self-signed cert is at ${DIR}/certs/"
        echo "        Clients must trust this CA or skip TLS verification."
    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! 🎉"

# ═════════════════════════════════════════════════════════════════════════════
# STEP 6 – Firewall (optional)
# ═════════════════════════════════════════════════════════════════════════════
header "Step 6 · Firewall Configuration (optional)"

$DRY_RUN && { warn "DRY-RUN: skipping firewall step."; exit 0; }

if ! ask_yn "Configure firewall rules now?" "y"; then
    info "Skipped. Remember to open the ports listed above manually."
    exit 0
fi

# ── Detect firewall ───────────────────────────────────────────────────────────
detect_firewall() {
    if systemctl is-active --quiet firewalld 2>/dev/null; then
        echo "firewalld"
    elif systemctl is-active --quiet ufw 2>/dev/null || command -v ufw &>/dev/null && ufw status 2>/dev/null | grep -q "Status: active"; then
        echo "ufw"
    elif command -v nft &>/dev/null && nft list ruleset 2>/dev/null | grep -q "table"; then
        echo "nftables"
    elif command -v iptables &>/dev/null; then
        echo "iptables"
    else
        echo "none"
    fi
}

FW="$(detect_firewall)"

case "$FW" in
    firewalld) info "Detected: firewalld" ;;
    ufw)       info "Detected: ufw" ;;
    nftables)  info "Detected: nftables" ;;
    iptables)  info "Detected: iptables" ;;
    none)      warn "No supported firewall detected (firewalld / ufw / nftables / iptables)."
               info "Open the ports listed in Step 5 manually."
               exit 0 ;;
esac

# ── Collect all ports to open across all relays ───────────────────────────────
declare -a FW_TCP=()
declare -a FW_UDP=()

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

    # TCP: relay port
    FW_TCP+=("$LP")
    # TCP: port 80 for Let's Encrypt
    [[ "$TM" == "1" ]] && FW_TCP+=("80")
    # UDP: STUN ports
    if [[ "$ES" == "true" ]] && [[ -n "$SP" ]]; then
        IFS=',' read -ra _SP <<< "$SP"
        for _P in "${_SP[@]}"; do
            FW_UDP+=("${_P// /}")
        done
    fi
done

# Deduplicate
mapfile -t FW_TCP < <(printf '%s\n' "${FW_TCP[@]}" | sort -un)
mapfile -t FW_UDP < <(printf '%s\n' "${FW_UDP[@]}" | sort -un)

echo ""
echo -e "${BOLD}Ports to open:${RESET}"
for p in "${FW_TCP[@]}"; do echo "  ${p}/tcp"; done
for p in "${FW_UDP[@]}"; do echo "  ${p}/udp"; done
echo ""

if ! ask_yn "Apply these rules?" "y"; then
    info "Skipped."
    exit 0
fi

# ── Apply rules ───────────────────────────────────────────────────────────────
apply_firewall_rules() {
    local fw="$1"

    case "$fw" in

    firewalld)
        for p in "${FW_TCP[@]}"; do
            firewall-cmd --permanent --add-port="${p}/tcp" && \
                success "  firewalld: opened ${p}/tcp" || \
                error   "  firewalld: failed to open ${p}/tcp"
        done
        for p in "${FW_UDP[@]}"; do
            firewall-cmd --permanent --add-port="${p}/udp" && \
                success "  firewalld: opened ${p}/udp" || \
                error   "  firewalld: failed to open ${p}/udp"
        done
        firewall-cmd --reload && success "  firewalld: reloaded" || error "  firewalld: reload failed"
        ;;

    ufw)
        for p in "${FW_TCP[@]}"; do
            ufw allow "${p}/tcp" && \
                success "  ufw: allowed ${p}/tcp" || \
                error   "  ufw: failed to allow ${p}/tcp"
        done
        for p in "${FW_UDP[@]}"; do
            ufw allow "${p}/udp" && \
                success "  ufw: allowed ${p}/udp" || \
                error   "  ufw: failed to allow ${p}/udp"
        done
        ;;

    nftables)
        # Add rules to the first inet/ip filter input chain found, or create one
        _NFT_TABLE="filter"
        _NFT_CHAIN="input"
        # Check if table/chain exist
        if ! nft list chain inet "${_NFT_TABLE}" "${_NFT_CHAIN}" &>/dev/null; then
            nft add table inet "${_NFT_TABLE}"
            nft add chain inet "${_NFT_TABLE}" "${_NFT_CHAIN}" \
                '{ type filter hook input priority 0 ; policy accept ; }'
            info "  nftables: created table inet ${_NFT_TABLE} chain ${_NFT_CHAIN}"
        fi
        for p in "${FW_TCP[@]}"; do
            nft add rule inet "${_NFT_TABLE}" "${_NFT_CHAIN}" tcp dport "${p}" accept && \
                success "  nftables: accepted ${p}/tcp" || \
                error   "  nftables: failed ${p}/tcp"
        done
        for p in "${FW_UDP[@]}"; do
            nft add rule inet "${_NFT_TABLE}" "${_NFT_CHAIN}" udp dport "${p}" accept && \
                success "  nftables: accepted ${p}/udp" || \
                error   "  nftables: failed ${p}/udp"
        done
        # Persist
        if command -v nft &>/dev/null; then
            if [[ -d /etc/nftables.d ]]; then
                nft list ruleset > /etc/nftables.d/netbird-relay.nft && \
                    success "  nftables: saved to /etc/nftables.d/netbird-relay.nft"
            elif [[ -f /etc/nftables.conf ]]; then
                nft list ruleset > /etc/nftables.conf && \
                    success "  nftables: saved to /etc/nftables.conf"
            else
                warn "  nftables: rules applied but not persisted — save manually with:"
                warn "    nft list ruleset > /etc/nftables.conf"
            fi
        fi
        ;;

    iptables)
        for p in "${FW_TCP[@]}"; do
            iptables -C INPUT -p tcp --dport "${p}" -j ACCEPT 2>/dev/null || {
                iptables -A INPUT -p tcp --dport "${p}" -j ACCEPT && \
                    success "  iptables: opened ${p}/tcp" || \
                    error   "  iptables: failed to open ${p}/tcp"
            }
        done
        for p in "${FW_UDP[@]}"; do
            iptables -C INPUT -p udp --dport "${p}" -j ACCEPT 2>/dev/null || {
                iptables -A INPUT -p udp --dport "${p}" -j ACCEPT && \
                    success "  iptables: opened ${p}/udp" || \
                    error   "  iptables: failed to open ${p}/udp"
            }
        done
        # Persist
        if command -v netfilter-persistent &>/dev/null; then
            netfilter-persistent save && success "  iptables: rules persisted via netfilter-persistent"
        elif command -v iptables-save &>/dev/null; then
            if [[ -f /etc/iptables/rules.v4 ]]; then
                iptables-save > /etc/iptables/rules.v4 && \
                    success "  iptables: saved to /etc/iptables/rules.v4"
            else
                warn "  iptables: rules applied but not persisted — save manually with:"
                warn "    iptables-save > /etc/iptables/rules.v4"
            fi
        fi
        ;;
    esac
}

if [[ "$EUID" -ne 0 ]]; then
    warn "Not running as root — firewall commands may fail."
    warn "Re-run with sudo if needed."
fi

apply_firewall_rules "$FW"
echo ""
success "Firewall rules applied. 🎉"