#!/usr/bin/env bash
# =============================================================================
# WireGuard VPS Tunnel — 1-Click Installer
# Автоматическая настройка туннеля между VPS (публичный IP) и домашним сервером
# (Raspberry Pi / Orange Pi за NAT, например Starlink)
# =============================================================================
set -euo pipefail

# ── Цвета для вывода ────────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m' # No Color

# ── Пути ────────────────────────────────────────────────────────────────────
CONFIG_DIR="/etc/wireguard-vps-tunnel"
CONFIG_FILE="$CONFIG_DIR/config"
WG_CONF="/etc/wireguard/wg0.conf"
WATCHDOG_BIN="/usr/local/bin/wg-watchdog.sh"
WATCHDOG_SERVICE="/etc/systemd/system/wg-watchdog.service"
WG_SERVICE="wg-quick@wg0"

# ── Значения по умолчанию ───────────────────────────────────────────────────
MODE=""                 # vps или client
VPS_PUBLIC_IP=""        # Публичный IP VPS
CLIENT_WG_IP="10.0.0.2" # WireGuard IP клиента
SERVER_WG_IP="10.0.0.1" # WireGuard IP сервера
WG_PORT="51820"         # Порт WireGuard
PORTS="80,443"          # Порты для проброса (только VPS)
KEEPALIVE="21"          # PersistentKeepalive (секунды)
INTERACTIVE=true

# ── Функции вывода ──────────────────────────────────────────────────────────
info()    { echo -e "${BLUE}[INFO]${NC}    $*"; }
success() { echo -e "${GREEN}[OK]${NC}      $*"; }
warn()    { echo -e "${YELLOW}[WARN]${NC}    $*"; }
error()   { echo -e "${RED}[ERROR]${NC}   $*"; }
header()  { echo -e "\n${BOLD}${CYAN}═══ $* ═══${NC}\n"; }

# ── Баннер ──────────────────────────────────────────────────────────────────
banner() {
    echo -e "${CYAN}${BOLD}"
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║     WireGuard VPS Tunnel — 1-Click Installer            ║"
    echo "║     VPS ↔ Домашний сервер (за NAT / Starlink)           ║"
    echo "╚══════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
}

# ── Проверка прав ──────────────────────────────────────────────────────────
check_root() {
    if [[ $EUID -ne 0 ]]; then
        error "Скрипт должен выполняться от root. Используйте: sudo $0"
        exit 1
    fi
}

# ── Определение ОС ─────────────────────────────────────────────────────────
detect_os() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS_ID="${ID}"
        OS_NAME="${PRETTY_NAME:-$NAME}"
    else
        error "Не удалось определить операционную систему"
        exit 1
    fi

    case "$OS_ID" in
        debian|ubuntu|armbian|raspbian)
            info "Обнаружена ОС: ${OS_NAME}"
            ;;
        *)
            warn "ОС '$OS_ID' не тестировалась. Продолжаем, но возможны проблемы."
            ;;
    esac
}

# ── Установка зависимостей ─────────────────────────────────────────────────
install_deps() {
    header "Установка зависимостей"

    if command -v wg &>/dev/null; then
        success "WireGuard уже установлен: $(wg --version 2>&1 | head -1)"
        return 0
    fi

    info "Установка wireguard-tools..."
    apt-get update -qq
    apt-get install -y -qq wireguard-tools iptables

    if command -v wg &>/dev/null; then
        success "WireGuard установлен успешно"
    else
        error "Не удалось установить WireGuard"
        exit 1
    fi
}

# ── Включение IP-форвардинга ───────────────────────────────────────────────
enable_ip_forward() {
    header "Включение IP-форвардинга"

    local sysctl_conf="/etc/sysctl.d/99-wireguard-vps-tunnel.conf"

    if [[ -f "$sysctl_conf" ]]; then
        info "IP-форвардинг уже настроен"
    else
        cat > "$sysctl_conf" <<'SYSCTL'
# Включение IP-форвардинга для WireGuard VPS Tunnel
net.ipv4.ip_forward = 1
net.ipv6.conf.all.forwarding = 1
SYSCTL
        sysctl --system &>/dev/null
        success "IP-форвардинг включён"
    fi
}

