#!/bin/bash
# ==============================================================================
# WALLARM RECONFIGURATION SCRIPT - V1.0 (Native deployment)
# ==============================================================================
# Purpose: Modify NGINX configuration of an existing native Wallarm node
# Features:
# - Update set_real_ip_from (trusted proxy IPs/CIDRs)
# - Change wallarm_mode (monitoring/block)
# - Backup current config before changes
# - NGINX test + reload instead of container restart
# - Interactive prompts with validation
# ==============================================================================

# Script location and shared library (colors, logging, validation)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=../common/wallarm-lib.sh
source "$SCRIPT_DIR/../common/wallarm-lib.sh"

set -euo pipefail
trap early_error_handler ERR

# ==============================================================================
# CHECK FOR SUDO / ROOT PRIVILEGES
# ==============================================================================
if [ "$EUID" -ne 0 ]; then
    echo -e "${RED}${BOLD}ERROR:${NC} This script must be run with sudo or as root."
    echo -e "${YELLOW}Please run: sudo $0${NC}"
    exit 1
fi

# ==============================================================================
# CONFIGURATION
# ==============================================================================
CONFIG_GLOB="/etc/nginx/conf.d/wallarm-*.conf"
CONFIG_FILE=""
INSTANCE_NAME=""

# ==============================================================================
# FUNCTIONS
# ==============================================================================

