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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
# shellcheck source=../lib/common.sh
source "${PROJECT_ROOT}/lib/common.sh"

LOG_FILE="/var/log/wireguard-client-install.log"
NON_INTERACTIVE=0

WG_INTERFACE="wg0"
CLIENT_NAME=""
CLIENT_DNS=""
TUNNEL_MODE=""
SPLIT_ALLOWED_IPS=""
ADVERTISE_SUBNETS=""

SERVER_HOST=""
SERVER_USER="root"
SSH_PORT="22"
SSH_AUTH_METHOD="key"
SSH_PASSWORD=""

KEEPALIVE="25"
CLIENT_ADDRESS_PREFIX="24"

usage() {
  cat <<'USAGE'
Установка WireGuard-клиента и автоматическая регистрация на сервере.
Каждый запуск выполняет полный reset предыдущей клиентской конфигурации для выбранного интерфейса.

Использование:
  install_client.sh [опции]

Опции:
  --non-interactive              Режим без вопросов
  --interface <name>             Интерфейс клиента (по умолчанию: wg0)
  --client-name <name>           Имя клиента (по умолчанию: hostname)
  --mode <full|split>            full: весь трафик, split: только выбранные сети
  --allowed-ips <cidr,cidr>      Для режима split
  --dns <ip>                     DNS для клиента (если не задан, берется с сервера)
  --client-address-prefix <1-32> Префикс маски адреса интерфейса клиента (по умолчанию: 24)
  --advertise-subnets <cidr,...> Сети за клиентом, которые нужно маршрутизировать через него (например 192.168.33.0/24)

  --server-host <host>           IP/домен WireGuard-сервера
  --server-user <user>           SSH-пользователь (по умолчанию: root)
  --ssh-port <port>              SSH-порт (по умолчанию: 22)
  --ssh-auth <key|password>      Тип SSH-аутентификации (по умолчанию: key)
  --ssh-password <password>      Пароль SSH (используйте осторожно)

  -h, --help                     Показать помощь
USAGE
}

on_error() {
  local code=$?
  log_error "Клиентская настройка завершилась с ошибкой (код: ${code}). Лог: ${LOG_FILE}"
  exit "$code"
}
trap on_error ERR

parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --non-interactive)
        NON_INTERACTIVE=1; shift ;;
      --interface)
        WG_INTERFACE="$2"; shift 2 ;;
      --client-name)
        CLIENT_NAME="$2"; shift 2 ;;
      --mode)
        TUNNEL_MODE="$2"; shift 2 ;;
      --allowed-ips)
        SPLIT_ALLOWED_IPS="$2"; shift 2 ;;
      --dns)
        CLIENT_DNS="$2"; shift 2 ;;
      --client-address-prefix)
        CLIENT_ADDRESS_PREFIX="$2"; shift 2 ;;
      --advertise-subnets)
        ADVERTISE_SUBNETS="$2"; shift 2 ;;
      --server-host)
        SERVER_HOST="$2"; shift 2 ;;
      --server-user)
        SERVER_USER="$2"; shift 2 ;;
      --ssh-port)
        SSH_PORT="$2"; shift 2 ;;
      --ssh-auth)
        SSH_AUTH_METHOD="$2"; shift 2 ;;
      --ssh-password)
        SSH_PASSWORD="$2"; shift 2 ;;
      -h|--help)
        usage; exit 0 ;;
      *)
        die "Неизвестный аргумент: $1"
        ;;
    esac
  done
}

validate_inputs() {
  is_valid_port "$SSH_PORT" || die "Некорректный SSH-порт: $SSH_PORT"
  [[ "$SSH_AUTH_METHOD" == "key" || "$SSH_AUTH_METHOD" == "password" ]] || die "--ssh-auth должен быть key или password"
  [[ "$TUNNEL_MODE" == "full" || "$TUNNEL_MODE" == "split" ]] || die "--mode должен быть full или split"
  [[ "$CLIENT_ADDRESS_PREFIX" =~ ^[0-9]+$ ]] || die "--client-address-prefix должен быть числом 1..32"
  ((CLIENT_ADDRESS_PREFIX >= 1 && CLIENT_ADDRESS_PREFIX <= 32)) || die "--client-address-prefix должен быть в диапазоне 1..32"
  if [[ -n "$ADVERTISE_SUBNETS" ]]; then
    is_valid_cidr_list "$ADVERTISE_SUBNETS" || die "Некорректный список --advertise-subnets"
  fi

  if [[ -z "$SERVER_HOST" ]]; then
    die "Не указан --server-host"
  fi

  if [[ -n "$SPLIT_ALLOWED_IPS" ]]; then
    is_valid_cidr_list "$SPLIT_ALLOWED_IPS" || die "Некорректный список --allowed-ips"
  fi
}