# ── Генерация ключей ───────────────────────────────────────────────────────
generate_keys() {
    header "Генерация ключей WireGuard"

    mkdir -p "$CONFIG_DIR"

    if [[ -f "$CONFIG_DIR/server_private.key" ]] && [[ -f "$CONFIG_DIR/client_private.key" ]]; then
        info "Ключи уже существуют, пропускаем генерацию"
    else
        # Генерируем ключи сервера
        wg genkey | tee "$CONFIG_DIR/server_private.key" | wg pubkey > "$CONFIG_DIR/server_public.key"
        chmod 600 "$CONFIG_DIR/server_private.key"

        # Генерируем ключи клиента
        wg genkey | tee "$CONFIG_DIR/client_private.key" | wg pubkey > "$CONFIG_DIR/client_public.key"
        chmod 600 "$CONFIG_DIR/client_private.key"

        success "Ключи сгенерированы"
    fi

    SERVER_PRIVATE_KEY=$(cat "$CONFIG_DIR/server_private.key")
    SERVER_PUBLIC_KEY=$(cat "$CONFIG_DIR/server_public.key")
    CLIENT_PRIVATE_KEY=$(cat "$CONFIG_DIR/client_private.key")
    CLIENT_PUBLIC_KEY=$(cat "$CONFIG_DIR/client_public.key")
}

# ── Загрузка существующей конфигурации ─────────────────────────────────────
load_config() {
    if [[ -f "$CONFIG_FILE" ]]; then
        info "Загружена существующая конфигурация из $CONFIG_FILE"
        source "$CONFIG_FILE"
    fi
}

# ── Сохранение конфигурации ────────────────────────────────────────────────
save_config() {
    mkdir -p "$CONFIG_DIR"
    cat > "$CONFIG_FILE" <<EOF
# WireGuard VPS Tunnel — конфигурация
# Создано: $(date '+%Y-%m-%d %H:%M:%S')
MODE="$MODE"
VPS_PUBLIC_IP="$VPS_PUBLIC_IP"
CLIENT_WG_IP="$CLIENT_WG_IP"
SERVER_WG_IP="$SERVER_WG_IP"
WG_PORT="$WG_PORT"
PORTS="$PORTS"
KEEPALIVE="$KEEPALIVE"
EOF
    chmod 600 "$CONFIG_FILE"
    success "Конфигурация сохранена в $CONFIG_FILE"
}

# ── Интерактивный опрос ────────────────────────────────────────────────────
interactive_setup() {
    header "Настройка параметров"

    echo -e "${BOLD}Где запущен этот скрипт?${NC}"
    echo "  1) VPS (сервер с публичным IP)"
    echo "  2) Домашний сервер / Raspberry Pi (за NAT)"
    echo
    while true; do
        read -r -p "Выберите [1/2]: " choice
        case "$choice" in
            1) MODE="vps"; break ;;
            2) MODE="client"; break ;;
            *) warn "Введите 1 или 2" ;;
        esac
    done

    if [[ "$MODE" == "vps" ]]; then
        # Автоопределение публичного IP
        local detected_ip
        detected_ip=$(curl -4 -s ifconfig.me 2>/dev/null || curl -4 -s icanhazip.com 2>/dev/null || echo "")
        if [[ -n "$detected_ip" ]]; then
            echo -e "${BOLD}Публичный IP VPS:${NC} [${detected_ip}]"
            read -r -p "Нажмите Enter для подтверждения или введите другой: " input
            VPS_PUBLIC_IP="${input:-$detected_ip}"
        else
            read -r -p "Введите публичный IP VPS: " VPS_PUBLIC_IP
        fi

        read -r -p "Порты для проброса (через запятую) [80,443]: " input
        PORTS="${input:-80,443}"

        read -r -p "WireGuard IP клиента [10.0.0.2]: " input
        CLIENT_WG_IP="${input:-10.0.0.2}"

        read -r -p "WireGuard порт [51820]: " input
        WG_PORT="${input:-51820}"
    else
        read -r -p "Публичный IP VPS: " VPS_PUBLIC_IP

        read -r -p "WireGuard IP этого сервера [10.0.0.2]: " input
        CLIENT_WG_IP="${input:-10.0.0.2}"

        read -r -p "WireGuard порт VPS [51820]: " input
        WG_PORT="${input:-51820}"

        read -r -p "PersistentKeepalive (сек) [21]: " input
        KEEPALIVE="${input:-21}"
    fi
}

