#!/bin/bash

################################################################################
# System Update and Maintenance Script
################################################################################
#
# Author: Peter Wood <peter@peterwood.dev>
# Description: Comprehensive system update script with advanced error detection,
#              automatic remediation of common package manager issues, and
#              intelligent service management.
#
# Features:
#   - Cross-platform package manager detection (apt/nala/dnf)
#   - Automatic detection and remediation of corrupt package lists
#   - APT database lock handling and recovery
#   - Broken dependency detection and repair
#   - Service-aware updating with Plex management
#   - Oh My Zsh upgrade integration
#   - Comprehensive error handling and logging
#   - Detailed debugging output for troubleshooting
#
# Common Issues Remediated:
#   - Corrupt package cache and lists
#   - APT database locks and conflicts
#   - Interrupted package installations
#   - Broken or missing dependencies
#   - GPG key validation errors
#   - Network connectivity issues during updates
#
# Usage:
#   ./update.sh                    # Standard system update
#   ./update.sh --debug           # Enhanced debugging output
#   ./update.sh --force-repair    # Force package manager repair
#   ./update.sh --skip-services   # Skip all service management
#   ./update.sh --skip-plex       # Skip only Plex-related operations
#
# Dependencies:
#   - sudo privileges
#   - Package manager (apt/nala/dnf)
#   - systemctl (for service management)
#   - Oh My Zsh (optional)
#
# Exit Codes:
#   0 - Success
#   1 - General error
#   2 - Package manager issues
#   3 - Service management failure
#   4 - Permission denied
#
################################################################################

set -e

# Color codes for output formatting
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[0;33m'
readonly RED='\033[0;31m'
readonly CYAN='\033[0;36m'
readonly NC='\033[0m' # No Color

# Configuration
if [[ -w "/var/log" ]]; then
    LOG_FILE="/var/log/system-update.log"
else
    LOG_FILE="$HOME/.local/share/system-update.log"
    mkdir -p "$(dirname "$LOG_FILE")"
fi
readonly LOG_FILE

# Global variables
ERRORS_DETECTED=0
REMEDIATION_ATTEMPTED=0
PKG_MANAGER=""
OS_NAME=""
OS_VERSION=""

# Command line flags
DEBUG_MODE=false
FORCE_REPAIR=false
SKIP_SERVICES=false
SKIP_PLEX=false
TEST_MODE=false

################################################################################
# Utility Functions
################################################################################

# Logging function with timestamps
log_message() {
    local level="$1"
    local message="$2"
    local timestamp
    local color=""

    timestamp=$(date '+%Y-%m-%d %H:%M:%S')

    case "$level" in
        "INFO")  color="$GREEN" ;;
        "WARN")  color="$YELLOW" ;;
        "ERROR") color="$RED" ;;
        "DEBUG") color="$CYAN" ;;
        *)       color="$NC" ;;
    esac

    echo -e "${color}[$timestamp] [$level] $message${NC}"

    # Log to file if possible
    if [[ -w "$(dirname "$LOG_FILE")" ]] 2>/dev/null; then
        echo "[$timestamp] [$level] $message" >> "$LOG_FILE" 2>/dev/null || true
    fi
}

# Debug logging (only when debug mode enabled)
debug_log() {
    if [[ "$DEBUG_MODE" == true ]]; then
        log_message "DEBUG" "$1"
    fi
}

# Error counter
increment_error() {
    ((ERRORS_DETECTED++))
    log_message "ERROR" "$1"
}

# Remediation counter
increment_remediation() {
    ((REMEDIATION_ATTEMPTED++))
    log_message "INFO" "REMEDIATION: $1"
}

# Check if running with sufficient privileges
check_privileges() {
    if [[ "$TEST_MODE" == true ]]; then
        debug_log "Test mode enabled - skipping privilege check"
        return 0
    fi

    if [[ $EUID -eq 0 ]]; then
        log_message "WARN" "Running as root - this is not recommended"
        return 0
    fi

    # Check if sudo is available, but don't require it upfront
    if sudo -n true 2>/dev/null; then
        debug_log "Sudo privileges confirmed and cached"
    else
        log_message "INFO" "Sudo credentials not cached - you may be prompted for password during updates"
        log_message "INFO" "To avoid prompts, run: sudo -v before running this script"
    fi

    return 0
}

# Cleanup function for script exit
cleanup() {
    local exit_code=$?
    if [[ $exit_code -ne 0 ]]; then
        log_message "ERROR" "Script exited with error code: $exit_code"
    fi

    log_message "INFO" "Script completed - Errors: $ERRORS_DETECTED, Remediations: $REMEDIATION_ATTEMPTED"
    exit $exit_code
}

