#!/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"
GUI_DB_FILE="/opt/wg-admin-gui/data/wgadmin.db"

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

  wg-peerctl.sh remove \
    --client-public-key <pubkey>

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

sql_escape() {
  local s="$1"
  s="${s//\'/\'\'}"
  printf "%s" "$s"
}

ensure_gui_db_schema() {
  command -v sqlite3 >/dev/null 2>&1 || return 0
  [[ -f "$GUI_DB_FILE" ]] || return 0
  sqlite3 "$GUI_DB_FILE" <<'SQL' >/dev/null 2>&1 || true
CREATE TABLE IF NOT EXISTS peers (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  name TEXT NOT NULL,
  public_key TEXT UNIQUE NOT NULL,
  client_address TEXT,
  advertised_routes TEXT,
  client_conf TEXT,
  peer_psk TEXT,
  peer_allowed_ips TEXT,
  enabled INTEGER NOT NULL DEFAULT 1,
  created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
ALTER TABLE peers ADD COLUMN client_conf TEXT;
ALTER TABLE peers ADD COLUMN peer_psk TEXT;
ALTER TABLE peers ADD COLUMN peer_allowed_ips TEXT;
ALTER TABLE peers ADD COLUMN enabled INTEGER NOT NULL DEFAULT 1;
SQL
}

sync_gui_db_upsert_peer() {
  local name="$1"
  local pubkey="$2"
  local address="$3"
  local routes="$4"
  local psk="$5"
  local peer_allowed_ips="$6"
  local enabled="${7:-1}"

  command -v sqlite3 >/dev/null 2>&1 || return 0
  [[ -f "$GUI_DB_FILE" ]] || return 0
  ensure_gui_db_schema

  local e_name e_pub e_addr e_routes e_psk e_allowed
  e_name="$(sql_escape "$name")"
  e_pub="$(sql_escape "$pubkey")"
  e_addr="$(sql_escape "$address")"
  e_routes="$(sql_escape "$routes")"
  e_psk="$(sql_escape "$psk")"
  e_allowed="$(sql_escape "$peer_allowed_ips")"

  sqlite3 "$GUI_DB_FILE" <<SQL >/dev/null 2>&1 || true
INSERT INTO peers(name, public_key, client_address, advertised_routes, peer_psk, peer_allowed_ips, enabled)
VALUES ('$e_name', '$e_pub', '$e_addr', '$e_routes', '$e_psk', '$e_allowed', $enabled)
ON CONFLICT(public_key)
DO UPDATE SET
  name=excluded.name,
  client_address=excluded.client_address,
  advertised_routes=excluded.advertised_routes,
  peer_psk=excluded.peer_psk,
  peer_allowed_ips=excluded.peer_allowed_ips,
  enabled=excluded.enabled;
SQL
}

sync_gui_db_set_enabled() {
  local pubkey="$1"
  local enabled="$2"
  command -v sqlite3 >/dev/null 2>&1 || return 0
  [[ -f "$GUI_DB_FILE" ]] || return 0
  local e_pub
  e_pub="$(sql_escape "$pubkey")"
  sqlite3 "$GUI_DB_FILE" "UPDATE peers SET enabled=${enabled} WHERE public_key='${e_pub}';" >/dev/null 2>&1 || true
}

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
}

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

routes_without_primary_address() {
  local allowed_ips="$1"
  local primary_addr="$2"
  local out=""
  local item
  local norm_primary
  norm_primary="$(echo "$primary_addr" | xargs)"

  IFS=',' read -ra items <<< "$allowed_ips"
  for item in "${items[@]}"; do
    item="$(echo "$item" | xargs)"
    [[ -z "$item" ]] && continue
    if [[ -n "$norm_primary" && "$item" == "$norm_primary" ]]; then
      continue
    fi
    if [[ -z "$out" ]]; then
      out="$item"
    else
      out="${out},${item}"
    fi
  done

  echo "$out"
}

apply_client_routes_now() {
  local routes="${1:-}"
  [[ -n "$routes" ]] || return 0

  local cidr
  IFS=',' read -ra cidrs <<< "$routes"
  for cidr in "${cidrs[@]}"; do
    cidr="$(echo "$cidr" | xargs)"
    [[ -n "$cidr" ]] || continue
    ip route replace "$cidr" dev "$WG_INTERFACE" proto static >/dev/null 2>&1 || true
  done
}