# ── Парсинг аргументов командной строки ────────────────────────────────────
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --vps)
                MODE="vps"
                INTERACTIVE=false
                shift
                ;;
            --client)
                MODE="client"
                INTERACTIVE=false
                shift
                ;;
            --vps-ip)
                VPS_PUBLIC_IP="$2"
                shift 2
                ;;
            --client-wg-ip)
                CLIENT_WG_IP="$2"
                shift 2
                ;;
            --ports)
                PORTS="$2"
                shift 2
                ;;
            --wg-port)
                WG_PORT="$2"
                shift 2
                ;;
            --keepalive)
                KEEPALIVE="$2"
                shift 2
                ;;
            --help|-h)
                echo "Использование: $0 [ОПЦИИ]"
                echo ""
                echo "Опции:"
                echo "  --vps               Режим VPS (сервер с публичным IP)"
                echo "  --client            Режим клиента (домашний сервер за NAT)"
                echo "  --vps-ip IP         Публичный IP-адрес VPS"
                echo "  --client-wg-ip IP   WireGuard IP клиента (по умолчанию: 10.0.0.2)"
                echo "  --ports PORT,...    Порты для проброса (по умолчанию: 80,443)"
                echo "  --wg-port PORT      Порт WireGuard (по умолчанию: 51820)"
                echo "  --keepalive SEC     PersistentKeepalive в секундах (по умолчанию: 21)"
                echo "  --help, -h          Показать эту справку"
                echo ""
                echo "Примеры:"
                echo "  # Интерактивный режим"
                echo "  sudo $0"
                echo ""
                echo "  # Неинтерактивный режим (VPS)"
                echo "  sudo $0 --vps --vps-ip 2.59.219.234 --ports 80,443,8080"
                echo ""
                echo "  # Неинтерактивный режим (клиент)"
                echo "  sudo $0 --client --vps-ip 2.59.219.234 --keepalive 21"
                exit 0
                ;;
            *)
                error "Неизвестный аргумент: $1"
                echo "Используйте --help для справки"
                exit 1
                ;;
        esac
    done
}

# ── Создание конфигурации WireGuard для VPS ────────────────────────────────
configure_vps() {
    header "Настройка WireGuard (VPS — сервер)"

    # Формируем правила iptables для каждого порта
    local dnat_rules=""
    local dnat_down_rules=""
    IFS=',' read -ra PORT_ARRAY <<< "$PORTS"

    for port in "${PORT_ARRAY[@]}"; do
        port=$(echo "$port" | xargs) # trim whitespace
        dnat_rules+="iptables -t nat -A PREROUTING -d ${VPS_PUBLIC_IP} -p tcp --dport ${port} -j DNAT --to-destination ${CLIENT_WG_IP}:${port}; "
        dnat_down_rules+="iptables -t nat -D PREROUTING -d ${VPS_PUBLIC_IP} -p tcp --dport ${port} -j DNAT --to-destination ${CLIENT_WG_IP}:${port}; "
    done

    # Определяем внешний интерфейс
    local ext_if
    ext_if=$(ip route get 8.8.8.8 | awk '{print $5; exit}')

    cat > "$WG_CONF" <<WGCONF
# WireGuard VPS Tunnel — конфигурация сервера (VPS)
# Создано: $(date '+%Y-%m-%d %H:%M:%S')

[Interface]
PrivateKey = ${SERVER_PRIVATE_KEY}
Address = ${SERVER_WG_IP}/24
ListenPort = ${WG_PORT}

# Правила проброса портов и форвардинга
PostUp = ${dnat_rules} iptables -t nat -A POSTROUTING -d ${CLIENT_WG_IP} -j MASQUERADE; iptables -I FORWARD 1 -i ${ext_if} -o wg0 -p tcp -m multiport --dports ${PORTS} -j ACCEPT; iptables -I FORWARD 2 -i wg0 -o ${ext_if} -m state --state ESTABLISHED,RELATED -j ACCEPT; iptables -A FORWARD -i wg0 -j ACCEPT; iptables -A FORWARD -o wg0 -j ACCEPT
PostDown = ${dnat_down_rules} iptables -t nat -D POSTROUTING -d ${CLIENT_WG_IP} -j MASQUERADE; iptables -D FORWARD -i ${ext_if} -o wg0 -p tcp -m multiport --dports ${PORTS} -j ACCEPT; iptables -D FORWARD -i wg0 -o ${ext_if} -m state --state ESTABLISHED,RELATED -j ACCEPT; iptables -D FORWARD -i wg0 -j ACCEPT; iptables -D FORWARD -o wg0 -j ACCEPT

[Peer]
# Клиент (домашний сервер / Raspberry Pi)
PublicKey = ${CLIENT_PUBLIC_KEY}
AllowedIPs = ${CLIENT_WG_IP}/32
WGCONF

    chmod 600 "$WG_CONF"
    success "Конфигурация WireGuard создана: $WG_CONF"
}

