#!/usr/bin/env bash
set -eu

echo "NexaVPN gateway helper starting"
mkdir -p /var/lib/nexavpn

IFACE="${NEXAVPN_GATEWAY_INTERFACE:-wg0}"
UPLINK_IFACE="${NEXAVPN_UPLINK_INTERFACE:-eth0}"
ENABLE_MASQUERADE="${NEXAVPN_ENABLE_MASQUERADE:-true}"
GATEWAY_NAME="${NEXAVPN_GATEWAY_NAME:-primary-gateway}"
BOOTSTRAP_URL="${NEXAVPN_GATEWAY_BOOTSTRAP_URL:-http://backend:8080/api/v1/gateway-agent/bootstrap}"
SYNC_BASE_URL="${NEXAVPN_GATEWAY_SYNC_URL:-http://backend:8080/api/v1/gateway-agent}"
GATEWAY_ID_FILE="/var/lib/nexavpn/gateway-id"
BACKEND_HOST="${NEXAVPN_BACKEND_HOST:-backend}"

if [ -z "${GATEWAY_BOOTSTRAP_TOKEN:-}" ]; then
  echo "GATEWAY_BOOTSTRAP_TOKEN is required."
  tail -f /dev/null
  exit 0
fi

if [ -z "${NEXAVPN_GATEWAY_PRIVATE_KEY:-}" ]; then
  if [ -f /var/lib/nexavpn/gateway-private.key ]; then
    NEXAVPN_GATEWAY_PRIVATE_KEY="$(cat /var/lib/nexavpn/gateway-private.key)"
  else
    wg genkey | tee /var/lib/nexavpn/gateway-private.key >/tmp/nexavpn-gateway-private.key
    NEXAVPN_GATEWAY_PRIVATE_KEY="$(cat /tmp/nexavpn-gateway-private.key)"
    rm -f /tmp/nexavpn-gateway-private.key
  fi
fi

if [ -z "${NEXAVPN_GATEWAY_ID:-}" ] && [ -f "${GATEWAY_ID_FILE}" ]; then
  NEXAVPN_GATEWAY_ID="$(cat "${GATEWAY_ID_FILE}")"
fi

bootstrap_gateway() {
  GATEWAY_PUBLIC_KEY="$(printf '%s' "${NEXAVPN_GATEWAY_PRIVATE_KEY}" | wg pubkey)"
  echo "Bootstrapping gateway ${GATEWAY_NAME}"
  BOOTSTRAP_RESPONSE="$(curl -fsSL \
    -H "Content-Type: application/json" \
    -H "X-Gateway-Bootstrap-Token: ${GATEWAY_BOOTSTRAP_TOKEN}" \
    -d "{\"name\":\"${GATEWAY_NAME}\",\"endpoint\":\"${DEFAULT_GATEWAY_ENDPOINT:-localhost:51820}\",\"public_key\":\"${GATEWAY_PUBLIC_KEY}\",\"listen_port\":51820,\"vpn_cidr\":\"${DEFAULT_VPN_CIDR:-100.96.0.0/24}\",\"dns_servers\":[\"10.20.0.53\"]}" \
    "${BOOTSTRAP_URL}")"
  NEXAVPN_GATEWAY_ID="$(printf '%s' "${BOOTSTRAP_RESPONSE}" | jq -r '.id')"
  if [ -z "${NEXAVPN_GATEWAY_ID:-}" ] || [ "${NEXAVPN_GATEWAY_ID}" = "null" ]; then
    echo "Gateway bootstrap did not return an id."
    return 1
  fi
  printf '%s' "${NEXAVPN_GATEWAY_ID}" > "${GATEWAY_ID_FILE}"
}

STATE_JSON="/var/lib/nexavpn/sync-bundle.json"
WG_CONF="/etc/wireguard/${IFACE}.conf"
WG_GENERATED="/var/lib/nexavpn/${IFACE}.generated.conf"
NFT_CONF="/var/lib/nexavpn/nftables.generated.conf"

mkdir -p /etc/wireguard