# Find the Wallarm NGINX config file(s)
find_wallarm_config() {
    local configs=()
    while IFS= read -r file; do
        configs+=("$file")
    done < <(ls $CONFIG_GLOB 2>/dev/null)

    if [ ${#configs[@]} -eq 0 ]; then
        echo -e "${RED}No Wallarm NGINX configuration found ($CONFIG_GLOB).${NC}"
        echo -e "${YELLOW}Has the native deployment been run? Check /etc/nginx/conf.d/wallarm-*.conf${NC}"
        exit 1
    elif [ ${#configs[@]} -eq 1 ]; then
        CONFIG_FILE="${configs[0]}"
        INSTANCE_NAME=$(basename "$CONFIG_FILE" | sed 's/^wallarm-//; s/\.conf$//')
        echo -e "${GREEN}Found Wallarm config: $CONFIG_FILE${NC}"
    else
        echo -e "${YELLOW}Multiple Wallarm NGINX configs found:${NC}"
        for i in "${!configs[@]}"; do
            echo "$((i+1)). ${configs[$i]}"
        done
        read -r -p "Select config number: " choice
        if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le ${#configs[@]} ]; then
            CONFIG_FILE="${configs[$((choice-1))]}"
            INSTANCE_NAME=$(basename "$CONFIG_FILE" | sed 's/^wallarm-//; s/\.conf$//')
        else
            echo -e "${RED}Invalid selection.${NC}"
            exit 1
        fi
    fi
}

# Parse current configuration to get existing values
parse_current_config() {
    local config_file="$1"
    current_mode=$(grep -oP 'wallarm_mode\s+\K\S+' "$config_file" | head -1)
    current_proxies=$(grep -oP 'set_real_ip_from\s+\K\S+' "$config_file")
}

# Update configuration
update_config() {
    local config_file="$1"
    local backup_file="$config_file.backup.$(date +%Y%m%d_%H%M%S)"

    echo -e "${YELLOW}Backing up current config to $backup_file${NC}"
    cp "$config_file" "$backup_file"

    echo -e "\n${CYAN}${BOLD}Current set_real_ip_from entries:${NC}"
    if [ -n "$current_proxies" ]; then
        while read -r proxy; do
            echo "  $proxy"
        done <<< "$current_proxies"
    else
        echo "  (none)"
    fi

    echo -e "\n${YELLOW}Do you want to change the trusted proxy IPs/CIDRs? (y/N)${NC}"
    read -r change_proxy
    new_proxies=()
    if [[ "$change_proxy" =~ ^[Yy]$ ]]; then
        echo -e "${YELLOW}Enter new trusted proxy IPs/CIDRs (space-separated, or empty to remove all):${NC}"
        read -r new_proxies_input
        if [[ -n "$new_proxies_input" ]]; then
            IFS=' ' read -ra proxy_array <<< "$new_proxies_input"
            for proxy in "${proxy_array[@]}"; do
                proxy=$(echo "$proxy" | xargs)
                if validate_ip_or_cidr "$proxy"; then
                    new_proxies+=("$proxy")
                else
                    echo -e "${RED}Invalid format: $proxy. Skipping.${NC}"
                fi
            done
        fi
    else
        # Keep existing
        while read -r proxy; do
            new_proxies+=("$proxy")
        done <<< "$current_proxies"
    fi

    echo -e "\n${CYAN}${BOLD}Current wallarm_mode:${NC} ${current_mode:-not set}"
    echo -e "${YELLOW}Do you want to change the wallarm_mode? (y/N)${NC}"
    read -r change_mode
    if [[ "$change_mode" =~ ^[Yy]$ ]]; then
        echo -e "${YELLOW}Select new mode:${NC}"
        echo "1. monitoring"
        echo "2. block"
        read -r mode_choice
        case "$mode_choice" in
            1) new_mode="monitoring" ;;
            2) new_mode="block" ;;
            *) echo -e "${RED}Invalid choice, keeping current mode.${NC}"; new_mode="$current_mode" ;;
        esac
    else
        new_mode="$current_mode"
    fi

    # Rebuild the config file
    temp_config=$(mktemp)

    in_server_block=false
    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]]*server[[:space:]]*{ ]]; then
            in_server_block=true
        fi

        if $in_server_block; then
            if [[ "$line" =~ ^[[:space:]]*set_real_ip_from[[:space:]]+ ]]; then
                continue
            fi
            if [[ "$line" =~ ^[[:space:]]*wallarm_mode[[:space:]]+ ]]; then
                continue
            fi
            if [[ "$line" =~ ^[[:space:]]*real_ip_header[[:space:]]+ ]]; then
                continue
            fi
            if [[ "$line" =~ ^[[:space:]]*real_ip_recursive[[:space:]]+ ]]; then
                continue
            fi
        fi

        echo "$line" >> "$temp_config"

        if $in_server_block && [[ "$line" =~ ^[[:space:]]*}$ ]]; then
            in_server_block=false
            if [ ${#new_proxies[@]} -gt 0 ]; then
                for proxy in "${new_proxies[@]}"; do
                    echo "    set_real_ip_from  $proxy;" >> "$temp_config"
                done
                echo "    real_ip_header    X-Real-IP;" >> "$temp_config"
                echo "    real_ip_recursive on;" >> "$temp_config"
            fi
            if [ -n "$new_mode" ]; then
                echo "    wallarm_mode $new_mode;" >> "$temp_config"
            fi
        fi
    done < "$config_file"

    mv "$temp_config" "$config_file"
    chmod 644 "$config_file"

    echo -e "${GREEN}Configuration updated.${NC}"
}

reload_nginx() {
    echo -e "${YELLOW}Testing NGINX configuration...${NC}"
    if ! sudo nginx -t; then
        echo -e "${RED}NGINX configuration test failed. Restoring backup...${NC}"
        cp "$CONFIG_FILE.backup."*. "$CONFIG_FILE" 2>/dev/null || true
        exit 1
    fi

    echo -e "${YELLOW}Reloading NGINX to apply changes...${NC}"
    sudo systemctl reload nginx 2>/dev/null || sudo nginx -s reload 2>/dev/null || {
        sudo systemctl restart nginx 2>/dev/null || {
            echo -e "${RED}Could not reload NGINX. Please do it manually: sudo systemctl reload nginx${NC}"
            exit 1
        }
    }
    echo -e "${GREEN}NGINX reloaded.${NC}"
}

main() {
    echo -e "${BLUE}${BOLD}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║      WALLARM RECONFIGURATION SCRIPT (Native) - V1.0         ║"
    echo "║      Modify NGINX config (trusted proxies / mode)           ║"
    echo "╚══════════════════════════════════════════════════════════════╝${NC}"

    find_wallarm_config

    parse_current_config "$CONFIG_FILE"
    update_config "$CONFIG_FILE"

    echo -e "${YELLOW}Do you want to reload NGINX now? (Y/n)${NC}"
    read -r reload_choice
    if [[ ! "$reload_choice" =~ ^[Nn]$ ]]; then
        reload_nginx
    else
        echo -e "${YELLOW}Changes will take effect after NGINX reload.${NC}"
        echo -e "You can reload later with: sudo systemctl reload nginx"
    fi

    echo -e "\n${GREEN}${BOLD}Reconfiguration completed.${NC}"
}

main "$@"