#!/usr/bin/env bash
# =============================================================================
# WireGuard VPS Tunnel — 1-Click Docker Installer
# Multi-arch: linux/amd64, linux/arm64 (ARM Raspberry Pi)
# =============================================================================
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'

# ── Пути ────────────────────────────────────────────────────────────────────
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONFIG_DIR="$SCRIPT_DIR/config"
REGISTRY="git.softuniq.eu"
IMAGE_PREFIX="opendoor/wireguard-vps-tunnel"
IMAGE_TAG="latest"

# ── Значения по умолчанию ──────────────────────────────────────────────────
MODE=""
VPS_PUBLIC_IP=""
CLIENT_WG_IP="10.0.0.2"
SERVER_WG_IP="10.0.0.1"
WG_PORT="51820"
PORTS="80,443"
KEEPALIVE="21"
INTERACTIVE=true
BUILD_LOCAL=false
SERVER_KEY=""
CLIENT_KEY=""

# ── Функции вывода ──────────────────────────────────────────────────────────
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"; }

# ── Определение архитектуры ────────────────────────────────────────────────
detect_arch() {
    local arch
    arch="$(uname -m)"
    case "$arch" in
        x86_64|amd64)   echo "amd64" ;;
        aarch64|arm64)   echo "arm64" ;;
        armv7l|armhf)   echo "armv7" ;;
        *)               echo "amd64" ;;
    esac
}

# ── Баннер ──────────────────────────────────────────────────────────────────
banner() {
    echo -e "${CYAN}${BOLD}"
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║     WireGuard VPS Tunnel — Docker Installer             ║"
    echo "║     Multi-arch: amd64 + arm64 (RPi / Orange Pi)         ║"
    echo "╚══════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
}

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

# ── Проверка и установка Docker ────────────────────────────────────────────
install_docker() {
    header "Проверка Docker"

    if command -v docker &>/dev/null; then
        success "Docker уже установлен: $(docker --version)"
    else
        info "Установка Docker..."
        curl -fsSL https://get.docker.com | sh
        success "Docker установлен"
    fi

    if docker compose version &>/dev/null; then
        success "Docker Compose доступен: $(docker compose version)"
    else
        error "Docker Compose не найден. Установите Docker Compose plugin."
        exit 1
    fi

    # Проверка multi-arch поддержки
    if docker buildx version &>/dev/null; then
        success "Docker Buildx доступен: $(docker buildx version | head -1)"
    else
        warn "Docker Buildx не найден. Локальная сборка будет только для текущей архитектуры."
    fi
}

# ── Включение IP-форвардинга (обязательно для VPS) ──────────────────────────
enable_ip_forwarding() {
    header "Настройка IP-форвардинга"

    # Включить временно
    echo 1 > /proc/sys/net/ipv4/ip_forward 2>/dev/null || true
    echo 1 > /proc/sys/net/ipv6/conf/all/forwarding 2>/dev/null || true

    # Включить постоянно
    if ! grep -q "net.ipv4.ip_forward" /etc/sysctl.conf 2>/dev/null; then
        echo "net.ipv4.ip_forward = 1" >> /etc/sysctl.conf
    else
        sed -i 's/.*net.ipv4.ip_forward.*/net.ipv4.ip_forward = 1/' /etc/sysctl.conf
    fi

    if ! grep -q "net.ipv6.conf.all.forwarding" /etc/sysctl.conf 2>/dev/null; then
        echo "net.ipv6.conf.all.forwarding = 1" >> /etc/sysctl.conf
    else
        sed -i 's/.*net.ipv6.conf.all.forwarding.*/net.ipv6.conf.all.forwarding = 1/' /etc/sysctl.conf
    fi

    sysctl -p 2>/dev/null || true

    success "IP-форвардинг включён"
}