apply_bundle() {
  if [ -z "${NEXAVPN_GATEWAY_ID:-}" ]; then
    bootstrap_gateway || return 1
  fi

  if [ -z "${NEXAVPN_GATEWAY_ID:-}" ] || [ -z "${NEXAVPN_GATEWAY_PRIVATE_KEY:-}" ]; then
    echo "Gateway sync is not configured yet."
    return 1
  fi

  SYNC_URL="${SYNC_BASE_URL}/${NEXAVPN_GATEWAY_ID}/sync"
  echo "Fetching bundle from ${SYNC_URL}"
  TMP_STATE_JSON="${STATE_JSON}.tmp"
  rm -f "${TMP_STATE_JSON}"
  curl -fsSL \
    -H "X-Gateway-Bootstrap-Token: ${GATEWAY_BOOTSTRAP_TOKEN}" \
    "${SYNC_URL}" \
    -o "${TMP_STATE_JSON}" || return 1
  mv "${TMP_STATE_JSON}" "${STATE_JSON}"

  INTERFACE_ADDRESS=$(jq -r '.interface.address' "${STATE_JSON}")
  NETWORK_CIDR=$(jq -r '.interface.network_cidr' "${STATE_JSON}")
  LISTEN_PORT=$(jq -r '.interface.listen_port' "${STATE_JSON}")

  cat > "${WG_GENERATED}" <<EOF
[Interface]
Address = ${INTERFACE_ADDRESS}
ListenPort = ${LISTEN_PORT}
PrivateKey = ${NEXAVPN_GATEWAY_PRIVATE_KEY}

EOF

  jq -c '.peers[]?' "${STATE_JSON}" | while read -r peer; do
    PUBLIC_KEY=$(printf '%s' "${peer}" | jq -r '.public_key')
    ASSIGNED_IP=$(printf '%s' "${peer}" | jq -r '.assigned_ip')

    cat >> "${WG_GENERATED}" <<EOF
[Peer]
PublicKey = ${PUBLIC_KEY}
AllowedIPs = ${ASSIGNED_IP}

EOF
  done

  cp "${WG_GENERATED}" "${WG_CONF}"

  {
    echo "flush ruleset"
    echo "table inet nexavpn {"
    echo "  chain forward {"
    echo "    type filter hook forward priority 0;"
    echo "    policy drop;"
    echo "    ct state established,related accept"
    echo "    iifname \"${IFACE}\" ip saddr ${NETWORK_CIDR} oifname \"${UPLINK_IFACE}\" accept"

    jq -c '.peers[]?' "${STATE_JSON}" | while read -r peer; do
      ASSIGNED_IP=$(printf '%s' "${peer}" | jq -r '.assigned_ip')
      printf '%s' "${peer}" | jq -r '.allowed_destinations[]?' | while read -r destination; do
        echo "    iifname \"${IFACE}\" ip saddr ${ASSIGNED_IP} ip daddr ${destination} accept"
      done
    done

    echo "  }"
    if [ "${ENABLE_MASQUERADE}" = "true" ]; then
      echo "  chain postrouting {"
      echo "    type nat hook postrouting priority 100;"
      echo "    oifname \"${UPLINK_IFACE}\" ip saddr ${NETWORK_CIDR} masquerade"
      echo "  }"
    fi
    echo "}"
  } > "${NFT_CONF}"

  if [ -w /proc/sys/net/ipv4/ip_forward ]; then
    sysctl -w net.ipv4.ip_forward=1 >/dev/null || true
  fi

  nft -f "${NFT_CONF}"

  if ip link show "${IFACE}" >/dev/null 2>&1; then
    wg syncconf "${IFACE}" <(wg-quick strip "${WG_CONF}")
    ip link set "${IFACE}" up
  else
    wg-quick up "${WG_CONF}"
  fi

  echo "Applied WireGuard config from ${WG_CONF}"
  echo "Applied nftables config from ${NFT_CONF}"
}

while true; do
  apply_bundle || echo "Gateway apply failed; retrying in 15 seconds"
  sleep 15
done