collect_inputs() {
  if [[ -z "$CLIENT_NAME" ]]; then
    CLIENT_NAME="$(hostname -s)"
  fi
  CLIENT_NAME="$(sanitize_name "$CLIENT_NAME")"
  [[ -n "$CLIENT_NAME" ]] || die "Некорректное имя клиента"

  if [[ -z "$SERVER_HOST" && ! $NON_INTERACTIVE -eq 1 ]]; then
    read -r -p "Введите IP/домен сервера: " SERVER_HOST
  fi

  if [[ -z "$TUNNEL_MODE" ]]; then
    if ((NON_INTERACTIVE)); then
      TUNNEL_MODE="full"
    else
      echo "Выберите режим маршрутизации:"
      echo "  1) весь трафик через VPN (full)"
      echo "  2) только выбранные сети (split)"
      read -r -p "Ваш выбор [1/2, по умолчанию 1]: " mode_choice
      case "${mode_choice:-1}" in
        1) TUNNEL_MODE="full" ;;
        2) TUNNEL_MODE="split" ;;
        *) die "Некорректный выбор" ;;
      esac
    fi
  fi

  if [[ "$SSH_AUTH_METHOD" == "password" && -z "$SSH_PASSWORD" ]]; then
    if ((NON_INTERACTIVE)); then
      die "Для --ssh-auth password нужно указать --ssh-password"
    fi
    ask_secret "Введите SSH-пароль для ${SERVER_USER}@${SERVER_HOST}" SSH_PASSWORD
  fi

  validate_inputs
}

detect_client_lan_subnets() {
  ip -o -4 route show scope link 2>/dev/null | awk '
    {
      cidr=$1
      dev=""
      for (i=1; i<=NF; i++) {
        if ($i=="dev") { dev=$(i+1); break }
      }
      if (cidr ~ /^127\./ || cidr ~ /^169\.254\./) next
      if (dev ~ /^(lo|wg|docker|veth|br-|tun|tap)/) next
      if (cidr ~ /^([0-9]{1,3}\.){3}[0-9]{1,3}\/[0-9]+$/) print cidr
    }
  ' | sort -u | paste -sd, -
}

install_packages() {
  pkg_has_install_candidate() {
    local pkg="$1"
    local candidate
    candidate="$(apt-cache policy "$pkg" 2>/dev/null | awk '/Candidate:/ {print $2; exit}')"
    [[ -n "$candidate" && "$candidate" != "(none)" ]]
  }

  apt_install_if_missing wireguard wireguard-tools iproute2 openssh-client sshpass ca-certificates

  if ! command -v resolvconf >/dev/null 2>&1; then
    if pkg_has_install_candidate openresolv; then
      apt_install_if_missing openresolv
    elif pkg_has_install_candidate resolvconf; then
      apt_install_if_missing resolvconf
    else
      log_warn "Пакеты openresolv/resolvconf недоступны. DNS-строка в wg0.conf будет пропущена."
    fi
  fi
}

reset_existing_client_install() {
  log_warn "Выполняю полный reset предыдущей конфигурации клиента (${WG_INTERFACE})"

  local unit="wg-quick@${WG_INTERFACE}.service"
  systemctl disable --now "$unit" >/dev/null 2>&1 || true
  wg-quick down "$WG_INTERFACE" >/dev/null 2>&1 || true

  rm -f \
    "/etc/wireguard/${WG_INTERFACE}.conf" \
    "/etc/wireguard/${WG_INTERFACE}_client_private.key" \
    "/etc/wireguard/${WG_INTERFACE}_client_public.key" \
    "/etc/wireguard/${WG_INTERFACE}_client_psk.key"

  log_info "Старые конфиг/ключи клиента удалены для интерфейса ${WG_INTERFACE}"
}

ssh_base_cmd() {
  echo "ssh -o StrictHostKeyChecking=accept-new -o ConnectTimeout=10 -p ${SSH_PORT} ${SERVER_USER}@${SERVER_HOST}"
}

run_ssh() {
  local remote_cmd="$1"
  local base
  base="$(ssh_base_cmd)"

  if [[ "$SSH_AUTH_METHOD" == "password" ]]; then
    require_cmd sshpass
    SSHPASS="$SSH_PASSWORD" sshpass -e bash -c "$base \"$remote_cmd\""
  else
    bash -c "$base \"$remote_cmd\""
  fi
}

