#!/bin/bash
# ==============================================================================
# WALLARM UNINSTALL SCRIPT - V1.0
# ==============================================================================
# Purpose: Safely remove Wallarm filtering node and cleanup Docker installation
# Features:
# - Interactive confirmation with safety checks
# - Stops and removes Wallarm container and image
# - Removes Docker service files created by deployment script
# - Optional cleanup of Docker binaries (if no other containers exist)
# - Preserves user data and logs (with option to remove)
# - DAU-friendly warnings and confirmations
# ==============================================================================

# Color definitions for better UX
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[1;34m'
CYAN='\033[0;36m'
MAGENTA='\033[0;35m'
BOLD='\033[1m'
NC='\033[0m' # No Color

# Strict error handling
set -euo pipefail
# Simple error handler for early failures (before log_message is defined)
early_error_handler() {
    echo -e "${RED}${BOLD}[ERROR]${NC} Script failed at line $LINENO. Command: $BASH_COMMAND" >&2
    exit 1
}
trap early_error_handler ERR

# Logging function
log_message() {
    local level="$1"
    local message="$2"
    local timestamp
    timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    case "$level" in
        "INFO")    color="${BLUE}" ;;
        "SUCCESS") color="${GREEN}" ;;
        "WARNING") color="${YELLOW}" ;;
        "ERROR")   color="${RED}" ;;
        "DEBUG")   color="${CYAN}" ;;
        *)         color="${NC}" ;;
    esac
    
    echo -e "${color}[${timestamp}] ${level}: ${message}${NC}" >&2
}

# Ask for confirmation
confirm() {
    local prompt="$1"
    local default="${2:-n}"
    local options="[y/N]"
    
    if [ "$default" = "y" ]; then
        options="[Y/n]"
    fi
    
    echo -e -n "${YELLOW}${prompt} ${options}${NC} "
    read -r response
    
    case "$response" in
        [yY][eE][sS]|[yY])
            return 0
            ;;
        [nN][oO]|[nN])
            return 1
            ;;
        "")
            # Use default
            if [ "$default" = "y" ]; then
                return 0
            else
                return 1
            fi
            ;;
        *)
            # Invalid input, treat as no
            return 1
            ;;
    esac
}

# Check if running as root or with sudo
check_sudo() {
    if [ "$EUID" -ne 0 ]; then
        log_message "INFO" "This script requires sudo privileges"
        if ! sudo -n true 2>/dev/null; then
            log_message "INFO" "Please enter your sudo password when prompted"
            sudo -v
        fi
    fi
}

# Detect init system
detect_init_system() {
    if command -v systemctl >/dev/null 2>&1 && systemctl --version >/dev/null 2>&1; then
        echo "systemd"
    elif [ -d /run/openrc ]; then
        echo "openrc"
    elif [ -f /etc/init.d/docker ]; then
        echo "sysvinit"
    else
        echo "unknown"
    fi
}

# Check if Docker is installed and running
check_docker() {
    if ! command -v docker >/dev/null 2>&1; then
        log_message "WARNING" "Docker command not found"
        return 1
    fi
    
    if ! sudo docker info >/dev/null 2>&1; then
        log_message "WARNING" "Docker is not running"
        return 1
    fi
    
    return 0
}

# Check for other Docker containers (besides Wallarm)
check_other_containers() {
    local wallarm_container="wallarm-node"
    local all_containers
    all_containers=$(sudo docker ps -a -q 2>/dev/null | wc -l)
    local wallarm_containers
    wallarm_containers=$(sudo docker ps -a --filter "name=${wallarm_container}" -q 2>/dev/null | wc -l)
    
    if [ "$all_containers" -gt "$wallarm_containers" ]; then
        log_message "WARNING" "Found other Docker containers besides Wallarm"
        sudo docker ps -a --format "table {{.Names}}\t{{.Image}}\t{{.Status}}" | grep -v "$wallarm_container" || true
        return 0  # Other containers exist
    fi
    
    return 1  # Only Wallarm containers or no containers
}

