#!/bin/bash

# Define color variables for formatted output
COLOR_DEBUG="\033[0;36m"   # Cyan
COLOR_INFO="\033[0;32m"    # Green
COLOR_WARNING="\033[0;33m" # Yellow
COLOR_ERROR="\033[0;31m"   # Red
COLOR_DRY_RUN="\033[0;35m" # Purple
COLOR_RESET="\033[0m"      # Reset to default

# Color Codes for all colors
COLOR_BLACK="\033[0;30m"
COLOR_RED="\033[0;31m"
COLOR_GREEN="\033[0;32m"
COLOR_YELLOW="\033[0;33m"
COLOR_BLUE="\033[0;34m"
COLOR_PURPLE="\033[0;35m"
COLOR_CYAN="\033[0;36m"
COLOR_WHITE="\033[0;37m"

# Color Codes for all colors with background
COLOR_BLACK_BG="\033[0;40m"
COLOR_RED_BG="\033[0;41m"
COLOR_GREEN_BG="\033[0;42m"
COLOR_YELLOW_BG="\033[0;43m"
COLOR_BLUE_BG="\033[0;44m"
COLOR_PURPLE_BG="\033[0;45m"
COLOR_CYAN_BG="\033[0;46m"
COLOR_WHITE_BG="\033[0;47m"

# Formatting Codes
BOLD="\033[1m"
DIM="\033[2m"
UNDERLINED="\033[4m"
BLINK="\033[5m"
INVERTED="\033[7m"
HIDDEN="\033[8m"

# Combined Formatting Codes
BOLD_BLACK="\033[1;30m"
BOLD_RED="\033[1;31m"
BOLD_GREEN="\033[1;32m"
BOLD_YELLOW="\033[1;33m"
BOLD_BLUE="\033[1;34m"
BOLD_PURPLE="\033[1;35m"
BOLD_CYAN="\033[1;36m"
BOLD_WHITE="\033[1;37m"

DIM_BLACK="\033[2;30m"
DIM_RED="\033[2;31m"
DIM_GREEN="\033[2;32m"
DIM_YELLOW="\033[2;33m"
DIM_BLUE="\033[2;34m"
DIM_PURPLE="\033[2;35m"
DIM_CYAN="\033[2;36m"
DIM_WHITE="\033[2;37m"

UNDERLINED_BLACK="\033[4;30m"
UNDERLINED_RED="\033[4;31m"
UNDERLINED_GREEN="\033[4;32m"
UNDERLINED_YELLOW="\033[4;33m"
UNDERLINED_BLUE="\033[4;34m"
UNDERLINED_PURPLE="\033[4;35m"
UNDERLINED_CYAN="\033[4;36m"
UNDERLINED_WHITE="\033[4;37m"

BLINK_BLACK="\033[5;30m"
BLINK_RED="\033[5;31m"
BLINK_GREEN="\033[5;32m"
BLINK_YELLOW="\033[5;33m"
BLINK_BLUE="\033[5;34m"
BLINK_PURPLE="\033[5;35m"
BLINK_CYAN="\033[5;36m"
BLINK_WHITE="\033[5;37m"

INVERTED_BLACK="\033[7;30m"
INVERTED_RED="\033[7;31m"
INVERTED_GREEN="\033[7;32m"
INVERTED_YELLOW="\033[7;33m"
INVERTED_BLUE="\033[7;34m"
INVERTED_PURPLE="\033[7;35m"
INVERTED_CYAN="\033[7;36m"
INVERTED_WHITE="\033[7;37m"

HIDDEN_BLACK="\033[8;30m"
HIDDEN_RED="\033[8;31m"
HIDDEN_GREEN="\033[8;32m"
HIDDEN_YELLOW="\033[8;33m"
HIDDEN_BLUE="\033[8;34m"
HIDDEN_PURPLE="\033[8;35m"
HIDDEN_CYAN="\033[8;36m"
HIDDEN_WHITE="\033[8;37m"

