#!/usr/bin/bash
# amnezia-himarc-connect.sh - Universal Amnezia HiMarc Client
# (amnezia-himarc-connect.sh - Универсальный клиент Amnezia HiMarc)
# Version: 1.2
# Исправления:
#   - Исправлена команда docker-compose → docker compose / docker-compose
#   - Добавлена поддержка AmneziaWG (awg вместо wg)
#   - Добавлен интерактивный выбор конфига
#   - Улучшена обработка ошибок
#   - Автоматическая установка AmneziaWG компонентов при первом запуске
#   - Поддержка DKMS для автоматической сборки модуля ядра
#   - Автоматическая загрузка модуля ядра
#   - Проверка прав конфигов (600)

set -e

# Colors (Цвета)
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# Directories (Директории)
CONFIG_DIR="/etc/amnezia-himarc"
CLIENTS_DIR="$CONFIG_DIR/clients"
LOG_DIR="/var/log/amnezia-himarc"
LOG_FILE="$LOG_DIR/connect.log"
DOCKER_COMPOSE_DIR="/etc/amnezia-himarc/client"
RPM_DIR="/var/lib/amnezia-himarc/rpms"

# Создаем директории если их нет
sudo mkdir -p "$CLIENTS_DIR" "$LOG_DIR" "$DOCKER_COMPOSE_DIR"

# Log function (Функция логирования)
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | sudo tee -a "$LOG_FILE"
}

# Функция установки AmneziaWG компонентов при первом запуске
install_amneziawg_components() {
    local needs_dkms_build=0
    
    # Устанавливаем tools если нет
    if ! rpm -q amneziawg-tools >/dev/null 2>&1; then
        echo -e "${YELLOW}📦 First start: installing AmneziaWG tools...${NC}"
        if sudo rpm -ivh --force --noscripts "$RPM_DIR/amneziawg-tools.rpm"; then
            echo -e "${GREEN}✅ AmneziaWG tools installed successfully${NC}"
        else
            echo -e "${RED}❌ Failed to install AmneziaWG tools${NC}"
            exit 1
        fi
    fi
    
    # Устанавливаем dkms если нет
    if ! rpm -q amneziawg-dkms >/dev/null 2>&1; then
        echo -e "${YELLOW}📦 First start: installing AmneziaWG DKMS module...${NC}"
        if sudo rpm -ivh --force --noscripts "$RPM_DIR/amneziawg-dkms.rpm"; then
            echo -e "${GREEN}✅ AmneziaWG DKMS package installed${NC}"
            needs_dkms_build=1
        else
            echo -e "${RED}❌ Failed to install AmneziaWG DKMS${NC}"
            exit 1
        fi
    fi
    
    # Собираем модуль через DKMS если нужно
    if [ $needs_dkms_build -eq 1 ] || [ ! -f "/lib/modules/$(uname -r)/extra/amneziawg.ko" ]; then
        echo -e "${BLUE}🔧 Building kernel module with DKMS for kernel $(uname -r)...${NC}"
        if sudo dkms autoinstall; then
            echo -e "${GREEN}✅ Kernel module built successfully${NC}"
        else
            echo -e "${YELLOW}⚠️ DKMS autoinstall failed, trying manual build...${NC}"
            sudo dkms add amneziawg/1.0.20260210 2>/dev/null || true
            sudo dkms build amneziawg/1.0.20260210 -k $(uname -r)
            sudo dkms install amneziawg/1.0.20260210 -k $(uname -r)
        fi
    fi
}

# Функция загрузки модуля ядра AmneziaWG
load_awg_module() {
    echo -e "${BLUE}🔧 Checking AmneziaWG kernel module...${NC}"
    
    # Проверяем, загружен ли модуль
    if lsmod | grep -q amneziawg; then
        echo -e "${GREEN}✅ AmneziaWG module already loaded${NC}"
        return 0
    fi
    
    # Проверяем наличие модуля в системе
    if [ ! -f "/lib/modules/$(uname -r)/extra/amneziawg.ko" ] && \
       [ ! -f "/lib/modules/$(uname -r)/kernel/drivers/net/amneziawg/amneziawg.ko" ]; then
        echo -e "${YELLOW}⚠️ Kernel module not found, rebuilding with DKMS...${NC}"
        sudo dkms autoinstall
    fi
    
    # Пробуем загрузить модуль
    echo -e "${YELLOW}⚠️ Loading AmneziaWG kernel module...${NC}"
    if sudo modprobe amneziawg 2>/dev/null; then
        echo -e "${GREEN}✅ AmneziaWG module loaded successfully${NC}"
        return 0
    else
        echo -e "${RED}❌ Failed to load AmneziaWG kernel module${NC}"
        echo -e "${YELLOW}Please check:${NC}"
        echo -e "  sudo dkms status"
        echo -e "  sudo journalctl -xe | grep amneziawg"
        return 1
    fi
}