# Stop and remove Wallarm container
remove_wallarm_container() {
    local container_name="wallarm-node"
    
    log_message "INFO" "Looking for Wallarm container..."
    
    if sudo docker ps -a --filter "name=${container_name}" --format "{{.Names}}" | grep -q "${container_name}"; then
        log_message "INFO" "Found Wallarm container: ${container_name}"
        
        # Stop container if running
        if sudo docker ps --filter "name=${container_name}" --filter "status=running" --format "{{.Names}}" | grep -q "${container_name}"; then
            log_message "INFO" "Stopping Wallarm container..."
            sudo docker stop "${container_name}" || {
                log_message "WARNING" "Failed to stop container, attempting force stop"
                sudo docker kill "${container_name}" 2>/dev/null || true
            }
        fi
        
        # Remove container
        log_message "INFO" "Removing Wallarm container..."
        sudo docker rm -f "${container_name}" 2>/dev/null || {
            log_message "WARNING" "Failed to remove container, it may already be removed"
        }
        
        log_message "SUCCESS" "Wallarm container removed"
    else
        log_message "INFO" "No Wallarm container found"
    fi
}

# Remove Wallarm image
remove_wallarm_image() {
    local image_name="wallarm/node"
    
    log_message "INFO" "Looking for Wallarm image..."
    
    if sudo docker images --format "{{.Repository}}" | grep -q "^${image_name}"; then
        log_message "INFO" "Found Wallarm image: ${image_name}"
        
        # Check if image is used by any containers
        local used_by
        used_by=$(sudo docker ps -a --filter "ancestor=${image_name}" -q 2>/dev/null | wc -l)
        
        if [ "$used_by" -gt 0 ]; then
            log_message "WARNING" "Image ${image_name} is still in use by containers, skipping removal"
            return
        fi
        
        # Remove image
        log_message "INFO" "Removing Wallarm image..."
        sudo docker rmi "${image_name}:latest" 2>/dev/null || {
            log_message "WARNING" "Failed to remove image, it may be in use or already removed"
        }
        
        # Also try to remove by ID if tag removal failed
        local image_id
        image_id=$(sudo docker images --filter "reference=${image_name}" --format "{{.ID}}" 2>/dev/null | head -1)
        if [ -n "$image_id" ]; then
            sudo docker rmi -f "$image_id" 2>/dev/null || true
        fi
        
        log_message "SUCCESS" "Wallarm image removed"
    else
        log_message "INFO" "No Wallarm image found"
    fi
}

# Remove Docker service files (created by deployment script)
remove_docker_service_files() {
    local init_system
    init_system=$(detect_init_system)
    
    log_message "INFO" "Removing Docker service files for init system: ${init_system}"
    
    case "$init_system" in
        "systemd")
            # Stop and disable Docker service
            if sudo systemctl is-active docker --quiet 2>/dev/null; then
                log_message "INFO" "Stopping Docker service..."
                sudo systemctl stop docker 2>/dev/null || true
            fi
            
            if sudo systemctl is-enabled docker --quiet 2>/dev/null; then
                log_message "INFO" "Disabling Docker service..."
                sudo systemctl disable docker 2>/dev/null || true
            fi
            
            # Remove systemd unit files (if they exist and were created by our script)
            local systemd_files=(
                "/etc/systemd/system/docker.socket"
                "/etc/systemd/system/docker.service"
                "/usr/lib/systemd/system/docker.socket"
                "/usr/lib/systemd/system/docker.service"
            )
            
            for file in "${systemd_files[@]}"; do
                if [ -f "$file" ]; then
                    log_message "INFO" "Removing systemd file: $file"
                    sudo rm -f "$file"
                fi
            done
            
            sudo systemctl daemon-reload 2>/dev/null || true
            ;;
            
        "openrc")
            # Stop and remove from runlevels
            if sudo rc-service docker status 2>/dev/null | grep -q "started"; then
                log_message "INFO" "Stopping Docker service (OpenRC)..."
                sudo rc-service docker stop 2>/dev/null || true
            fi
            
            if [ -f /etc/init.d/docker ]; then
                log_message "INFO" "Removing OpenRC init script..."
                sudo rc-update del docker default 2>/dev/null || true
                sudo rm -f /etc/init.d/docker
            fi
            ;;
            
        "sysvinit")
            # Stop service
            if [ -f /etc/init.d/docker ]; then
                log_message "INFO" "Stopping Docker service (SysV init)..."
                sudo service docker stop 2>/dev/null || true
                
                # Remove from startup
                if command -v update-rc.d >/dev/null 2>&1; then
                    sudo update-rc.d -f docker remove 2>/dev/null || true
                elif command -v chkconfig >/dev/null 2>&1; then
                    sudo chkconfig --del docker 2>/dev/null || true
                fi
                
                log_message "INFO" "Removing SysV init script..."
                sudo rm -f /etc/init.d/docker
            fi
            ;;
            
        *)
            log_message "WARNING" "Unknown init system, skipping service file cleanup"
            ;;
    esac
    
    log_message "SUCCESS" "Docker service files removed"
}