trap cleanup EXIT ERR

################################################################################
# Argument Parsing
################################################################################

# Parse command line arguments
parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --debug)
                DEBUG_MODE=true
                debug_log "Debug mode enabled"
                shift
                ;;
            --force-repair)
                FORCE_REPAIR=true
                shift
                ;;
            --skip-services)
                SKIP_SERVICES=true
                shift
                ;;
            --skip-plex)
                SKIP_PLEX=true
                shift
                ;;
            --test-mode)
                TEST_MODE=true
                shift
                ;;
            -h|--help)
                show_help
                exit 0
                ;;
            *)
                log_message "WARN" "Unknown argument: $1"
                shift
                ;;
        esac
    done
}

# Show help message
show_help() {
    cat << EOF
System Update and Maintenance Script

Usage: $0 [OPTIONS]

OPTIONS:
    --debug           Enable enhanced debugging output
    --force-repair    Force package manager repair before updates
    --skip-services   Skip all service management (includes Plex)
    --skip-plex       Skip only Plex-related operations
    --test-mode       Test flag parsing without running privileged commands
    -h, --help        Show this help message

Examples:
    $0                    # Standard system update
    $0 --debug           # Enhanced debugging output
    $0 --skip-plex       # Skip Plex service management
    $0 --force-repair    # Force package manager repair
    $0 --skip-services   # Skip all service management
    $0 --test-mode       # Test flag parsing without privileged commands

Exit Codes:
    0 - Success
    1 - General error
    2 - Package manager issues
    3 - Service management failure
    4 - Permission denied
EOF
}

################################################################################
# System Detection and Package Manager Setup
################################################################################

# Detect operating system
detect_os() {
    log_message "INFO" "Detecting operating system..."

    if [[ -f /etc/os-release ]]; then
        # shellcheck disable=SC1091
        . /etc/os-release
        OS_NAME="$ID"
        OS_VERSION="$VERSION_ID"
        log_message "INFO" "Detected OS: $OS_NAME $OS_VERSION"
    else
        log_message "WARN" "Unable to detect OS from /etc/os-release, assuming Ubuntu"
        OS_NAME="ubuntu"
        OS_VERSION="unknown"
    fi

    debug_log "OS detection completed: $OS_NAME ($OS_VERSION)"
}

# Determine the best package manager to use
determine_pkg_manager() {
    log_message "INFO" "Determining package manager..."

    if command -v nala &> /dev/null; then
        PKG_MANAGER="nala"
        log_message "INFO" "Using nala (enhanced apt frontend)"
    elif [[ "$OS_NAME" == "fedora" ]] || [[ "$OS_NAME" == "rhel" ]] || [[ "$OS_NAME" == "centos" ]]; then
        PKG_MANAGER="dnf"
        log_message "INFO" "Using dnf (Red Hat family)"
    elif command -v apt &> /dev/null; then
        PKG_MANAGER="apt"
        log_message "INFO" "Using apt (Debian family)"
    else
        increment_error "No supported package manager found"
        exit 2
    fi

    debug_log "Package manager determination completed: $PKG_MANAGER"
}

################################################################################
# Error Detection and Remediation Functions
################################################################################

# Check for and fix APT database locks
fix_apt_locks() {
    log_message "INFO" "Checking for APT database locks..."

    local lock_files=(
        "/var/lib/dpkg/lock"
        "/var/lib/dpkg/lock-frontend"
        "/var/cache/apt/archives/lock"
        "/var/lib/apt/lists/lock"
    )

    local locks_found=false

    for lock_file in "${lock_files[@]}"; do
        if [[ -f "$lock_file" ]]; then
            debug_log "Found lock file: $lock_file"
            locks_found=true
        fi
    done

    if $locks_found; then
        increment_remediation "Removing APT database locks"

        # Kill any running apt processes
        sudo pkill -f apt &>/dev/null || true
        sudo pkill -f dpkg &>/dev/null || true

        # Remove lock files
        for lock_file in "${lock_files[@]}"; do
            if [[ -f "$lock_file" ]]; then
                sudo rm -f "$lock_file"
                debug_log "Removed lock file: $lock_file"
            fi
        done

        # Reconfigure interrupted packages
        sudo dpkg --configure -a
        log_message "INFO" "APT locks cleared and packages reconfigured"
    else
        debug_log "No APT locks detected"
    fi
}