BOLD_DIM_BLACK="\033[1;2;30m"
BOLD_DIM_RED="\033[1;2;31m"
BOLD_DIM_GREEN="\033[1;2;32m"
BOLD_DIM_YELLOW="\033[1;2;33m"
BOLD_DIM_BLUE="\033[1;2;34m"
BOLD_DIM_PURPLE="\033[1;2;35m"
BOLD_DIM_CYAN="\033[1;2;36m"
BOLD_DIM_WHITE="\033[1;2;37m"

BOLD_UNDERLINED_BLACK="\033[1;4;30m"
BOLD_UNDERLINED_RED="\033[1;4;31m"
BOLD_UNDERLINED_GREEN="\033[1;4;32m"
BOLD_UNDERLINED_YELLOW="\033[1;4;33m"
BOLD_UNDERLINED_BLUE="\033[1;4;34m"
BOLD_UNDERLINED_PURPLE="\033[1;4;35m"
BOLD_UNDERLINED_CYAN="\033[1;4;36m"
BOLD_UNDERLINED_WHITE="\033[1;4;37m"

BOLD_BLINK_BLACK="\033[1;5;30m"
BOLD_BLINK_RED="\033[1;5;31m"
BOLD_BLINK_GREEN="\033[1;5;32m"
BOLD_BLINK_YELLOW="\033[1;5;33m"
BOLD_BLINK_BLUE="\033[1;5;34m"
BOLD_BLINK_PURPLE="\033[1;5;35m"
BOLD_BLINK_CYAN="\033[1;5;36m"
BOLD_BLINK_WHITE="\033[1;5;37m"

DIM_UNDERLINED_BLACK="\033[2;4;30m"
DIM_UNDERLINED_RED="\033[2;4;31m"
DIM_UNDERLINED_GREEN="\033[2;4;32m"
DIM_UNDERLINED_YELLOW="\033[2;4;33m"
DIM_UNDERLINED_BLUE="\033[2;4;34m"
DIM_UNDERLINED_PURPLE="\033[2;4;35m"
DIM_UNDERLINED_CYAN="\033[2;4;36m"
DIM_UNDERLINED_WHITE="\033[2;4;37m"

DIM_BLINK_BLACK="\033[2;5;30m"
DIM_BLINK_RED="\033[2;5;31m"
DIM_BLINK_GREEN="\033[2;5;32m"
DIM_BLINK_YELLOW="\033[2;5;33m"
DIM_BLINK_BLUE="\033[2;5;34m"
DIM_BLINK_PURPLE="\033[2;5;35m"
DIM_BLINK_CYAN="\033[2;5;36m"
DIM_BLINK_WHITE="\033[2;5;37m"

BOLD_DIM_UNDERLINED_BLACK="\033[1;2;4;30m"
BOLD_DIM_UNDERLINED_RED="\033[1;2;4;31m"
BOLD_DIM_UNDERLINED_GREEN="\033[1;2;4;32m"
BOLD_DIM_UNDERLINED_YELLOW="\033[1;2;4;33m"
BOLD_DIM_UNDERLINED_BLUE="\033[1;2;4;34m"
BOLD_DIM_UNDERLINED_PURPLE="\033[1;2;4;35m"
BOLD_DIM_UNDERLINED_CYAN="\033[1;2;4;36m"
BOLD_DIM_UNDERLINED_WHITE="\033[1;2;4;37m"

# Function to print help
print_help() {
    echo "Usage: $0 [options]"
    echo "Options:"
    echo "  -s            Get the package status of each container"
    echo "  -u            Update the package cache, but do not upgrade"
    echo "  -U            Perform upgrades on all containers"
    echo "  -x <ID>       Exclude a container by its ID (can be used multiple times)"
    echo "  -i <ID>       Include only specific containers by their IDs (can be used multiple times)"
    echo "  -y            Assume 'yes' for all operations (asks once)"
    echo "  -Y            Assume 'yes' for all operations without asking (for automation)"
    echo "  -b            Boot non-running containers before performing actions"
    echo "  -B <ID>       Boot a specific container by its ID (can be used multiple times)"
    echo "  -d            Dry-run mode (simulate actions without making changes)"
    echo "  -v            Enable verbose mode"
    echo "  -h            Print this help and exit"
}