# Remove Docker binaries (optional, only if no other containers exist)
remove_docker_binaries() {
    local docker_binaries=(
        "/usr/bin/docker"
        "/usr/bin/dockerd"
        "/usr/bin/docker-init"
        "/usr/bin/docker-proxy"
        "/usr/bin/containerd"
        "/usr/bin/containerd-shim"
        "/usr/bin/containerd-shim-runc-v1"
        "/usr/bin/containerd-shim-runc-v2"
        "/usr/bin/runc"
    )
    
    log_message "INFO" "Checking Docker binaries..."
    
    local binaries_found=0
    for binary in "${docker_binaries[@]}"; do
        if [ -f "$binary" ]; then
            binaries_found=$((binaries_found + 1))
        fi
    done
    
    if [ "$binaries_found" -eq 0 ]; then
        log_message "INFO" "No Docker binaries found in /usr/bin/"
        return
    fi
    
    if confirm "Remove Docker binaries from /usr/bin/? (Only do this if Docker was installed by wallarm-ct-deploy.sh)" "n"; then
        log_message "WARNING" "Removing Docker binaries..."
        
        for binary in "${docker_binaries[@]}"; do
            if [ -f "$binary" ]; then
                log_message "INFO" "Removing $binary"
                sudo rm -f "$binary"
            fi
        done
        
        # Also remove CNI plugins if they exist
        if [ -d "/opt/cni/bin" ]; then
            log_message "INFO" "Removing CNI plugins from /opt/cni/bin/"
            sudo rm -rf /opt/cni/bin/*
        fi
        
        log_message "SUCCESS" "Docker binaries removed"
    else
        log_message "INFO" "Skipping Docker binary removal"
    fi
}

# Remove Docker configuration files
remove_docker_config() {
    local config_files=(
        "/etc/docker/daemon.json"
        "/etc/containerd/config.toml"
        "/var/lib/docker"  # Warning: This removes all Docker data!
    )
    
    log_message "INFO" "Checking Docker configuration files..."
    
    # Only remove daemon.json if it was created by our script
    if [ -f "/etc/docker/daemon.json" ]; then
        log_message "INFO" "Found /etc/docker/daemon.json"
        if grep -q "storage-driver.*vfs" "/etc/docker/daemon.json" 2>/dev/null; then
            log_message "INFO" "This appears to be the VFS configuration from wallarm-ct-deploy.sh"
            if confirm "Remove /etc/docker/daemon.json?" "n"; then
                sudo rm -f "/etc/docker/daemon.json"
                log_message "SUCCESS" "Docker configuration removed"
            fi
        else
            log_message "WARNING" "/etc/docker/daemon.json doesn't appear to be from wallarm-ct-deploy.sh, skipping"
        fi
    fi
    
    # Warn about Docker data directory
    if [ -d "/var/lib/docker" ]; then
        log_message "WARNING" "/var/lib/docker contains Docker data (images, containers, volumes)"
        log_message "WARNING" "Removing this directory will delete ALL Docker data on the system"
        if confirm "Remove /var/lib/docker? (WARNING: Deletes ALL Docker data)" "n"; then
            log_message "WARNING" "Removing /var/lib/docker - this may take a while..."
            sudo rm -rf /var/lib/docker
            log_message "SUCCESS" "Docker data directory removed"
        fi
    fi
}

# Remove docker group (if empty)
remove_docker_group() {
    log_message "INFO" "Checking docker group..."
    
    if getent group docker >/dev/null; then
        local group_users
        group_users=$(getent group docker | cut -d: -f4)
        
        if [ -z "$group_users" ]; then
            log_message "INFO" "Docker group exists and has no users"
            if confirm "Remove docker group?" "n"; then
                sudo groupdel docker 2>/dev/null || {
                    log_message "WARNING" "Failed to remove docker group (may be system group)"
                }
                log_message "SUCCESS" "Docker group removed"
            fi
        else
            log_message "WARNING" "Docker group has users: $group_users"
            log_message "INFO" "Skipping docker group removal (users still present)"
        fi
    else
        log_message "INFO" "Docker group not found"
    fi
}

# Remove Wallarm-specific files and logs
remove_wallarm_files() {
    local wallarm_files=(
        "$HOME/wallarm-start.sh"
        "$HOME/wallarm-stop.sh"
        "$HOME/wallarm-status.sh"
        "/usr/local/bin/wallarm-start"
        "/usr/local/bin/wallarm-stop"
        "/usr/local/bin/wallarm-status"
    )
    
    log_message "INFO" "Removing Wallarm scripts and logs..."
    
    # Remove scripts
    for file in "${wallarm_files[@]}"; do
        if [ -f "$file" ]; then
            log_message "INFO" "Removing $file"
            sudo rm -f "$file"
        fi
    done
    
    # Remove log directory (if empty)
    local log_dir="$HOME/logs"
    if [ -d "$log_dir" ]; then
        log_message "INFO" "Found log directory: $log_dir"
        if [ -z "$(ls -A "$log_dir" 2>/dev/null)" ]; then
            log_message "INFO" "Log directory is empty, removing..."
            sudo rmdir "$log_dir" 2>/dev/null || true
        else
            log_message "INFO" "Log directory contains files, preserving..."
        fi
    fi
    
    # Remove .env file if it exists
    if [ -f ".env" ]; then
        log_message "INFO" "Removing .env file..."
        rm -f ".env"
    fi
    
    log_message "SUCCESS" "Wallarm files cleaned up"
}

# Main uninstall function
main() {
    echo -e "${CYAN}${BOLD}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║                  WALLARM UNINSTALLATION                      ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
    
    echo -e "${YELLOW}This script will remove Wallarm filtering node and cleanup Docker installation.${NC}"
    echo -e "${YELLOW}You will be asked for confirmation before each destructive operation.${NC}"
    echo ""
    
    if ! confirm "Do you want to continue with the uninstallation?" "n"; then
        log_message "INFO" "Uninstallation cancelled by user"
        exit 0
    fi
    
    # Check sudo
    check_sudo
    
    # Check Docker
    if check_docker; then
        log_message "INFO" "Docker is installed and running"
        
        # Check for other containers
        if check_other_containers; then
            log_message "WARNING" "Other Docker containers exist on this system"
            echo -e "${YELLOW}Warning: Removing Docker may affect other containers.${NC}"
            echo -e "${YELLOW}Consider leaving Docker installed if you need it for other purposes.${NC}"
            echo ""
        fi
    else
        log_message "WARNING" "Docker is not running or not installed"
    fi
    
    # Step 1: Remove Wallarm container and image
    echo ""
    echo -e "${CYAN}${BOLD}Step 1: Remove Wallarm container and image${NC}"
    if confirm "Stop and remove Wallarm container and image?" "y"; then
        remove_wallarm_container
        remove_wallarm_image
    else
        log_message "INFO" "Skipping Wallarm container/image removal"
    fi
    
    # Step 2: Remove Docker service files
    echo ""
    echo -e "${CYAN}${BOLD}Step 2: Remove Docker service files${NC}"
    if confirm "Remove Docker service files (systemd/OpenRC/SysV init scripts)?" "y"; then
        remove_docker_service_files
    else
        log_message "INFO" "Skipping Docker service file removal"
    fi
    
    # Step 3: Optional Docker binary removal
    echo ""
    echo -e "${CYAN}${BOLD}Step 3: Docker binaries and configuration${NC}"
    remove_docker_binaries
    remove_docker_config
    
    # Step 4: Remove docker group
    echo ""
    echo -e "${CYAN}${BOLD}Step 4: System cleanup${NC}"
    remove_docker_group
    
    # Step 5: Remove Wallarm files
    echo ""
    echo -e "${CYAN}${BOLD}Step 5: Wallarm files and logs${NC}"
    if confirm "Remove Wallarm scripts and log files?" "y"; then
        remove_wallarm_files
    else
        log_message "INFO" "Skipping Wallarm file cleanup"
    fi
    
    # Final message
    echo ""
    echo -e "${GREEN}${BOLD}╔══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${GREEN}${BOLD}║               UNINSTALLATION COMPLETE                        ║${NC}"
    echo -e "${GREEN}${BOLD}╚══════════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo -e "${GREEN}Wallarm filtering node has been removed.${NC}"
    echo ""
    echo -e "${YELLOW}Note:${NC}"
    echo -e "  • Docker may still be installed on your system"
    echo -e "  • Docker data in /var/lib/docker may still exist"
    echo -e "  • User may still be in docker group (check with 'groups')"
    echo ""
    echo -e "To completely remove Docker, you may need to:"
    echo -e "  1. Remove Docker package using your system's package manager"
    echo -e "  2. Remove /var/lib/docker directory (contains all Docker data)"
    echo -e "  3. Remove user from docker group: sudo gpasswd -d \$USER docker"
    echo ""
}

# Run main function
main "$@"