# ── Создание конфигурации WireGuard для клиента ────────────────────────────
configure_client() {
    header "Настройка WireGuard (клиент — домашний сервер)"

    cat > "$WG_CONF" <<WGCONF
# WireGuard VPS Tunnel — конфигурация клиента (домашний сервер)
# Создано: $(date '+%Y-%m-%d %H:%M:%S')

[Interface]
PrivateKey = ${CLIENT_PRIVATE_KEY}
Address = ${CLIENT_WG_IP}/24

[Peer]
# VPS (сервер)
PublicKey = ${SERVER_PUBLIC_KEY}
Endpoint = ${VPS_PUBLIC_IP}:${WG_PORT}
AllowedIPs = ${SERVER_WG_IP}/24
PersistentKeepalive = ${KEEPALIVE}
WGCONF

    chmod 600 "$WG_CONF"
    success "Конфигурация WireGuard создана: $WG_CONF"
}

# ── Запуск WireGuard ───────────────────────────────────────────────────────
start_wireguard() {
    header "Запуск WireGuard"

    # Останавливаем если уже запущен (идемпотентность)
    if systemctl is-active --quiet "$WG_SERVICE" 2>/dev/null; then
        info "WireGuard уже запущен, перезапускаем..."
        systemctl stop "$WG_SERVICE"
    fi

    systemctl enable "$WG_SERVICE" &>/dev/null
    systemctl start "$WG_SERVICE"

    sleep 2

    if systemctl is-active --quiet "$WG_SERVICE"; then
        success "WireGuard запущен и добавлен в автозагрузку"
    else
        error "Не удалось запустить WireGuard. Проверьте: systemctl status $WG_SERVICE"
        exit 1
    fi

    # Показываем статус интерфейса
    echo ""
    wg show wg0 2>/dev/null || warn "Интерфейс wg0 не найден"
}

# ── Установка watchdog (только для клиента) ────────────────────────────────
install_watchdog() {
    header "Установка Watchdog (мониторинг туннеля)"

    # Копируем watchdog из директории скрипта
    local script_dir
    script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

    if [[ -f "$script_dir/watchdog.sh" ]]; then
        cp "$script_dir/watchdog.sh" "$WATCHDOG_BIN"
    else
        error "Файл watchdog.sh не найден в $script_dir"
        error "Поместите watchdog.sh рядом с install.sh и запустите снова"
        return 1
    fi

    chmod +x "$WATCHDOG_BIN"

    # Подставляем IP VPS в watchdog
    sed -i "s/VPS_IP=\"\"/VPS_IP=\"${VPS_PUBLIC_IP}\"/" "$WATCHDOG_BIN" 2>/dev/null || true
    sed -i "s/VPS_WG_IP=\"\"/VPS_WG_IP=\"${SERVER_WG_IP}\"/" "$WATCHDOG_BIN" 2>/dev/null || true

    # Создаём systemd-сервис для watchdog
    cat > "$WATCHDOG_SERVICE" <<SYSTEMD
[Unit]
Description=WireGuard Watchdog — оптимизирован для Starlink
After=network-online.target wg-quick@wg0.service
Wants=network-online.target

[Service]
Type=simple
ExecStart=${WATCHDOG_BIN}
Restart=always
RestartSec=10
StandardOutput=journal
StandardError=journal
SyslogIdentifier=wg-watchdog

[Install]
WantedBy=multi-user.target
SYSTEMD

    systemctl daemon-reload
    systemctl enable wg-watchdog &>/dev/null
    systemctl start wg-watchdog

    if systemctl is-active --quiet wg-watchdog; then
        success "Watchdog установлен и запущен"
    else
        warn "Watchdog установлен, но не запустился. Проверьте: systemctl status wg-watchdog"
    fi
}