# Docker compose command (определяем правильную команду)
if command -v docker-compose &>/dev/null; then
    DOCKER_COMPOSE="docker-compose"
    log "Using docker-compose command"
elif command -v docker &>/dev/null && docker compose version &>/dev/null; then
    DOCKER_COMPOSE="docker compose"
    log "Using docker compose plugin"
else
    echo -e "${RED}❌ Docker Compose not found. Please install docker-compose or docker compose plugin.${NC}"
    echo -e "${YELLOW}Install with: sudo zypper install docker-compose${NC}"
    exit 1
fi

# Проверка наличия Docker
check_docker() {
    if ! command -v docker &>/dev/null; then
        echo -e "${RED}❌ Docker not found. Please install Docker.${NC}"
        echo -e "${YELLOW}Install with: sudo zypper install docker docker-compose${NC}"
        exit 1
    fi
    
    # Проверяем, запущен ли Docker
    if ! systemctl is-active --quiet docker; then
        echo -e "${YELLOW}⚠️ Docker is not running. Starting Docker...${NC}"
        sudo systemctl start docker
        sleep 2
    fi
}

# Docker management (Управление Docker)
docker_up() {
    cd "$DOCKER_COMPOSE_DIR"
    if ! docker ps | grep -q amnezia-client; then
        echo -e "${YELLOW}Starting Docker container... (Запуск Docker контейнера...)${NC}"
        $DOCKER_COMPOSE up -d --build
        sleep 3
        if ! docker ps | grep -q amnezia-client; then
            echo -e "${RED}❌ Failed to start container. Check logs: $0 logs${NC}"
            $DOCKER_COMPOSE logs
            exit 1
        fi
        echo -e "${GREEN}✅ Container started${NC}"
    fi
}

docker_exec() {
    cd "$DOCKER_COMPOSE_DIR"
    docker exec amnezia-client /usr/bin/amnezia "$@"
}

