#!/bin/bash

#-----------------------------------#
#             VARIABLES             #
#-----------------------------------#

this_script_url="https://git.hhf.technology/hhf/script-management-cloudpanel/raw/branch/main/maintenance/crowdsec_removal.sh"
this_script_name="CrowdSec Removal Script"
formatter_url="https://git.hhf.technology/hhf/TaskFormatter/raw/branch/main/bash_task_formatter/task_formatter.sh"
scriptname=$0

# Script version
VERSION="2.0.0"

# Initialize success flag
success=0

# Determine the user (use the first argument if provided, otherwise fallback)
USER_TO_RUN_AS="${1:-$SUDO_USER}"
USER_HOME=$(eval echo ~$USER_TO_RUN_AS)

# Logging setup
LOGFILE="/var/log/crowdsec_removal_$(date +%Y%m%d_%H%M%S).log"
BACKUP_DIR="/var/backup/crowdsec_$(date +%Y%m%d_%H%M%S)"

#-----------------------------------#
#             FORMATTER             #
#-----------------------------------#

# Download and source the formatter with error handling
download_formatter() {
    if [ ! -f "task_formatter.sh" ]; then
        if ! wget "$formatter_url" --no-check-certificate -O task_formatter.sh > /dev/null 2>&1; then
            echo "Error: Failed to download task_formatter.sh"
            exit 1
        fi
    fi

    if [ ! -f "task_formatter.sh" ]; then
        echo "Error: task_formatter.sh not found after download attempt"
        exit 1
    fi

    # Make the formatter executable
    chmod +x task_formatter.sh

    # Source the formatter with error checking
    source ./task_formatter.sh || { echo "Error: Failed to source task_formatter.sh"; exit 1; }
    
    # Check if print_header is available after sourcing
    if ! declare -f print_header > /dev/null; then
        echo "Error: print_header function not found after sourcing."
        exit 1
    fi
}

# Call the download_formatter function
download_formatter

#-----------------------------------#
#             FUNCTIONS             #
#-----------------------------------#

# Trap errors and interrupts
trap 'error_handler $? $LINENO' ERR
trap 'cleanup' EXIT
trap 'interrupt_handler' INT TERM

# Function to handle errors
error_handler() {
    local exit_code=$1
    local line_number=$2
    log_message "ERROR" "Error $exit_code occurred on line $line_number"
    cleanup
    exit "$exit_code"
}

# Function to handle interrupts
interrupt_handler() {
    log_message "WARNING" "Script interrupted by user"
    cleanup
    exit 130
}

# Function to clean up temporary files
cleanup_files() {
    # Remove temporary files if they exist
    [ -f "/tmp/crowdsec_services.tmp" ] && rm -f "/tmp/crowdsec_services.tmp"
    [ -f "/tmp/crowdsec_packages.tmp" ] && rm -f "/tmp/crowdsec_packages.tmp"
    return 0
}

# Function to remove the script
remove_script() {
    if [ -f "$0" ]; then
        rm -f -- "$0"
    fi
    if [ -f "task_formatter.sh" ]; then
        rm -f task_formatter.sh
    fi
    return 0
}

# Main cleanup function
cleanup() {
    format_output cleanup_files "Cleaning up temporary files"
    format_output remove_script "Removing script"
}

# Function to check if a command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Enhanced logging function with formatter integration
log_message() {
    local level=$1
    local message=$2
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] [$level] $message" >> "$LOGFILE"
    case $level in
        "INFO")    log "$message" ;;
        "WARNING") log "$message" ;;
        "ERROR")   log "$message" ;;
        "SUCCESS") log "${CHECK_MARK} $message" ;;
    esac
}