# Check for and fix corrupt package lists
fix_corrupt_lists() {
    log_message "INFO" "Checking for corrupt package lists..."

    local lists_dir="/var/lib/apt/lists"
    local corrupt_found=false

    # Check if lists directory exists and is accessible
    if [[ ! -d "$lists_dir" ]]; then
        increment_error "APT lists directory not found: $lists_dir"
        return 1
    fi

    # Look for partial or corrupt files
    if find "$lists_dir" -name "*_Packages" -size 0 2>/dev/null | grep -q .; then
        corrupt_found=true
        debug_log "Found zero-size package files"
    fi

    if find "$lists_dir" -name "partial" -type d 2>/dev/null | grep -q .; then
        if [[ -n "$(ls -A "$lists_dir/partial" 2>/dev/null)" ]]; then
            corrupt_found=true
            debug_log "Found files in partial directory"
        fi
    fi

    if $corrupt_found; then
        increment_remediation "Cleaning corrupt package lists"

        # Remove partial downloads and corrupt files
        sudo rm -rf "$lists_dir/partial/"*
        sudo find "$lists_dir" -name "*_Packages" -size 0 -delete

        log_message "INFO" "Corrupt package lists cleaned"
    else
        debug_log "No corrupt package lists detected"
    fi
}

# Check network connectivity
check_network_connectivity() {
    log_message "INFO" "Checking network connectivity..."

    local test_hosts=("8.8.8.8" "1.1.1.1")
    local connectivity=false

    for host in "${test_hosts[@]}"; do
        if ping -c 1 -W 3 "$host" &>/dev/null; then
            connectivity=true
            debug_log "Network connectivity confirmed via $host"
            break
        fi
    done

    if ! $connectivity; then
        increment_error "No network connectivity detected"
        return 1
    fi

    debug_log "Network connectivity check passed"
    return 0
}

# Repair broken dependencies
fix_broken_dependencies() {
    log_message "INFO" "Checking for broken dependencies..."

    case "$PKG_MANAGER" in
        apt|nala)
            # Check for broken packages
            local broken_output
            broken_output=$(sudo apt-get check 2>&1 || true)

            if echo "$broken_output" | grep -q "broken"; then
                increment_remediation "Fixing broken dependencies"
                sudo apt-get install -f -y
                sudo apt-get autoremove -y
                log_message "INFO" "Broken dependencies repair attempted"
            else
                debug_log "No broken dependencies detected"
            fi
            ;;
        dnf)
            # DNF has built-in dependency resolution
            debug_log "DNF handles dependencies automatically"
            ;;
    esac
}

# Comprehensive package manager health check
check_package_manager_health() {
    log_message "INFO" "Performing package manager health check..."

    case "$PKG_MANAGER" in
        apt|nala)
            check_network_connectivity || return 1
            fix_apt_locks
            fix_corrupt_lists
            fix_broken_dependencies
            ;;
        dnf)
            check_network_connectivity || return 1
            # DNF is generally more robust, fewer common issues
            debug_log "DNF health check completed"
            ;;
    esac

    log_message "INFO" "Package manager health check completed"
}

################################################################################
# Service Management Functions
################################################################################

# Manage Plex service during updates
manage_plex_service() {
    local action="$1"
    local plex_script="/home/acedanger/shell/plex/plex.sh"

    if ! systemctl list-unit-files plexmediaserver.service &>/dev/null; then
        debug_log "Plex Media Server service not found on this system"
        return 0
    fi

    case "$action" in
        "stop")
            if systemctl is-active --quiet plexmediaserver.service 2>/dev/null; then
                log_message "INFO" "Stopping Plex Media Server for system update"
                if [[ -x "$plex_script" ]]; then
                    PLEX_SIMPLE_OUTPUT=1 "$plex_script" stop
                else
                    sudo systemctl stop plexmediaserver.service
                fi
                debug_log "Plex service stopped successfully"
            else
                debug_log "Plex service is not running"
            fi
            ;;
        "start")
            if systemctl is-enabled --quiet plexmediaserver.service 2>/dev/null; then
                log_message "INFO" "Starting Plex Media Server after system update"
                if [[ -x "$plex_script" ]]; then
                    PLEX_SIMPLE_OUTPUT=1 "$plex_script" start
                else
                    sudo systemctl start plexmediaserver.service
                fi
                debug_log "Plex service started successfully"
            else
                debug_log "Plex service is not enabled"
            fi
            ;;
    esac
}

################################################################################
# Update Functions
################################################################################

# Upgrade Oh My Zsh
upgrade_oh_my_zsh() {
    local omz_upgrade_script="$HOME/.oh-my-zsh/tools/upgrade.sh"

    log_message "INFO" "Checking for Oh My Zsh upgrade..."

    if [[ -x "$omz_upgrade_script" ]]; then
        log_message "INFO" "Upgrading Oh My Zsh..."
        if "$omz_upgrade_script"; then
            log_message "INFO" "Oh My Zsh upgrade completed successfully"
        else
            increment_error "Oh My Zsh upgrade failed"
        fi
    else
        debug_log "Oh My Zsh not found or upgrade script not executable"
    fi
}