remove_client_routes_now() {
  local routes="${1:-}"
  [[ -n "$routes" ]] || return 0

  local cidr
  IFS=',' read -ra cidrs <<< "$routes"
  for cidr in "${cidrs[@]}"; do
    cidr="$(echo "$cidr" | xargs)"
    [[ -n "$cidr" ]] || continue
    ip route del "$cidr" dev "$WG_INTERFACE" >/dev/null 2>&1 || true
  done
}

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_routes=""
  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-routes)
        client_routes="$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
    local existing_allowed
    existing_addr="$(extract_peer_address_by_pubkey "$client_pubkey")"
    existing_allowed="${existing_addr:-}"
    if [[ -n "$client_routes" ]]; then
      existing_allowed="${existing_allowed},${client_routes}"
      apply_client_routes_now "$client_routes"
    fi
    sync_gui_db_upsert_peer "$client_name" "$client_pubkey" "${existing_addr:-}" "$client_routes" "$client_psk" "${existing_allowed}" 1
    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}
WG_NETWORK=${WG_NETWORK}
EOF_OUT
    return 0
  fi

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

  local peer_allowed_ips="$client_address"
  if [[ -n "$client_routes" ]]; then
    is_valid_cidr_list "$client_routes" || die "Некорректный список --client-routes"
    client_routes="$(echo "$client_routes" | tr -d ' ')"
    peer_allowed_ips="${peer_allowed_ips},${client_routes}"
  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 = ${peer_allowed_ips}"
    echo "PersistentKeepalive = ${keepalive}"
  } >> "$WG_CONF"

  apply_config
  apply_client_routes_now "$client_routes"
  sync_gui_db_upsert_peer "$client_name" "$client_pubkey" "$client_address" "$client_routes" "$client_psk" "$peer_allowed_ips" 1

  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}
WG_NETWORK=${WG_NETWORK}
EOF_OUT
}

cmd_remove() {
  local client_pubkey=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --client-public-key)
        client_pubkey="$2"; shift 2 ;;
      *)
        die "Неизвестный аргумент: $1"
        ;;
    esac
  done
  [[ -n "$client_pubkey" ]] || die "Не указан --client-public-key"

  load_meta
  [[ -f "$WG_CONF" ]] || die "Не найден конфиг WireGuard: $WG_CONF"
  local existing_allowed existing_addr existing_routes
  existing_allowed="$(extract_peer_allowed_ips_by_pubkey "$client_pubkey")"
  existing_addr="$(extract_peer_address_by_pubkey "$client_pubkey")"
  existing_routes="$(routes_without_primary_address "$existing_allowed" "$existing_addr")"
  backup_file "$WG_CONF"

  local tmp
  tmp="$(mktemp)"
  awk -v pk="$client_pubkey" '
    BEGIN {in=0; block=""; keep=1}
    /^\[Peer\]/ {
      if (in && keep) printf "%s", block
      in=1; block=$0 ORS; keep=1; next
    }
    {
      if (in) {
        block = block $0 ORS
        if ($0 ~ /^PublicKey[[:space:]]*=/) {
          line=$0
          sub(/^[^=]*=[[:space:]]*/, "", line)
          if (line == pk) keep=0
        }
        next
      }
      print
    }
    END {
      if (in && keep) printf "%s", block
    }
  ' "$WG_CONF" > "$tmp"
  mv "$tmp" "$WG_CONF"
  safe_chmod_600 "$WG_CONF"

  apply_config
  remove_client_routes_now "$existing_routes"
  sync_gui_db_set_enabled "$client_pubkey" 0

  cat <<EOF_OUT
STATUS=removed
PUBLIC_KEY=${client_pubkey}
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 "$@"
      ;;
    remove)
      require_root
      check_os_supported
      cmd_remove "$@"
      ;;
    -h|--help|help)
      usage
      ;;
    *)
      die "Неизвестная команда: $cmd"
      ;;
  esac
}

main "$@"