# ── Остановка host-level WireGuard ────────────────────────────────────────
stop_host_wireguard() {
    if systemctl is-active --quiet wg-quick@wg0 2>/dev/null; then
        warn "Обнаружен host-level WireGuard (wg-quick@wg0). Останавливаем..."
        systemctl stop wg-quick@wg0 2>/dev/null || true
        systemctl disable wg-quick@wg0 2>/dev/null || true
        ip link del wg0 2>/dev/null || true
        success "Host-level WireGuard остановлен и отключён"
    fi

    # Остановить watchdog если есть
    if systemctl is-active --quiet wg-watchdog 2>/dev/null; then
        systemctl stop wg-watchdog 2>/dev/null || true
        systemctl disable wg-watchdog 2>/dev/null || true
        success "wg-watchdog остановлен"
    fi
}

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

    mkdir -p "$CONFIG_DIR"

    # Если ключи переданы через аргументы — используем их
    if [[ -n "$SERVER_KEY" ]] && [[ -n "$CLIENT_KEY" ]]; then
        info "Использование ключей из аргументов командной строки"
        SERVER_PUBLIC_KEY="$SERVER_KEY"
        CLIENT_PUBLIC_KEY="$CLIENT_KEY"
        # Приватные ключи не нужны для .env, только публичные
        SERVER_PRIVATE_KEY=""
        CLIENT_PRIVATE_KEY=""
        success "Ключи загружены из аргументов"
        return
    fi

    if [[ -n "$SERVER_KEY" ]]; then
        info "Использование SERVER_PUBLIC_KEY из аргумента --server-key"
        SERVER_PUBLIC_KEY="$SERVER_KEY"
    fi

    if [[ -n "$CLIENT_KEY" ]]; then
        info "Использование CLIENT_PUBLIC_KEY из аргумента --client-key"
        CLIENT_PUBLIC_KEY="$CLIENT_KEY"
    fi

    if [[ -f "$CONFIG_DIR/server_private.key" ]] && [[ -f "$CONFIG_DIR/client_private.key" ]]; then
        info "Ключи уже существуют, пропускаем генерацию"
    else
        if ! command -v wg &>/dev/null; then
            info "Установка wireguard-tools для генерации ключей..."
            apt-get update -qq 2>/dev/null && apt-get install -y -qq wireguard-tools 2>/dev/null || \
            yum install -y wireguard-tools 2>/dev/null || {
                # Fallback: генерация через Docker
                info "Генерация ключей через Docker (wireguard-tools не установлен на хосте)..."
                docker run --rm alpine:3.20 sh -c "
                    apk add --no-cache wireguard-tools && \
                    wg genkey | tee /tmp/server_private.key | wg pubkey > /tmp/server_public.key && \
                    wg genkey | tee /tmp/client_private.key | wg pubkey > /tmp/client_public.key && \
                    chmod 600 /tmp/server_private.key /tmp/client_private.key
                "
            }
        fi

        if command -v wg &>/dev/null; then
            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"
        fi

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

    # Загружаем ключи из файлов, если не были переданы через аргументы
    if [[ -z "${SERVER_PRIVATE_KEY:-}" ]] && [[ -f "$CONFIG_DIR/server_private.key" ]]; then
        SERVER_PRIVATE_KEY=$(cat "$CONFIG_DIR/server_private.key")
    fi
    if [[ -z "${SERVER_PUBLIC_KEY:-}" ]] && [[ -f "$CONFIG_DIR/server_public.key" ]]; then
        SERVER_PUBLIC_KEY=$(cat "$CONFIG_DIR/server_public.key")
    fi
    if [[ -z "${CLIENT_PRIVATE_KEY:-}" ]] && [[ -f "$CONFIG_DIR/client_private.key" ]]; then
        CLIENT_PRIVATE_KEY=$(cat "$CONFIG_DIR/client_private.key")
    fi
    if [[ -z "${CLIENT_PUBLIC_KEY:-}" ]] && [[ -f "$CONFIG_DIR/client_public.key" ]]; then
        CLIENT_PUBLIC_KEY=$(cat "$CONFIG_DIR/client_public.key")
    fi
}

# ── Создание .env файла ────────────────────────────────────────────────────
create_env_file() {
    if [[ "$MODE" == "vps" ]]; then
        cat > "$SCRIPT_DIR/.env" <<EOF
# WireGuard VPS Tunnel — конфигурация сервера (VPS)
CLIENT_PUBLIC_KEY=${CLIENT_PUBLIC_KEY}
SERVER_PUBLIC_IP=${VPS_PUBLIC_IP:-}
SERVER_PUBLIC_IFACE=
FORWARD_PORTS=${PORTS}
EOF
    else
        cat > "$SCRIPT_DIR/.env" <<EOF
# WireGuard VPS Tunnel — конфигурация клиента
VPS_PUBLIC_IP=${VPS_PUBLIC_IP}
SERVER_PUBLIC_KEY=${SERVER_PUBLIC_KEY}
PERSISTENT_KEEPALIVE=${KEEPALIVE}
EOF
    fi

    chmod 600 "$SCRIPT_DIR/.env"
    success "Файл .env создан"
}