# Perform system package updates
perform_system_update() {
    log_message "INFO" "Starting system package update with $PKG_MANAGER..."

    case "$PKG_MANAGER" in
        nala)
            log_message "INFO" "Updating package lists with nala..."
            if ! sudo nala update; then
                increment_error "Failed to update package lists with nala"
                return 1
            fi

            log_message "INFO" "Upgrading packages with nala..."
            if ! sudo nala upgrade -y; then
                increment_error "Failed to upgrade packages with nala"
                return 1
            fi

            log_message "INFO" "Cleaning up unused packages with nala..."
            sudo nala autoremove -y
            ;;
        dnf)
            log_message "INFO" "Checking for updates with dnf..."
            sudo dnf check-update -y || true  # Exit code 100 means updates available

            log_message "INFO" "Upgrading packages with dnf..."
            if ! sudo dnf upgrade -y; then
                increment_error "Failed to upgrade packages with dnf"
                return 1
            fi

            log_message "INFO" "Cleaning up unused packages with dnf..."
            sudo dnf autoremove -y
            ;;
        apt)
            log_message "INFO" "Updating package lists with apt..."
            if ! sudo apt update; then
                increment_error "Failed to update package lists with apt"
                return 1
            fi

            log_message "INFO" "Upgrading packages with apt..."
            if ! sudo apt upgrade -y; then
                increment_error "Failed to upgrade packages with apt"
                return 1
            fi

            log_message "INFO" "Cleaning up unused packages with apt..."
            sudo apt autoremove -y && sudo apt autoclean
            ;;
    esac

    # Universal packages
    if command -v flatpak &> /dev/null; then
        log_message "INFO" "Updating Flatpak packages..."
        flatpak update -y

        log_message "INFO" "Cleaning up unused Flatpak runtimes..."
        flatpak uninstall --unused -y
    fi

    if command -v snap &> /dev/null; then
        log_message "INFO" "Updating Snap packages..."
        sudo snap refresh
    fi

    log_message "INFO" "System package update completed successfully"
}

update_signal() {
    # check if hostname is `mini`
    if [[ "$(hostname)" != "mini" ]]; then
        debug_log "Signal update is only available on host 'mini'"
        return 0
    fi

    # check if distrobox is installed
    if ! command -v distrobox-upgrade &> /dev/null; then
        debug_log "distrobox is not installed"
        return 0
    fi

    # Capture failure to prevent script exit due to set -e
    # Known issue: distrobox-upgrade may throw a stat error at the end despite success
    if ! distrobox-upgrade signal; then
        log_message "WARN" "Signal update reported an error (likely benign 'stat' issue). Continuing..."
    fi
}

################################################################################
# Main Execution
################################################################################

main() {
    # Parse command line arguments first
    parse_arguments "$@"

    log_message "INFO" "Starting system update process..."

    # Check privileges
    check_privileges

    # Detect system and package manager
    detect_os
    determine_pkg_manager

    # Pre-update health checks and repairs
    if [[ "$FORCE_REPAIR" == true ]] || ! check_package_manager_health; then
        log_message "WARN" "Package manager health issues detected, attempting repairs..."
        check_package_manager_health
    fi

    # Stop services that might interfere with updates
    if [[ "$SKIP_SERVICES" != true ]]; then
        if [[ "$SKIP_PLEX" != true ]]; then
            manage_plex_service "stop"
        else
            debug_log "Skipping Plex service stop due to --skip-plex flag"
        fi
    else
        debug_log "Skipping all service management due to --skip-services flag"
    fi

    # Perform updates
    upgrade_oh_my_zsh
    perform_system_update

    # signal is made available using distrobox and is only available on `mini`
    update_signal

    # Restart services
    if [[ "$SKIP_SERVICES" != true ]]; then
        if [[ "$SKIP_PLEX" != true ]]; then
            manage_plex_service "start"
        else
            debug_log "Skipping Plex service start due to --skip-plex flag"
        fi
    else
        debug_log "Skipping all service management due to --skip-services flag"
    fi

    # Check for reboot requirement
    if [[ -f /var/run/reboot-required ]]; then
        log_message "WARN" "A system reboot is required to complete the update."
    fi

    # Final status
    if [[ $ERRORS_DETECTED -eq 0 ]]; then
        log_message "INFO" "System update completed successfully!"
    else
        log_message "WARN" "System update completed with $ERRORS_DETECTED errors"
        exit 1
    fi
}

# Execute main function
main "$@"