# Function to check system requirements
check_system_requirements() {
    local missing_deps=()
    
    # Check for required commands
    for cmd in systemctl grep awk sed; do
        if ! command_exists "$cmd"; then
            missing_deps+=("$cmd")
        fi
    done
    
    if [ ${#missing_deps[@]} -ne 0 ]; then
        log_message "ERROR" "Missing required dependencies: ${missing_deps[*]}"
        return 1
    fi
    return 0
}

# Function to check if running as root
check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        log_message "ERROR" "Please run as root"
        return 1
    fi
    return 0
}

# Function to create backup
create_backup() {
    log_message "INFO" "Creating backup of CrowdSec configurations..."
    
    mkdir -p "$BACKUP_DIR"
    
    # Backup configuration files
    if [ -d "/etc/crowdsec" ]; then
        cp -r "/etc/crowdsec" "$BACKUP_DIR/"
    fi
    
    if [ -d "/var/lib/crowdsec" ]; then
        cp -r "/var/lib/crowdsec" "$BACKUP_DIR/"
    fi
    
    # Backup service files
    if [ -d "/etc/systemd/system" ]; then
        find "/etc/systemd/system" -name "*crowdsec*" -exec cp {} "$BACKUP_DIR/" \;
    fi
    
    log_message "SUCCESS" "Backup created at $BACKUP_DIR"
    return 0
}

# Enhanced service check function
check_services() {
    log_message "INFO" "Checking for CrowdSec services..."
    systemctl list-units --type=service --all | grep -i crowdsec | awk '{print $1}' > /tmp/crowdsec_services.tmp
    
    if [ ! -s /tmp/crowdsec_services.tmp ]; then
        log_message "SUCCESS" "No CrowdSec services found."
        services=""
        return 0
    else
        services=$(cat /tmp/crowdsec_services.tmp)
        log_message "WARNING" "Found CrowdSec services:"
        echo "$services" | while read -r service; do
            log_message "INFO" "Found service: $service"
        done
        return 0
    fi
}

# Enhanced package check function
check_packages() {
    log_message "INFO" "Checking for CrowdSec packages..."
    
    if command_exists dpkg; then
        dpkg -l | grep -i crowdsec | awk '{print $2}' > /tmp/crowdsec_packages.tmp
    elif command_exists rpm; then
        rpm -qa | grep -i crowdsec > /tmp/crowdsec_packages.tmp
    else
        log_message "ERROR" "Unable to determine package manager."
        return 1
    fi
    
    if [ ! -s /tmp/crowdsec_packages.tmp ]; then
        log_message "SUCCESS" "No CrowdSec packages found."
        packages=""
    else
        packages=$(cat /tmp/crowdsec_packages.tmp)
        log_message "WARNING" "Found CrowdSec packages:"
        echo "$packages" | while read -r package; do
            log_message "INFO" "Found package: $package"
        done
    fi
    return 0
}

# Enhanced service disable function
disable_services() {
    log_message "INFO" "Disabling CrowdSec services..."
    if [ -z "$services" ]; then
        log_message "WARNING" "No CrowdSec services to disable."
        return 0
    fi
    
    echo "$services" | while read -r service; do
        systemctl stop "$service" 2>/dev/null
        if [ $? -eq 0 ]; then
            systemctl disable "$service" 2>/dev/null
            log_message "SUCCESS" "Disabled $service"
        else
            log_message "ERROR" "Failed to stop $service"
            return 1
        fi
    done
    return 0
}

# Enhanced package removal function
remove_packages() {
    log_message "INFO" "Removing CrowdSec packages..."
    if [ -z "$packages" ]; then
        log_message "WARNING" "No CrowdSec packages to remove."
        return 0
    fi
    
    local remove_cmd=""
    if command_exists apt-get; then
        remove_cmd="apt-get remove --purge -y"
        cleanup_cmd="apt-get autoremove -y"
    elif command_exists yum; then
        remove_cmd="yum remove -y"
        cleanup_cmd="yum autoremove -y"
    elif command_exists dnf; then
        remove_cmd="dnf remove -y"
        cleanup_cmd="dnf autoremove -y"
    else
        log_message "ERROR" "No supported package manager found."
        return 1
    fi
    
    local failed=0
    echo "$packages" | while read -r package; do
        if $remove_cmd "$package"; then
            log_message "SUCCESS" "Removed package: $package"
        else
            log_message "ERROR" "Failed to remove package: $package"
            failed=1
        fi
    done
    
    if [ $failed -eq 0 ] && $cleanup_cmd; then
        log_message "SUCCESS" "Cleaned up dependencies"
        return 0
    fi
    return 1
}

# Enhanced repository removal function
remove_apt_repos() {
    log_message "INFO" "Removing CrowdSec repository entries..."
    if ! command_exists apt-get; then
        log_message "INFO" "System is not Debian/Ubuntu based, skipping repo removal"
        return 0
    fi
    
    local repo_files=("/etc/apt/sources.list" "/etc/apt/sources.list.d/"*)
    local failed=0
    
    for file in "${repo_files[@]}"; do
        if [ -f "$file" ] && grep -q "crowdsec" "$file"; then
            if sed -i '/crowdsec/d' "$file"; then
                log_message "SUCCESS" "Removed CrowdSec entry from $file"
            else
                log_message "ERROR" "Failed to remove CrowdSec entry from $file"
                failed=1
            fi
        fi
    done
    
    if [ $failed -eq 0 ] && apt-get update >/dev/null 2>&1; then
        log_message "SUCCESS" "APT repositories updated"
        return 0
    fi
    return 1
}

# Enhanced configuration removal function
remove_config() {
    log_message "INFO" "Removing CrowdSec configuration files..."
    
    local config_dirs=(
        "/etc/crowdsec"
        "/var/lib/crowdsec"
        "/usr/local/lib/crowdsec"
        "/usr/share/crowdsec"
    )
    
    local failed=0
    for dir in "${config_dirs[@]}"; do
        if [ -d "$dir" ]; then
            if rm -rf "$dir" 2>/dev/null; then
                log_message "SUCCESS" "Removed directory: $dir"
            else
                log_message "ERROR" "Failed to remove directory: $dir"
                failed=1
            fi
        fi
    done
    
    [ $failed -eq 0 ] && return 0 || return 1
}

#-----------------------------------#
#             MAIN LOGIC            #
#-----------------------------------#

# Print header
print_header "$this_script_name" "$this_script_url"

echo -e "Running as User: $USER_TO_RUN_AS\nUser Home: $USER_HOME\n"

# Check requirements
if ! format_output check_root "Checking root privileges"; then
    cleanup
    exit 1
fi

if ! format_output check_system_requirements "Checking system requirements"; then
    cleanup
    exit 1
fi

# Initialize log file
touch "$LOGFILE" || { echo "Cannot create log file"; exit 1; }
log_message "INFO" "Starting CrowdSec removal script v${VERSION}"

# Check current state
format_output check_services "Checking CrowdSec services"
format_output check_packages "Checking CrowdSec packages"

if [ -z "$services" ] && [ -z "$packages" ]; then
    log_message "SUCCESS" "No CrowdSec components found on your system."
    cleanup
    exit 0
fi

# Interactive menu
log_message "INFO" "What would you like to do?"
echo "1) Disable CrowdSec services"
echo "2) Remove CrowdSec packages"
echo "3) Remove everything (services, packages, configuration, and repositories)"
echo "4) Exit without changes"

read -r -p "Enter your choice (1-4): " choice

case $choice in
    1) format_output disable_services "Disabling CrowdSec services" ;;
    2) format_output remove_packages "Removing CrowdSec packages" ;;
    3)
        format_output create_backup "Creating backup"
        format_output disable_services "Disabling services"
        format_output remove_packages "Removing packages"
        format_output remove_apt_repos "Removing repositories"
        format_output remove_config "Removing configuration"
        ;;
    4) log_message "INFO" "Exiting without changes."; cleanup; exit 0 ;;
    *) log_message "ERROR" "Invalid choice."; cleanup; exit 1 ;;
esac

log_message "SUCCESS" "Operations completed successfully."
echo "Log file available at: $LOGFILE"

# Exit with appropriate status
exit $success