# ── Запуск Docker Compose ──────────────────────────────────────────────────
start_docker() {
    header "Запуск Docker-контейнера"

    local ARCH
    ARCH="$(detect_arch)"
    info "Архитектура: ${ARCH}"

    if [[ "$BUILD_LOCAL" == true ]]; then
        # Локальная сборка для текущей архитектуры
        if [[ "$MODE" == "vps" ]]; then
            info "Локальная сборка серверного контейнера (${ARCH})..."
            docker compose -f "$SCRIPT_DIR/docker-compose.server.yml" up -d --build
        else
            info "Локальная сборка клиентского контейнера (${ARCH})..."
            docker compose -f "$SCRIPT_DIR/docker-compose.client.yml" up -d --build
        fi
    else
        # Использование готового образа из registry
        local image_name
        if [[ "$MODE" == "vps" ]]; then
            image_name="${REGISTRY}/${IMAGE_PREFIX}-server:${IMAGE_TAG}"
        else
            image_name="${REGISTRY}/${IMAGE_PREFIX}-client:${IMAGE_TAG}"
        fi

        info "Загрузка образа: ${image_name} (${ARCH})..."

        # Подтянуть образ для текущей архитектуры
        if docker pull "${image_name}" 2>/dev/null; then
            success "Образ загружен из registry"
        else
            warn "Не удалось загрузить образ из registry. Локальная сборка..."
            BUILD_LOCAL=true
            start_docker
            return
        fi

        # Обновить docker-compose для использования загруженного образа
        if [[ "$MODE" == "vps" ]]; then
            docker compose -f "$SCRIPT_DIR/docker-compose.server.yml" up -d
        else
            docker compose -f "$SCRIPT_DIR/docker-compose.client.yml" up -d
        fi
    fi

    sleep 3

    local container_name
    if [[ "$MODE" == "vps" ]]; then
        container_name="wireguard-server"
    else
        container_name="wireguard-client"
    fi

    if docker ps --format '{{.Names}}' | grep -q "$container_name"; then
        success "Контейнер ${container_name} запущен (${ARCH})"
    else
        error "Контейнер не запустился. Проверьте логи:"
        if [[ "$MODE" == "vps" ]]; then
            echo "  docker compose -f docker-compose.server.yml logs"
        else
            echo "  docker compose -f docker-compose.client.yml logs"
        fi
        exit 1
    fi
}

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

    local arch
    arch="$(detect_arch)"
    echo -e "${BOLD}Текущая архитектура:${NC} ${arch}"

    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
        local detected_ip
        detected_ip=$(curl -4 -s --connect-timeout 5 ifconfig.me 2>/dev/null || curl -4 -s --connect-timeout 5 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 порт [51820]: " input
        WG_PORT="${input:-51820}"
    else
        read -r -p "Публичный IP VPS: " VPS_PUBLIC_IP

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

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

    echo ""
    read -r -p "Собрать образ локально? [y/N]: " build_choice
    if [[ "$build_choice" =~ ^[Yy]$ ]]; then
        BUILD_LOCAL=true
    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 ;;
            --ports)     PORTS="$2"; shift 2 ;;
            --wg-port)   WG_PORT="$2"; shift 2 ;;
            --keepalive) KEEPALIVE="$2"; shift 2 ;;
            --build)     BUILD_LOCAL=true; shift ;;
            --server-key) SERVER_KEY="$2"; shift 2 ;;
            --client-key) CLIENT_KEY="$2"; shift 2 ;;
            --help|-h)
                echo "Использование: $0 [ОПЦИИ]"
                echo ""
                echo "Опции:"
                echo "  --vps               Режим VPS (сервер с публичным IP)"
                echo "  --client            Режим клиента (домашний сервер за NAT)"
                echo "  --vps-ip IP         Публичный IP-адрес VPS"
                echo "  --ports PORT,...    Порты для проброса (по умолчанию: 80,443)"
                echo "  --wg-port PORT      Порт WireGuard (по умолчанию: 51820)"
                echo "  --keepalive SEC     PersistentKeepalive в секундах (по умолчанию: 21)"
                echo "  --build             Собрать образ локально (вместо загрузки из registry)"
                echo "  --server-key KEY    Публичный ключ сервера (для неинтерактивной настройки клиента)"
                echo "  --client-key KEY    Публичный ключ клиента (для неинтерактивной настройки сервера)"
                echo "  --help, -h          Показать эту справку"
                echo ""
                echo "Примеры:"
                echo "  # Интерактивный режим"
                echo "  sudo $0"
                echo ""
                echo "  # Неинтерактивный: VPS (amd64 или arm64)"
                echo "  sudo $0 --vps --vps-ip 2.59.219.234 --ports 80,443"
                echo ""
                echo "  # Неинтерактивный: клиент (Raspberry Pi arm64)"
                echo "  sudo $0 --client --vps-ip 2.59.219.234 --keepalive 21"
                echo ""
                echo "  # Локальная сборка вместо загрузки из registry"
                echo "  sudo $0 --vps --build"
                exit 0
                ;;
            *) error "Неизвестный аргумент: $1"; echo "Используйте --help"; exit 1 ;;
        esac
    done
}

