#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
if [[ -f "${SCRIPT_DIR}/../lib/common.sh" ]]; then
  # shellcheck source=../lib/common.sh
  source "${SCRIPT_DIR}/../lib/common.sh"
elif [[ -f "/usr/local/lib/wireguard-automation/lib/common.sh" ]]; then
  # shellcheck source=/usr/local/lib/wireguard-automation/lib/common.sh
  source "/usr/local/lib/wireguard-automation/lib/common.sh"
else
  echo "Не найден common.sh" >&2
  exit 1
fi

LOG_FILE="/var/log/wireguard-peerctl.log"
WG_META_FILE="/etc/wireguard/wg-meta.env"

usage() {
  cat <<'USAGE'
Использование:
  wg-peerctl.sh add \
    --client-name <name> \
    --client-public-key <pubkey> \
    [--client-address <10.66.66.X/32>] \
    [--client-preshared-key <psk>] \
    [--persistent-keepalive 25]

Описание:
  Скрипт добавляет peer в конфигурацию WireGuard-сервера идемпотентно.
  Если peer с таким public key уже существует, повторно не добавляет.
USAGE
}

load_meta() {
  [[ -f "$WG_META_FILE" ]] || die "Не найден $WG_META_FILE. Сначала выполните install_server.sh"
  # shellcheck disable=SC1090
  source "$WG_META_FILE"

  WG_INTERFACE="${WG_INTERFACE:-wg0}"
  WG_NETWORK="${WG_NETWORK:-10.66.66.0/24}"
  WG_ADDRESS="${WG_ADDRESS:-10.66.66.1/24}"
  WG_PORT="${WG_PORT:-51820}"
  SERVER_PUBLIC_IP="${SERVER_PUBLIC_IP:-}"
  SERVER_DNS="${SERVER_DNS:-1.1.1.1}"
  WG_CONF="/etc/wireguard/${WG_INTERFACE}.conf"
}

ip_to_int() {
  local ip="$1"
  local o1 o2 o3 o4
  IFS='.' read -r o1 o2 o3 o4 <<< "$ip"
  echo $(( (o1 << 24) + (o2 << 16) + (o3 << 8) + o4 ))
}

int_to_ip() {
  local n="$1"
  printf '%d.%d.%d.%d' \
    $(( (n >> 24) & 255 )) \
    $(( (n >> 16) & 255 )) \
    $(( (n >> 8) & 255 )) \
    $(( n & 255 ))
}

cidr_bounds() {
  local cidr="$1"
  local ip prefix
  IFS='/' read -r ip prefix <<< "$cidr"
  [[ -n "$ip" && -n "$prefix" ]] || return 1

  local ip_int mask net broadcast
  ip_int="$(ip_to_int "$ip")"
  if ((prefix == 0)); then
    mask=0
  else
    mask=$(( (0xFFFFFFFF << (32 - prefix)) & 0xFFFFFFFF ))
  fi
  net=$(( ip_int & mask ))
  broadcast=$(( net | ((~mask) & 0xFFFFFFFF) ))
  echo "${net} ${broadcast}"
}

next_client_ip() {
  local network="$1"
  local net_start net_end
  read -r net_start net_end < <(cidr_bounds "$network") || return 1
  ((net_end - net_start >= 3)) || return 1

  local server_ip server_ip_int
  server_ip="${WG_ADDRESS%%/*}"
  server_ip_int="$(ip_to_int "$server_ip")"

  local first_host last_host
  first_host=$((net_start + 1))
  last_host=$((net_end - 1))

  local used
  used="$(grep -E '^AllowedIPs\s*=\s*' "$WG_CONF" | awk -F'=' '{print $2}' | tr ',' '\n' | sed 's/ //g' | grep -E '^([0-9]{1,3}\.){3}[0-9]{1,3}/32$' | sed 's#/32##' || true)"

  local candidate_int candidate_ip
  for ((candidate_int = first_host; candidate_int <= last_host; candidate_int++)); do
    ((candidate_int == server_ip_int)) && continue
    candidate_ip="$(int_to_ip "$candidate_int")"
    if ! grep -qx "$candidate_ip" <<< "$used"; then
      echo "${candidate_ip}/32"
      return 0
    fi
  done
  return 1
}