resolve_server_ip() {
  getent ahostsv4 "$SERVER_HOST" | awk '{print $1; exit}'
}

generate_keys() {
  mkdir -p /etc/wireguard
  chmod 700 /etc/wireguard

  CLIENT_PRIV_KEY_PATH="/etc/wireguard/${WG_INTERFACE}_client_private.key"
  CLIENT_PUB_KEY_PATH="/etc/wireguard/${WG_INTERFACE}_client_public.key"
  CLIENT_PSK_PATH="/etc/wireguard/${WG_INTERFACE}_client_psk.key"

  if [[ ! -f "$CLIENT_PRIV_KEY_PATH" ]]; then
    umask 077
    wg genkey | tee "$CLIENT_PRIV_KEY_PATH" | wg pubkey > "$CLIENT_PUB_KEY_PATH"
    wg genpsk > "$CLIENT_PSK_PATH"
    log_success "Сгенерированы клиентские ключи"
  else
    if [[ ! -f "$CLIENT_PUB_KEY_PATH" ]]; then
      wg pubkey < "$CLIENT_PRIV_KEY_PATH" > "$CLIENT_PUB_KEY_PATH"
    fi
    if [[ ! -f "$CLIENT_PSK_PATH" ]]; then
      wg genpsk > "$CLIENT_PSK_PATH"
    fi
    log_info "Ключи клиента уже существуют, переиспользую"
  fi

  safe_chmod_600 "$CLIENT_PRIV_KEY_PATH"
  safe_chmod_600 "$CLIENT_PUB_KEY_PATH"
  safe_chmod_600 "$CLIENT_PSK_PATH"
}

register_peer_on_server() {
  local pub psk remote_cmd response
  pub="$(cat "$CLIENT_PUB_KEY_PATH")"
  psk="$(cat "$CLIENT_PSK_PATH")"

  remote_cmd="/usr/local/sbin/wg-peerctl add --client-name '${CLIENT_NAME}' --client-public-key '${pub}' --client-preshared-key '${psk}' --persistent-keepalive '${KEEPALIVE}'"
  if [[ -n "$ADVERTISE_SUBNETS" ]]; then
    remote_cmd+=" --client-routes '${ADVERTISE_SUBNETS}'"
  fi
  response="$(run_ssh "$remote_cmd")"

  SERVER_RESPONSE_RAW="$response"

  kv_get() {
    local key="$1"
    printf '%s\n' "$response" | sed -n "s/^${key}=//p" | head -n1
  }

  SERVER_STATUS="$(kv_get STATUS)"
  CLIENT_ADDRESS="$(kv_get CLIENT_ADDRESS)"
  SERVER_PUBLIC_KEY="$(kv_get SERVER_PUBLIC_KEY)"
  SERVER_ENDPOINT="$(kv_get SERVER_ENDPOINT)"
  SERVER_DNS_REMOTE="$(kv_get SERVER_DNS)"
  SERVER_WG_NETWORK="$(kv_get WG_NETWORK)"

  [[ -n "$SERVER_PUBLIC_KEY" ]] || die "Сервер не вернул SERVER_PUBLIC_KEY"
  [[ -n "$CLIENT_ADDRESS" ]] || die "Сервер не вернул CLIENT_ADDRESS"
  [[ -n "$SERVER_ENDPOINT" ]] || die "Сервер не вернул SERVER_ENDPOINT"
}

build_allowed_ips() {
  if [[ "$TUNNEL_MODE" == "full" ]]; then
    ALLOWED_IPS="0.0.0.0/0,::/0"
  else
    if [[ -n "$SPLIT_ALLOWED_IPS" ]]; then
      ALLOWED_IPS="$SPLIT_ALLOWED_IPS"
    elif [[ -n "${SERVER_WG_NETWORK:-}" ]]; then
      ALLOWED_IPS="$SERVER_WG_NETWORK"
      log_info "Режим split: --allowed-ips не задан, использую сеть WG сервера: ${ALLOWED_IPS}"
    else
      die "Режим split: не удалось определить --allowed-ips автоматически"
    fi
  fi
}