# ── Вывод итоговой информации ──────────────────────────────────────────────
print_summary() {
    local ARCH
    ARCH="$(detect_arch)"

    header "Установка завершена"

    echo -e "${BOLD}Режим:${NC}             ${MODE} (Docker, ${ARCH})"
    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}"
    echo -e "${BOLD}Образ:${NC}               ${BUILD_LOCAL:-pull} из ${REGISTRY}/${IMAGE_PREFIX}"

    if [[ "$MODE" == "vps" ]]; then
        echo -e "${BOLD}Проброс портов:${NC}      ${PORTS} → ${CLIENT_WG_IP}"
        echo ""
        echo -e "${GREEN}${BOLD}✅ IP-форвардинг включён (net.ipv4.ip_forward=1)${NC}"
    fi

    echo ""
    echo -e "${BOLD}Полезные команды:${NC}"
    if [[ "$MODE" == "vps" ]]; then
        echo "  docker compose -f docker-compose.server.yml logs -f  — логи сервера"
        echo "  docker compose -f docker-compose.server.yml ps       — статус"
        echo "  docker exec wireguard-server wg show                 — статус туннеля"
    else
        echo "  docker compose -f docker-compose.client.yml logs -f   — логи клиента"
        echo "  docker compose -f docker-compose.client.yml ps       — статус"
        echo "  docker exec wireguard-client wg show                  — статус туннеля"
    fi

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

    if [[ "$MODE" == "vps" ]]; then
        echo ""
        echo -e "${YELLOW}${BOLD}⚠️  ВАЖНО:${NC} ${YELLOW}Двухэтапный обмен ключами:${NC}"
        echo -e "   1. Запишите публичный ключ сервера:"
        echo -e "      ${BOLD}cat ${CONFIG_DIR}/server_public.key${NC}"
        echo -e "   2. Передайте его на клиент (в .env как SERVER_PUBLIC_KEY)"
        echo -e "   3. После запуска клиента, получите CLIENT_PUBLIC_KEY:"
        echo -e "      ${BOLD}cat ${CONFIG_DIR}/client_public.key${NC} (на клиенте)"
        echo -e "   4. Впишите CLIENT_PUBLIC_KEY в .env на VPS и перезапустите:"
        echo -e "      ${BOLD}docker compose -f docker-compose.server.yml restart${NC}"
    else
        echo ""
        echo -e "${YELLOW}${BOLD}⚠️  ВАЖНО:${NC} ${YELLOW}Скопируйте публичный ключ клиента на VPS:${NC}"
        echo -e "   ${BOLD}cat ${CONFIG_DIR}/client_public.key${NC}"
        echo ""
        echo -e "${YELLOW}И добавьте в .env на VPS:${NC}"
        echo -e "   ${BOLD}CLIENT_PUBLIC_KEY=<ключ>${NC}"
        echo -e "${YELLOW}Затем перезапустите сервер:${NC}"
        echo -e "   ${BOLD}docker compose -f docker-compose.server.yml restart${NC}"
    fi

    echo ""
    echo -e "${CYAN}Поддерживаемые архитектуры: linux/amd64 (Intel/AMD), linux/arm64 (RPi 4/5, Orange Pi)${NC}"
}

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

    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

    if [[ "$MODE" == "vps" ]] && [[ -z "$VPS_PUBLIC_IP" ]]; then
        VPS_PUBLIC_IP=$(curl -4 -s --connect-timeout 5 ifconfig.me 2>/dev/null || curl -4 -s --connect-timeout 5 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_docker
    stop_host_wireguard
    generate_keys
    create_env_file
    if [[ "$MODE" == "vps" ]]; then
        enable_ip_forwarding
    fi
    start_docker
    print_summary
}

main "$@"