# ── Вывод итоговой информации ──────────────────────────────────────────────
print_summary() {
    header "Установка завершена"

    echo -e "${BOLD}Режим:${NC}         ${MODE}"
    echo -e "${BOLD}Публичный IP VPS:${NC} ${VPS_PUBLIC_IP}"
    echo -e "${BOLD}WireGuard IP:${NC}    ${CLIENT_WG_IP} (клиент) ↔ ${SERVER_WG_IP} (сервер)"
    echo -e "${BOLD}Порт WG:${NC}         ${WG_PORT}"

    if [[ "$MODE" == "vps" ]]; then
        echo -e "${BOLD}Проброс портов:${NC}  ${PORTS} → ${CLIENT_WG_IP}"
    fi

    echo ""
    echo -e "${BOLD}Полезные команды:${NC}"
    echo "  wg show                          — статус туннеля"
    echo "  systemctl status wg-quick@wg0    — статус сервиса"
    echo "  ping ${SERVER_WG_IP}             — проверка туннеля"

    if [[ "$MODE" == "client" ]]; then
        echo "  systemctl status wg-watchdog     — статус watchdog"
        echo "  tail -f /var/log/wg-watchdog.log — логи watchdog"
    fi

    echo ""
    echo -e "${BOLD}Конфигурация:${NC} ${CONFIG_FILE}"
    echo -e "${BOLD}Ключи:${NC}         ${CONFIG_DIR}/"

    if [[ "$MODE" == "vps" ]]; then
        echo ""
        echo -e "${YELLOW}${BOLD}⚠️  ВАЖНО:${NC} ${YELLOW}Скопируйте публичный ключ сервера на клиент:${NC}"
        echo -e "   ${BOLD}cat ${CONFIG_DIR}/server_public.key${NC}"
        echo ""
        echo -e "${YELLOW}Затем запустите install.sh на клиенте:${NC}"
        echo -e "   ${BOLD}sudo ./install.sh --client --vps-ip ${VPS_PUBLIC_IP}${NC}"
    fi

    echo ""
}

# ── Главная функция ────────────────────────────────────────────────────────
main() {
    banner
    check_root
    parse_args "$@"
    detect_os
    load_config

    if [[ "$INTERACTIVE" == true ]] && [[ -z "$MODE" ]]; then
        interactive_setup
    fi

    if [[ -z "$MODE" ]]; then
        error "Не указан режим. Используйте --vps или --client, либо интерактивный режим."
        exit 1
    fi

    if [[ "$MODE" == "client" ]] && [[ -z "$VPS_PUBLIC_IP" ]]; then
        error "Не указан публичный IP VPS. Используйте --vps-ip IP"
        exit 1
    fi

    # Автоопределение публичного IP для VPS если не указан
    if [[ "$MODE" == "vps" ]] && [[ -z "$VPS_PUBLIC_IP" ]]; then
        VPS_PUBLIC_IP=$(curl -4 -s ifconfig.me 2>/dev/null || curl -4 -s icanhazip.com 2>/dev/null || echo "")
        if [[ -z "$VPS_PUBLIC_IP" ]]; then
            error "Не удалось определить публичный IP. Укажите вручную: --vps-ip IP"
            exit 1
        fi
        info "Автоопределён публичный IP: ${VPS_PUBLIC_IP}"
    fi

    install_deps
    enable_ip_forward
    generate_keys

    if [[ "$MODE" == "vps" ]]; then
        configure_vps
    else
        configure_client
    fi

    start_wireguard
    save_config

    if [[ "$MODE" == "client" ]]; then
        install_watchdog
    fi

    print_summary
}

main "$@"