#!/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}"
GATEWAY_LISTEN_PORT="${NEXAVPN_GATEWAY_LISTEN_PORT:-51900}"
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:${GATEWAY_LISTEN_PORT}}\",\"public_key\":\"${GATEWAY_PUBLIC_KEY}\",\"listen_port\":${GATEWAY_LISTEN_PORT},\"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 "table inet nexavpn {"
    echo "  chain forward {"
    echo "    type filter hook forward priority 0;"
    echo "    policy accept;"
    echo "    ct state established,related accept"
    echo "    iifname != \"${IFACE}\" 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 '.dns_servers[]?' | while read -r dns_server; do
        echo "    iifname \"${IFACE}\" ip saddr ${ASSIGNED_IP} ip daddr ${dns_server} udp dport 53 accept"
        echo "    iifname \"${IFACE}\" ip saddr ${ASSIGNED_IP} ip daddr ${dns_server} tcp dport 53 accept"
      done
      printf '%s' "${peer}" | jq -r '.web_proxy_targets[]?' | while read -r proxy_target; do
        echo "    iifname \"${IFACE}\" ip saddr ${ASSIGNED_IP} ip daddr ${proxy_target} tcp dport 80 accept"
        echo "    iifname \"${IFACE}\" ip saddr ${ASSIGNED_IP} ip daddr ${proxy_target} tcp dport 443 accept"
      done
      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 "    iifname \"${IFACE}\" drop"

    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 delete table inet nexavpn >/dev/null 2>&1 || true
  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}"
  wg show "${IFACE}" latest-handshakes transfer 2>/dev/null || true
  post_telemetry || true
}

post_telemetry() {
  if [ ! -f "${STATE_JSON}" ]; then
    return 0
  fi

  TELEMETRY_URL="${SYNC_BASE_URL}/${NEXAVPN_GATEWAY_ID}/telemetry"
  TMP_TELEMETRY_JSON="/tmp/nexavpn-gateway-telemetry.json"

  {
    printf '{\"collected_at\":\"%s\",\"peers\":[' "$(date -u +"%Y-%m-%dT%H:%M:%SZ")"
    FIRST=1

    while IFS="$(printf '\t')" read -r PUBLIC_KEY _PRESHARED _ENDPOINT _ALLOWED_IPS LATEST_HANDSHAKE RX_BYTES TX_BYTES _KEEPALIVE; do
      if [ -z "${PUBLIC_KEY:-}" ] || [ "${PUBLIC_KEY}" = "private_key" ]; then
        continue
      fi

      DEVICE_ID="$(jq -r --arg public_key "${PUBLIC_KEY}" '.peers[]? | select(.public_key == $public_key) | .device_id' "${STATE_JSON}" | head -n1)"
      if [ -z "${DEVICE_ID:-}" ] || [ "${DEVICE_ID}" = "null" ]; then
        continue
      fi

      if [ "${FIRST}" -eq 0 ]; then
        printf ','
      fi
      FIRST=0

      if [ "${LATEST_HANDSHAKE:-0}" -gt 0 ] 2>/dev/null; then
        printf '{\"device_id\":\"%s\",\"public_key\":\"%s\",\"rx_bytes\":%s,\"tx_bytes\":%s,\"latest_handshake_at\":%s}' \
          "${DEVICE_ID}" "${PUBLIC_KEY}" "${RX_BYTES:-0}" "${TX_BYTES:-0}" "${LATEST_HANDSHAKE}"
      else
        printf '{\"device_id\":\"%s\",\"public_key\":\"%s\",\"rx_bytes\":%s,\"tx_bytes\":%s}' \
          "${DEVICE_ID}" "${PUBLIC_KEY}" "${RX_BYTES:-0}" "${TX_BYTES:-0}"
      fi
    done < <(wg show "${IFACE}" dump 2>/dev/null | tail -n +2)

    printf ']}'
  } > "${TMP_TELEMETRY_JSON}"

  curl -fsSL \
    -H "Content-Type: application/json" \
    -H "X-Gateway-Bootstrap-Token: ${GATEWAY_BOOTSTRAP_TOKEN}" \
    -X POST \
    --data @"${TMP_TELEMETRY_JSON}" \
    "${TELEMETRY_URL}" >/dev/null
}

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