# Функция для показа доступных конфигов
show_available_configs() {
    echo -e "${BLUE}📋 Available AmneziaWG configs:${NC}"
    
    # Проверяем конфиги в директории
    if [ -d "$CLIENTS_DIR" ] && [ "$(ls -A $CLIENTS_DIR/*.conf 2>/dev/null)" ]; then
        for cfg in "$CLIENTS_DIR"/*.conf; do
            name=$(basename "$cfg" .conf)
            perms=$(stat -c %a "$cfg" 2>/dev/null)
            if [ "$perms" = "600" ]; then
                echo -e "   ${GREEN}✓${NC} $name (permissions: $perms)"
            else
                echo -e "   ${YELLOW}⚠${NC} $name (permissions: $perms - should be 600)"
            fi
        done
    else
        echo -e "   ${YELLOW}No configs found. Import one first:${NC}"
        echo -e "   ${BLUE}amnezia-himarc-connect import <name> <file_or_url>${NC}"
        echo -e "   ${BLUE}or use test config: amnezia-himarc-connect import def /usr/share/doc/packages/amnezia-himarc-client/def.conf${NC}"
    fi
}

# Функция для интерактивного выбора конфига
select_config() {
    local configs=($(ls "$CLIENTS_DIR"/*.conf 2>/dev/null | sed 's/.*\///;s/.conf//'))
    
    if [ ${#configs[@]} -eq 0 ]; then
        echo -e "${RED}❌ No configs available. Please import a config first.${NC}"
        echo -e "${YELLOW}Usage: amnezia-himarc-connect import <name> <file_or_url>${NC}"
        echo -e "${YELLOW}Or use test config: amnezia-himarc-connect import def /usr/share/doc/packages/amnezia-himarc-client/def.conf${NC}"
        return 1
    fi
    
    echo -e "${BLUE}Available configs:${NC}"
    for i in "${!configs[@]}"; do
        echo "  $((i+1))) ${configs[$i]}"
    done
    
    read -p "Select config number (or 0 to cancel): " choice
    
    if [ "$choice" = "0" ]; then
        return 1
    fi
    
    if [ "$choice" -ge 1 ] && [ "$choice" -le ${#configs[@]} ]; then
        SELECTED_CONFIG="${configs[$((choice-1))]}"
        return 0
    else
        echo -e "${RED}Invalid selection${NC}"
        return 1
    fi
}

# Проверка AmneziaWG в системе
check_awg() {
    # Устанавливаем компоненты при первом запуске
    install_amneziawg_components
    
    if ! command -v awg &>/dev/null; then
        echo -e "${RED}❌ AmneziaWG tools not found in system.${NC}"
        echo -e "${YELLOW}Please check installation:${NC}"
        echo -e "  ls -la /usr/bin/awg"
        exit 1
    fi
    
    # Загружаем модуль ядра
    load_awg_module
}

# Проверка и загрузка модуля ядра
check_kernel_module() {
    if ! lsmod | grep -q amneziawg; then
        echo -e "${YELLOW}⚠️ AmneziaWG kernel module not loaded${NC}"
        load_awg_module
    fi
}

# Commands (Команды)
case "$1" in
    start)
        check_awg
        check_docker
        docker_up
        echo -e "${GREEN}✅ Client started (Клиент запущен)${NC}"
        log "Client started"
        ;;
    stop)
        cd "$DOCKER_COMPOSE_DIR"
        $DOCKER_COMPOSE down
        echo -e "${GREEN}✅ Client stopped (Клиент остановлен)${NC}"
        log "Client stopped"
        ;;
    restart)
        cd "$DOCKER_COMPOSE_DIR"
        $DOCKER_COMPOSE down
        check_awg
        $DOCKER_COMPOSE up -d --build
        echo -e "${GREEN}✅ Client restarted (Клиент перезапущен)${NC}"
        log "Client restarted"
        ;;
    connect)
        check_awg
        check_docker
        docker_up
        check_kernel_module
        
        # Если имя конфига не указано, показываем список и просим выбрать
        if [ -z "$2" ]; then
            echo -e "${YELLOW}No config specified. Showing available configs:${NC}"
            show_available_configs
            echo ""
            if select_config; then
                CONFIG_NAME="$SELECTED_CONFIG"
            else
                exit 1
            fi
        else
            CONFIG_NAME="$2"
        fi
        
        # Проверяем наличие конфига
        CONFIG_FILE="$CLIENTS_DIR/$CONFIG_NAME.conf"
        if [ ! -f "$CONFIG_FILE" ]; then
            echo -e "${RED}❌ Config $CONFIG_NAME.conf not found in $CLIENTS_DIR${NC}"
            echo -e "${YELLOW}Import it first: amnezia-himarc-connect import $CONFIG_NAME <file_or_url>${NC}"
            exit 1
        fi
        
        # Проверяем права конфига
        PERMS=$(stat -c %a "$CONFIG_FILE")
        if [ "$PERMS" != "600" ]; then
            echo -e "${YELLOW}⚠️ Fixing config permissions (was $PERMS, should be 600)${NC}"
            sudo chmod 600 "$CONFIG_FILE"
        fi
        
        # Копируем конфиг в контейнер
        docker cp "$CONFIG_FILE" amnezia-client:/config/
        
        echo -e "${BLUE}🔌 Connecting to $CONFIG_NAME using AmneziaWG...${NC}"
        if docker_exec connect "$CONFIG_NAME"; then
            log "Connected to $CONFIG_NAME"
            echo -e "${GREEN}✅ Successfully connected to $CONFIG_NAME${NC}"
            echo ""
            docker_exec status
        else
            echo -e "${RED}❌ Failed to connect to $CONFIG_NAME${NC}"
            log "Failed to connect to $CONFIG_NAME"
            exit 1
        fi
        ;;
    disconnect)
        check_awg
        check_docker
        docker_up
        echo -e "${BLUE}🔌 Disconnecting...${NC}"
        if docker_exec disconnect; then
            log "Disconnected"
            echo -e "${GREEN}✅ Disconnected${NC}"
        else
            echo -e "${YELLOW}⚠️ Not connected or disconnect failed${NC}"
        fi
        ;;
    status)
        check_awg
        check_docker
        docker_up
        echo -e "${BLUE}📊 AmneziaWG Connection Status:${NC}"
        echo "==================================="
        docker_exec status
        ;;
    list)
        check_awg
        check_docker
        docker_up
        echo -e "${BLUE}📋 Available AmneziaWG configs:${NC}"
        echo "==================================="
        show_available_configs
        ;;
    import)
        if [ -z "$2" ]; then
            echo -e "${RED}❌ Error: config name required${NC}"
            echo -e "${YELLOW}Usage: $0 import <name> <file_or_url>${NC}"
            echo ""
            echo "Examples:"
            echo "  $0 import myphone ./myphone.conf"
            echo "  $0 import myserver https://example.com/config.conf"
            echo "  $0 import def /usr/share/doc/packages/amnezia-himarc-client/def.conf"
            exit 1
        fi
        
        if [ -z "$3" ]; then
            echo -e "${RED}❌ Error: source file or URL required${NC}"
            echo -e "${YELLOW}Usage: $0 import <name> <file_or_url>${NC}"
            exit 1
        fi
        
        check_awg
        check_docker
        docker_up
        
        TEMP_FILE="/tmp/$2.conf"
        
        # Если это локальный файл, копируем
        if [ -f "$3" ]; then
            echo -e "${BLUE}📁 Importing from file: $3${NC}"
            cp "$3" "$TEMP_FILE"
        elif [[ "$3" =~ ^https?:// ]]; then
            echo -e "${BLUE}🌐 Importing from URL: $3${NC}"
            curl -L -s -o "$TEMP_FILE" "$3"
            if [ $? -ne 0 ]; then
                echo -e "${RED}❌ Download failed${NC}"
                exit 1
            fi
        else
            echo -e "${RED}❌ Source not found: $3${NC}"
            exit 1
        fi
        
        # Устанавливаем правильные права (600)
        chmod 600 "$TEMP_FILE"
        
        # Копируем в директорию конфигов
        sudo mv "$TEMP_FILE" "$CLIENTS_DIR/$2.conf"
        
        # Копируем в контейнер
        docker cp "$CLIENTS_DIR/$2.conf" amnezia-client:/config/
        
        echo -e "${GREEN}✅ Imported $2 with secure permissions (600)${NC}"
        log "Imported $2"
        
        # Показываем список после импорта
        echo ""
        show_available_configs
        ;;
    logs)
        cd "$DOCKER_COMPOSE_DIR"
        docker logs amnezia-client -f
        ;;
    shell)
        check_awg
        check_docker
        docker_up
        echo -e "${BLUE}🐚 Entering AmneziaWG container shell...${NC}"
        echo -e "${YELLOW}Type 'exit' to return${NC}"
        echo ""
        docker exec -it amnezia-client /bin/bash
        ;;
    gui)
        # Проверяем наличие dialog
        if command -v dialog &>/dev/null; then
            exec "$0" menu
        else
            # Запускаем в новом терминале
            if command -v gnome-terminal &>/dev/null; then
                gnome-terminal -- "$0" menu
            elif command -v xterm &>/dev/null; then
                xterm -e "$0" menu
            else
                "$0" menu
            fi
        fi
        ;;
    menu)
        while true; do
            clear
            echo -e "${BLUE}========================================${NC}"
            echo -e "${GREEN}     Amnezia HiMarc Client Menu${NC}"
            echo -e "${BLUE}========================================${NC}"
            echo ""
            
            # Проверяем статус Docker
            if ! docker ps | grep -q amnezia-client; then
                echo -e "  ${RED}●${NC} Docker: ${RED}STOPPED${NC}"
            else
                echo -e "  ${GREEN}●${NC} Docker: ${GREEN}RUNNING${NC}"
                # Показываем статус подключения
                if docker exec amnezia-client awg show 2>/dev/null | grep -q "interface"; then
                    echo -e "  ${GREEN}●${NC} VPN: ${GREEN}CONNECTED${NC}"
                else
                    echo -e "  ${RED}○${NC} VPN: ${RED}DISCONNECTED${NC}"
                fi
            fi
            echo ""
            
            echo -e "  ${YELLOW}1)${NC} 🔌 Connect (Подключиться)"
            echo -e "  ${YELLOW}2)${NC} 🔌 Disconnect (Отключиться)"
            echo -e "  ${YELLOW}3)${NC} 📊 Status (Статус)"
            echo -e "  ${YELLOW}4)${NC} 📋 List configs (Список конфигов)"
            echo -e "  ${YELLOW}5)${NC} 📁 Import config (Импорт конфига)"
            echo -e "  ${YELLOW}6)${NC} 📜 View logs (Логи)"
            echo -e "  ${YELLOW}7)${NC} 🐚 Shell (Войти в контейнер)"
            echo -e "  ${YELLOW}8)${NC} 🚀 Start client (Запустить клиент)"
            echo -e "  ${YELLOW}9)${NC} ⏹️ Stop client (Остановить клиент)"
            echo -e "  ${YELLOW}0)${NC} ❌ Exit (Выход)"
            echo ""
            echo -e "${BLUE}========================================${NC}"
            read -p "Choice (Выбор): " choice
            
            case $choice in
                1)
                    clear
                    show_available_configs
                    echo ""
                    read -p "Config name (or press Enter to select from list): " c
                    if [ -z "$c" ]; then
                        if select_config; then
                            "$0" connect "$SELECTED_CONFIG"
                        fi
                    else
                        [ -n "$c" ] && "$0" connect "$c"
                    fi
                    read -p "Press Enter..."
                    ;;
                2)
                    "$0" disconnect
                    read -p "Press Enter..."
                    ;;
                3)
                    clear
                    "$0" status
                    read -p "Press Enter..."
                    ;;
                4)
                    clear
                    "$0" list
                    read -p "Press Enter..."
                    ;;
                5)
                    clear
                    echo -e "${BLUE}📁 Import AmneziaWG Config${NC}"
                    echo "================================"
                    read -p "Config name (e.g., myphone): " n
                    if [ -n "$n" ]; then
                        read -p "File path or URL: " f
                        if [ -n "$f" ]; then
                            "$0" import "$n" "$f"
                        fi
                    fi
                    read -p "Press Enter..."
                    ;;
                6)
                    "$0" logs
                    ;;
                7)
                    "$0" shell
                    ;;
                8)
                    "$0" start
                    read -p "Press Enter..."
                    ;;
                9)
                    "$0" stop
                    read -p "Press Enter..."
                    ;;
                0)
                    exit 0
                    ;;
            esac
        done
        ;;
    *)
        echo "Amnezia HiMarc Client (AmneziaWG) v1.2"
        echo "========================================"
        echo ""
        echo "Commands (Команды):"
        echo "  $0 start          - Start client (Запустить клиент)"
        echo "  $0 stop           - Stop client (Остановить клиент)"
        echo "  $0 restart        - Restart client (Перезапустить)"
        echo "  $0 connect [name] - Connect to VPN (Подключиться)"
        echo "  $0 disconnect     - Disconnect (Отключиться)"
        echo "  $0 status         - Show status (Статус)"
        echo "  $0 list           - List configs (Список конфигов)"
        echo "  $0 import <n> <f> - Import config (Импорт)"
        echo "  $0 logs           - View logs (Логи)"
        echo "  $0 shell          - Enter container (Войти)"
        echo "  $0 gui            - GUI mode (Графический режим)"
        echo "  $0 menu           - Menu mode (Меню)"
        echo ""
        echo "Test config:"
        echo "  $0 import def /usr/share/doc/packages/amnezia-himarc-client/def.conf"
        echo ""
        echo "Examples (Примеры):"
        echo "  $0 import myphone ./myphone.conf"
        echo "  $0 import def /usr/share/doc/packages/amnezia-himarc-client/def.conf"
        echo "  $0 connect        # будет предложен выбор"
        echo "  $0 connect def    # подключиться к тестовому конфигу"
        echo ""
        
        # Если есть контейнер, показываем доступные конфиги
        if docker ps 2>/dev/null | grep -q amnezia-client; then
            show_available_configs
        fi
        ;;
esac