# Logging functions
log_dbg() { [ "$VERBOSE" = true ] && echo -e "${COLOR_DEBUG}[$(date '+%Y-%m-%d %H:%M:%S')][DEBUG] $1${COLOR_RESET}"; }
log_inf() { echo -e "${COLOR_INFO}[$(date '+%Y-%m-%d %H:%M:%S')][INFO] $1${COLOR_RESET}"; }
log_wrn() { echo -e "${COLOR_WARNING}[$(date '+%Y-%m-%d %H:%M:%S')][WARNING] $1${COLOR_RESET}"; }
log_err() { echo -e "${COLOR_ERROR}[$(date '+%Y-%m-%d %H:%M:%S')][ERROR] $1${COLOR_RESET}"; }
log_dry_run() { echo -e "${COLOR_DRY_RUN}[DRY-RUN] $1${COLOR_RESET}"; }

# Check for required commands
for cmd in awk grep pct; do
    if ! command -v $cmd &> /dev/null; then
        log_err "Required command $cmd is not available."
        exit 1
    fi
done

# Warn if running as root
if [ "$(id -u)" -eq 0 ]; then
    log_wrn "Running as root. Ensure this is necessary for your operations."
fi

# Initialize variables
EXCLUDE_IDS=()
INCLUDE_IDS=()
BOOT_IDS=()
ACTION=""
ASSUME_YES=false
FORCE_YES=false
BOOT_CONTAINERS=false
DRY_RUN=false
VERBOSE=false

# Internal variables
STATUS_SUMMARIES=()

# Parse arguments
while getopts ":suUx:i:hyYbdvB:" opt; do
    case ${opt} in
        s|u|U)
            if [[ -n "$ACTION" && "$ACTION" != "$opt" ]]; then
                log_err "Multiple actions specified. Please specify only one of -s, -u, or -U."
                exit 1
            fi
            ACTION="$opt"
            ;;
        x) EXCLUDE_IDS+=("$OPTARG") ;;
        i) INCLUDE_IDS+=("$OPTARG") ;;
        y) ASSUME_YES=true ;;
        Y) FORCE_YES=true ;;
        b) BOOT_CONTAINERS=true ;;
        B) BOOT_IDS+=("$OPTARG") ;;
        d) DRY_RUN=true ;;
        v) VERBOSE=true ;;
        h) print_help; exit 0 ;;
        \?) log_err "Invalid option: -$OPTARG"; exit 1 ;;
        :) log_err "Option -$OPTARG requires an argument."; exit 1 ;;
    esac
done