build_route_hooks_if_needed() {
  PRE_UP=""
  POST_DOWN=""

  if [[ "$TUNNEL_MODE" != "full" ]]; then
    return
  fi

  local server_ip gw iface
  server_ip="$(resolve_server_ip || true)"
  gw="$(detect_default_gateway || true)"
  iface="$(detect_default_iface || true)"

  if [[ -n "$server_ip" && -n "$gw" && -n "$iface" ]]; then
    PRE_UP="PreUp = ip route add ${server_ip}/32 via ${gw} dev ${iface} 2>/dev/null || true"
    POST_DOWN="PostDown = ip route del ${server_ip}/32 via ${gw} dev ${iface} 2>/dev/null || true"
  else
    log_warn "Не удалось вычислить route-exception до сервера. Использую стандартное поведение wg-quick."
  fi
}

build_client_interface_address() {
  local ip_only
  ip_only="${CLIENT_ADDRESS%%/*}"
  CLIENT_INTERFACE_ADDRESS="${ip_only}/${CLIENT_ADDRESS_PREFIX}"
}

write_client_config() {
  local conf="/etc/wireguard/${WG_INTERFACE}.conf"
  local dns

  dns="$CLIENT_DNS"
  if [[ -z "$dns" ]]; then
    dns="${SERVER_DNS_REMOTE:-1.1.1.1}"
  fi

  if [[ -f "$conf" ]]; then
    backup_file "$conf"
  fi

  {
    echo "[Interface]"
    echo "PrivateKey = $(cat "$CLIENT_PRIV_KEY_PATH")"
    echo "Address = ${CLIENT_INTERFACE_ADDRESS}"
    if command -v resolvconf >/dev/null 2>&1; then
      echo "DNS = ${dns}"
    else
      log_warn "Команда resolvconf не найдена. Пропускаю строку DNS в конфиге WireGuard."
    fi
    if [[ -n "$PRE_UP" ]]; then
      echo "$PRE_UP"
    fi
    if [[ -n "$POST_DOWN" ]]; then
      echo "$POST_DOWN"
    fi
    echo
    echo "[Peer]"
    echo "PublicKey = ${SERVER_PUBLIC_KEY}"
    echo "PresharedKey = $(cat "$CLIENT_PSK_PATH")"
    echo "Endpoint = ${SERVER_ENDPOINT}"
    echo "AllowedIPs = ${ALLOWED_IPS}"
    echo "PersistentKeepalive = ${KEEPALIVE}"
  } > "$conf"

  safe_chmod_600 "$conf"
}

apply_client_config() {
  local unit="wg-quick@${WG_INTERFACE}.service"
  if systemctl list-unit-files | grep -q "^${unit}"; then
    systemctl restart "$unit" || true
  fi
  systemd_enable_now "$unit"
}

print_routes_info() {
  log_info "Текущая таблица маршрутов клиента (до применения уже была прочитана системой):"
  ip route show | sed 's/^/  /'
}

print_summary() {
  local st
  st="$(systemctl is-active "wg-quick@${WG_INTERFACE}" 2>/dev/null || true)"

  cat <<EOF_SUMMARY

================ ИТОГОВАЯ СВОДКА ================
Клиент:               ${CLIENT_NAME}
Интерфейс:            ${WG_INTERFACE}
Сервис:               ${st}
Конфиг клиента:       /etc/wireguard/${WG_INTERFACE}.conf
Endpoint сервера:     ${SERVER_ENDPOINT}
Маршруты (AllowedIPs):${ALLOWED_IPS}
Адрес интерфейса:     ${CLIENT_INTERFACE_ADDRESS}
Режим туннеля:        ${TUNNEL_MODE}
SSH сервер:           ${SERVER_USER}@${SERVER_HOST}:${SSH_PORT}
Статус регистрации:   ${SERVER_STATUS}
Сети за клиентом:     ${ADVERTISE_SUBNETS:-не объявлены}
Лог:                  ${LOG_FILE}
=================================================
EOF_SUMMARY
}

main() {
  parse_args "$@"

  require_root
  check_os_supported
  require_cmd awk
  require_cmd sed
  require_cmd ip
  require_cmd ssh

  collect_inputs

  reset_existing_client_install
  print_routes_info
  install_packages
  generate_keys
  if [[ -z "$ADVERTISE_SUBNETS" ]]; then
    ADVERTISE_SUBNETS="$(detect_client_lan_subnets || true)"
    if [[ -n "$ADVERTISE_SUBNETS" ]]; then
      log_info "Автоопределены сети за клиентом: ${ADVERTISE_SUBNETS}"
    fi
  fi
  register_peer_on_server
  build_allowed_ips
  build_client_interface_address
  build_route_hooks_if_needed
  write_client_config
  apply_client_config
  print_summary

  # Не держим пароль в переменной дольше необходимого.
  unset SSH_PASSWORD
  log_success "Настройка клиента завершена"
}

main "$@"