peer_exists_by_pubkey() {
  local pubkey="$1"
  grep -Fq "PublicKey = $pubkey" "$WG_CONF"
}

extract_peer_address_by_pubkey() {
  local pubkey="$1"
  awk -v pk="$pubkey" '
    $0 ~ /^\[Peer\]/ {in_peer=1; key=""; addr=""}
    in_peer && $0 ~ /^PublicKey[[:space:]]*=/ {
      sub(/^[^=]*=[[:space:]]*/, "", $0); key=$0
    }
    in_peer && $0 ~ /^AllowedIPs[[:space:]]*=/ {
      sub(/^[^=]*=[[:space:]]*/, "", $0); addr=$0
    }
    in_peer && key==pk && addr!="" {print addr; exit}
  ' "$WG_CONF" | awk -F',' '{print $1}' | xargs
}

apply_config() {
  if systemctl is-active --quiet "wg-quick@${WG_INTERFACE}"; then
    wg syncconf "$WG_INTERFACE" <(wg-quick strip "$WG_CONF")
  else
    systemctl restart "wg-quick@${WG_INTERFACE}"
  fi
}

cmd_add() {
  local client_name=""
  local client_pubkey=""
  local client_address=""
  local client_psk=""
  local keepalive="25"

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --client-name)
        client_name="$2"; shift 2 ;;
      --client-public-key)
        client_pubkey="$2"; shift 2 ;;
      --client-address)
        client_address="$2"; shift 2 ;;
      --client-preshared-key)
        client_psk="$2"; shift 2 ;;
      --persistent-keepalive)
        keepalive="$2"; shift 2 ;;
      *)
        die "Неизвестный аргумент: $1"
        ;;
    esac
  done

  [[ -n "$client_name" ]] || die "Не указан --client-name"
  [[ -n "$client_pubkey" ]] || die "Не указан --client-public-key"

  client_name="$(sanitize_name "$client_name")"
  [[ -n "$client_name" ]] || die "Некорректное имя клиента"

  load_meta
  [[ -f "$WG_CONF" ]] || die "Не найден конфиг WireGuard: $WG_CONF"

  local server_pubkey
  server_pubkey="$(cat /etc/wireguard/server_public.key)"

  if peer_exists_by_pubkey "$client_pubkey"; then
    local existing_addr
    existing_addr="$(extract_peer_address_by_pubkey "$client_pubkey")"
    cat <<EOF_OUT
STATUS=exists
CLIENT_NAME=$client_name
CLIENT_ADDRESS=${existing_addr:-unknown}
SERVER_PUBLIC_KEY=$server_pubkey
SERVER_ENDPOINT=${SERVER_PUBLIC_IP}:${WG_PORT}
SERVER_DNS=${SERVER_DNS}
WG_INTERFACE=${WG_INTERFACE}
EOF_OUT
    return 0
  fi

  if [[ -z "$client_address" ]]; then
    client_address="$(next_client_ip "$WG_NETWORK")" || die "Не удалось выделить IP клиенту в сети $WG_NETWORK"
  fi

  backup_file "$WG_CONF"

  {
    echo
    echo "# managed-by=wg-peerctl client=${client_name} created_at=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
    echo "[Peer]"
    echo "PublicKey = ${client_pubkey}"
    if [[ -n "$client_psk" ]]; then
      echo "PresharedKey = ${client_psk}"
    fi
    echo "AllowedIPs = ${client_address}"
    echo "PersistentKeepalive = ${keepalive}"
  } >> "$WG_CONF"

  apply_config

  cat <<EOF_OUT
STATUS=created
CLIENT_NAME=$client_name
CLIENT_ADDRESS=$client_address
SERVER_PUBLIC_KEY=$server_pubkey
SERVER_ENDPOINT=${SERVER_PUBLIC_IP}:${WG_PORT}
SERVER_DNS=${SERVER_DNS}
WG_INTERFACE=${WG_INTERFACE}
EOF_OUT
}

main() {
  local cmd="${1:-}"
  if [[ -z "$cmd" ]]; then
    usage
    exit 0
  fi
  shift || true

  case "$cmd" in
    add)
      require_root
      check_os_supported
      cmd_add "$@"
      ;;
    -h|--help|help)
      usage
      ;;
    *)
      die "Неизвестная команда: $cmd"
      ;;
  esac
}

main "$@"