# Check for conflicts between -b and -B
if [ "$BOOT_CONTAINERS" = true ] && [ ${#BOOT_IDS[@]} -gt 0 ]; then
    log_err "Conflict: Both -b and -B options are set. Use only one. -b will boot all containers, which may not be intended."
    exit 1
fi

# Check if an action is specified
if [ -z "$ACTION" ]; then
    log_err "No action specified. Use -h for help."
    exit 1
fi

# Check for conflicts between include and exclude lists
for id in "${INCLUDE_IDS[@]}"; do
    if [[ " ${EXCLUDE_IDS[@]} " =~ " $id " ]]; then
        log_err "Container ID $id is both included and excluded. This is not allowed."
        exit 1
    fi
done

# Get list of all containers
log_inf "Retrieving list of containers..."
CONTAINERS=$(pct list | awk 'NR>1 {print $1}' | grep -E '^[0-9]+$')
log_dbg "Containers found: $CONTAINERS"

# Function to check if a container ID is in the exclude list
is_excluded() {
    local id="$1"
    for exclude_id in "${EXCLUDE_IDS[@]}"; do
        if [ "$exclude_id" == "$id" ]; then
            return 0
        fi
    done
    return 1
}

# Function to check if a container ID is in the include list
is_included() {
    local id="$1"
    if [ ${#INCLUDE_IDS[@]} -eq 0 ]; then
        return 0
    fi
    for include_id in "${INCLUDE_IDS[@]}"; do
        if [ "$include_id" == "$id" ]; then
            return 0
        fi
    done
    return 1
}

# Function to check if a container ID is in the boot list
is_booted() {
    local id="$1"
    for boot_id in "${BOOT_IDS[@]}"; do
        if [ "$boot_id" == "$id" ]; then
            return 0
        fi
    done
    return 1
}

# Function to confirm actions
confirm_action() {
    local message="$1"
    if $FORCE_YES; then
        return 0
    elif $ASSUME_YES; then
        log_wrn "Assuming 'yes' for all operations. Proceed with caution."
        return 0
    fi
    read -p "$message (y/n): " choice
    case "$choice" in
        y|Y ) return 0 ;;
        * ) return 1 ;;
    esac
}

# Function to validate container ID
validate_container_id() {
    local id="$1"
    if ! [[ "$id" =~ ^[0-9]+$ ]]; then
        log_err "Invalid container ID: $id"
        return 1
    fi
    return 0
}

# Function to check if a container ID exists
container_exists() {
    local id="$1"
    if pct list | awk 'NR>1 {print $1}' | grep -q "^$id$"; then
        return 0
    else
        log_err "Container ID $id does not exist."
        return 1
    fi
}

# Function to check if there are enough resources to boot a container
can_boot_container() {
    local id="$1"
    # Example check: Ensure there is enough free memory (this is a placeholder and should be replaced with actual checks)
    local required_memory=$(pct config "$id" | grep -i "memory" | awk '{print $2}')
    local free_memory=$(free -m | awk '/^Mem:/{print $7}')
    if [ "$free_memory" -lt "$required_memory" ]; then
        log_err "Not enough memory to boot container $id. Required: $required_memory MB, Available: $free_memory MB."
        return 1
    fi
    return 0
}

color() {
    local color="$1"
    local message="$2"
    local previous_color="${3:-$COLOR_RESET}"
    
    return "${color}$message${previous_color}"
}

# Filter containers based on include/exclude lists
FILTERED_CONTAINERS=()
for CTID in $CONTAINERS; do
    if ! validate_container_id "$CTID"; then
        continue
    fi
    if ! container_exists "$CTID"; then
        continue
    fi
    if is_excluded "$CTID"; then
        continue
    fi
    if is_included "$CTID"; then
        FILTERED_CONTAINERS+=("$CTID")
    fi
done

# Check if there are any containers to process
if [ ${#FILTERED_CONTAINERS[@]} -eq 0 ]; then
    log_wrn "No containers match the specified criteria."
    exit 0
fi

# Iterate over each filtered container
for CTID in "${FILTERED_CONTAINERS[@]}"; do
    log_dbg "Processing container ID: $CTID"

    # Check if the container is running
    STATUS=$(pct status "$CTID" 2>/dev/null | awk '{print $2}')
    if [ $? -ne 0 ]; then
        log_err "Failed to retrieve status for container $CTID."
        continue
    fi
    log_dbg "Container $CTID status: $STATUS"
    if [ "$STATUS" != "running" ]; then
        if $BOOT_CONTAINERS || is_booted "$CTID"; then
            log_inf "Container $CTID is not running."
            if $DRY_RUN; then
                log_dry_run "Would boot container $CTID."
            else
                if can_boot_container "$CTID"; then
                    if confirm_action "Boot container $CTID?"; then
                        if pct start "$CTID"; then
                            log_inf "Container $CTID started."
                        else
                            log_err "Failed to start container $CTID."
                            continue
                        fi
                    else
                        log_wrn "Boot skipped for container $CTID."
                        continue
                    fi
                else
                    log_wrn "Insufficient resources to boot container $CTID."
                    continue
                fi
            fi
        else
            log_wrn "Container $CTID is not running and will be skipped."
            continue
        fi
    fi

    # Get the OS of the container
    OS=$(pct config "$CTID" | grep -i "ostype" | awk '{print $2}' | tr '[:upper:]' '[:lower:]')
    CTNAME=$(pct config "$CTID" | grep -i "hostname" | awk '{print $2}')
    log_dbg "Container $CTID OS: $OS, Hostname: $CTNAME"

    # Determine package manager and commands based on OS
    case "$OS" in
        debian|ubuntu)
            PKG_STATUS_CMD="apt list --upgradable 2>/dev/null | grep -v 'Listing' | wc -l"
            UPDATE_CMD="apt update"
            UPGRADE_CMD="apt upgrade -y"
            ;;
        alpine)
            PKG_STATUS_CMD="apk version -l '<' | wc -l"
            UPDATE_CMD="apk update"
            UPGRADE_CMD="apk upgrade -y"
            ;;
        centos|fedora|rhel)
            if pct exec "$CTID" -- /bin/sh -c "command -v dnf" &> /dev/null; then
                PKG_STATUS_CMD="dnf check-update | wc -l"
                UPDATE_CMD="dnf makecache"
                UPGRADE_CMD="dnf upgrade -y"
            else
                PKG_STATUS_CMD="yum check-update | wc -l"
                UPDATE_CMD="yum makecache"
                UPGRADE_CMD="yum upgrade -y"
            fi
            ;;
        *)
            log_err "Unsupported OS for container $CTID: $OS"
            continue
            ;;
    esac

    log_dbg "Package status command: $PKG_STATUS_CMD"
    log_dbg "Update command: $UPDATE_CMD"
    log_dbg "Upgrade command: $UPGRADE_CMD"

    # Execute the appropriate action
    case "$ACTION" in
        s)
            log_inf "Checking package status for container $CTNAME ($CTID)..."
            if $DRY_RUN; then
                log_dry_run "Would check package status for container $CTNAME ($CTID)."
            else
                UPDATES=$(pct exec "$CTID" -- /bin/sh -c "$PKG_STATUS_CMD")
                if [ $? -ne 0 ]; then
                    log_err "Failed to check package status for container $CTNAME ($CTID)."
                    continue
                fi
                log_dbg "Updates available for container $CTNAME ($CTID): $UPDATES"
                if [ "$UPDATES" -gt 0 ]; then
                    log_inf "Container $CTNAME ($CTID) STATUS: \033[1mHas Updates ($UPDATES)\033[0m"
                else
                    log_dbg "Container $CTNAME ($CTID) STATUS: No Updates available"
                fi

                # Save the summary for later
                STATUS_SUMMARIES+=("$CTNAME ($CTID): $UPDATES")
            fi
            ;;
        u)
            log_inf "Preparing to update package cache for container $CTNAME ($CTID)..."
            if $DRY_RUN; then
                log_dry_run "Would execute update on $CTNAME ($CTID) with command: $UPDATE_CMD"
            else
                if confirm_action "Execute update on $CTNAME ($CTID) with command: $UPDATE_CMD"; then
                    if pct exec "$CTID" -- /bin/sh -c "$UPDATE_CMD"; then
                        log_inf "Update completed for $CTNAME ($CTID)"
                    else
                        log_err "Update failed for $CTNAME ($CTID)"
                        continue
                    fi
                else
                    log_wrn "Update skipped for $CTNAME ($CTID)"
                fi
            fi
            ;;
        U)
            log_inf "Preparing to upgrade packages for container $CTNAME ($CTID)..."
            if $DRY_RUN; then
                log_dry_run "Would execute upgrade on $CTNAME ($CTID) with command: $UPDATE_CMD && $UPGRADE_CMD"
            else
                if confirm_action "Execute upgrade on $CTNAME ($CTID) with command: $UPDATE_CMD && $UPGRADE_CMD"; then
                    if pct exec "$CTID" -- /bin/sh -c "$UPDATE_CMD && $UPGRADE_CMD"; then
                        log_inf "Upgrade completed for $CTNAME ($CTID)"
                    else
                        log_err "Upgrade failed for $CTNAME ($CTID)"
                        continue
                    fi
                else
                    log_wrn "Upgrade skipped for $CTNAME ($CTID)"
                fi
            fi
            ;;
    esac
done

# Print summary of statuses
if [ "$ACTION" == "s" ] && [ ${#STATUS_SUMMARIES[@]} -gt 0 ]; then
    log_inf "Summary of container statuses:"
    for summary in "${STATUS_SUMMARIES[@]}"; do
        log_inf "$summary"
    done
fi