update.sh

#!/bin/bash
#
# ██████╗ ██████╗  █████╗ ██████╗ ██╗  ██╗██╗   ██╗ ██████╗ ███████╗
# ██╔══██╗██╔══██╗██╔══██╗██╔══██╗██║  ██║╚██╗ ██╔╝██╔═══██╗██╔════╝
# ██║  ██║██████╔╝███████║██████╔╝███████║ ╚████╔╝ ██║   ██║███████╗
# ██║  ██║██╔══██╗██╔══██║██╔═══╝ ██╔══██║  ╚██╔╝  ██║   ██║╚════██║
# ██████╔╝██║  ██║██║  ██║██║     ██║  ██║   ██║   ╚██████╔╝███████║
# ╚═════╝ ╚═╝  ╚═╝╚═╝  ╚═╝╚═╝     ╚═╝  ╚═╝   ╚═╝    ╚═════╝ ╚══════╝
#
# draphyOS System Updater
# https://github.com/draphy/draphyOS
#
# A safe, interactive system update tool for Fedora
#
# Options:
#   --vm      Enable VM mode (xrender backend for picom)
#   --no-vm   Disable VM mode (glx backend for picom)
#

set -u

# Validate HOME directory
if [ -z "${HOME:-}" ] || [ ! -d "$HOME" ]; then
    echo "ERROR: HOME directory not set or does not exist"
    exit 1
fi

if [ ! -w "$HOME" ]; then
    echo "ERROR: HOME directory is not writable"
    exit 1
fi

# draphyOS paths
INSTALL_DIR="$HOME/.draphyOS"
MARKER_FILE="$HOME/.draphyOS-installed"
CHECKSUM_FILE="$HOME/.draphyOS-config-checksums"
UPGRADE_STATE_FILE="$HOME/.draphyOS-upgrade-state"

# Colors (with fallback for non-color terminals)
if [[ -t 1 ]] && [[ "${TERM:-dumb}" != "dumb" ]]; then
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    YELLOW='\033[1;33m'
    BLUE='\033[0;34m'
    CYAN='\033[0;36m'
    MINT='\033[38;2;154;184;124m'
    BOLD='\033[1m'
    NC='\033[0m'
else
    RED=''
    GREEN=''
    YELLOW=''
    BLUE=''
    CYAN=''
    MINT=''
    BOLD=''
    NC=''
fi

# Print functions (matching install.sh style)
print_header() {
    clear
    echo ""
    echo -e "${MINT}"
    echo "  ██████╗ ██████╗  █████╗ ██████╗ ██╗  ██╗██╗   ██╗ ██████╗ ███████╗"
    echo "  ██╔══██╗██╔══██╗██╔══██╗██╔══██╗██║  ██║╚██╗ ██╔╝██╔═══██╗██╔════╝"
    echo "  ██║  ██║██████╔╝███████║██████╔╝███████║ ╚████╔╝ ██║   ██║███████╗"
    echo "  ██║  ██║██╔══██╗██╔══██║██╔═══╝ ██╔══██║  ╚██╔╝  ██║   ██║╚════██║"
    echo "  ██████╔╝██║  ██║██║  ██║██║     ██║  ██║   ██║   ╚██████╔╝███████║"
    echo "  ╚═════╝ ╚═╝  ╚═╝╚═╝  ╚═╝╚═╝     ╚═╝  ╚═╝   ╚═╝    ╚═════╝ ╚══════╝"
    echo -e "${NC}"
    echo -e "                    ${CYAN}System Updater${NC}"
    echo ""
}

print_step() {
    echo -e "${GREEN}[*]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[!]${NC} $1"
}

print_error() {
    echo -e "${RED}[✗]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[✓]${NC} $1"
}

print_info() {
    echo -e "${CYAN}[i]${NC} $1"
}

# Get system information
get_system_info() {
    FEDORA_VERSION=$(rpm -E %fedora 2>/dev/null || echo "unknown")
    KERNEL_VERSION=$(uname -r)
    LAST_UPDATE=$(rpm -qa --last 2>/dev/null | head -1 | awk '{print $2, $3, $4}' || echo "unknown")
}

# Display system status with proper alignment
show_system_status() {
    get_system_info

    # Calculate dynamic box width based on content
    local header="System Information"
    local fedora_text="Fedora Version: $FEDORA_VERSION"
    local kernel_text="Kernel: $KERNEL_VERSION"
    local update_text="Last Package Change: $LAST_UPDATE"

    # Find the longest line (add 4 for "  " padding on each side)
    local max_len=${#header}
    [ ${#fedora_text} -gt $max_len ] && max_len=${#fedora_text}
    [ ${#kernel_text} -gt $max_len ] && max_len=${#kernel_text}
    [ ${#update_text} -gt $max_len ] && max_len=${#update_text}

    # Box width = max content + 4 (2 spaces padding each side)
    local box_width=$((max_len + 4))
    # Minimum width of 49 for aesthetics
    [ $box_width -lt 49 ] && box_width=49

    # Generate horizontal border
    local border=""
    for ((i=0; i<box_width; i++)); do border+="─"; done

    echo -e "  ┌${border}┐"
    printf "  │  ${BOLD}%-$((box_width - 4))s${NC}  │\n" "$header"
    echo -e "  ├${border}┤"

    # Content lines with dynamic padding
    # Format: "  │  <label>: " + value + "  │" = box_width inside
    printf "  │  Fedora Version: ${GREEN}%-$((box_width - 20))s${NC}  │\n" "$FEDORA_VERSION"
    printf "  │  Kernel: ${CYAN}%-$((box_width - 12))s${NC}  │\n" "$KERNEL_VERSION"
    printf "  │  Last Package Change: ${YELLOW}%-$((box_width - 25))s${NC}  │\n" "$LAST_UPDATE"

    echo -e "  └${border}┘"
    echo ""
}

# Check if running as root (we need sudo, not root)
check_not_root() {
    if [ "$(id -u)" -eq 0 ]; then
        print_error "Please run this script as a normal user, not as root."
        print_info "The script will ask for sudo when needed."
        exit 1
    fi
}

# Check and get sudo access
check_sudo() {
    print_step "Checking sudo access..."

    if ! sudo -v 2>/dev/null; then
        print_error "Sudo access required."
        exit 1
    fi

    # Keep sudo alive in background (with error handling)
    (while true; do sudo -v 2>/dev/null || break; sleep 50; done) &
    SUDO_KEEPER_PID=$!

    # Cleanup function that handles sudo keeper
    cleanup_sudo() {
        if [ -n "${SUDO_KEEPER_PID:-}" ] && kill -0 "$SUDO_KEEPER_PID" 2>/dev/null; then
            kill "$SUDO_KEEPER_PID" 2>/dev/null
        fi
    }
    trap cleanup_sudo EXIT INT TERM

    print_success "Sudo access OK"
    echo ""
}

# Check for updates without installing
check_updates() {
    print_step "Checking for available updates..."
    echo ""

    local updates
    updates=$(dnf check-update 2>/dev/null)
    local exit_code=$?

    # dnf check-update returns 100 if updates available, 0 if none, 1 on error
    if [ $exit_code -eq 100 ]; then
        local count
        count=$(echo "$updates" | grep -v "^$" | grep -v "^Last metadata" | grep -v "^Obsoleting" | wc -l)
        print_info "Found ${YELLOW}$count${NC} packages with updates available"
        echo ""
        echo "$updates" | head -30
        if [ "$count" -gt 30 ]; then
            echo -e "${YELLOW}... and $((count - 30)) more${NC}"
        fi
        return 0
    elif [ $exit_code -eq 0 ]; then
        print_success "System is up to date! No updates available."
        return 1
    else
        print_error "Error checking for updates"
        return 1
    fi
}

# Update all packages
do_update() {
    print_step "Preparing to update packages..."
    echo ""

    # First check what updates are available
    print_info "Refreshing package metadata..."
    local refresh_exit_code
    sudo dnf check-update --refresh 2>/dev/null
    refresh_exit_code=$?
    # Exit code 100 means updates available, 0 means no updates, other = error
    if [ $refresh_exit_code -ne 100 ] && [ $refresh_exit_code -ne 0 ]; then
        print_warning "Could not refresh metadata, continuing anyway..."
    fi
    echo ""

    # Confirm with user
    echo -e "Do you want to ${GREEN}install all available updates${NC}? (y/n)"
    read -r response </dev/tty || response="n"
    if [[ ! "$response" =~ ^[Yy](es)?$ ]]; then
        print_step "Update cancelled"
        return
    fi

    echo ""
    print_step "Installing updates..."
    print_info "This may take a while depending on the number of updates."
    echo ""

    if sudo dnf upgrade --refresh -y; then
        print_success "Updates installed successfully!"
    else
        print_error "Some updates may have failed. Check the output above."
    fi

    # Update Flatpak apps if installed
    if command -v flatpak &>/dev/null; then
        local flatpak_count
        flatpak_count=$(flatpak list --app --columns=application 2>/dev/null | wc -l)
        if [ "$flatpak_count" -gt 0 ]; then
            echo ""
            print_step "Updating Flatpak apps ($flatpak_count installed)..."
            if flatpak update -y 2>/dev/null; then
                print_success "Flatpak apps updated"
            else
                print_warning "Some Flatpak apps may have failed to update"
            fi
        fi
    fi

    check_reboot_needed
}

# Security updates only
do_security_update() {
    print_step "Checking for security updates..."
    echo ""

    # Check if security updates are available
    local sec_updates
    sec_updates=$(dnf updateinfo list --security 2>/dev/null | grep -v "^Last metadata" | grep -v "^$")

    if [ -z "$sec_updates" ]; then
        print_success "No security updates available."
        return
    fi

    print_info "Security updates available:"
    echo "$sec_updates" | head -20
    echo ""

    echo -e "Do you want to install ${RED}security updates only${NC}? (y/n)"
    read -r response </dev/tty || response="n"
    if [[ ! "$response" =~ ^[Yy](es)?$ ]]; then
        print_step "Security update cancelled"
        return
    fi

    echo ""
    print_step "Installing security updates..."

    if sudo dnf upgrade --security --refresh -y; then
        print_success "Security updates installed!"
    else
        print_error "Some security updates may have failed."
    fi

    check_reboot_needed
}

# Full maintenance (update + clean + autoremove)
do_full_maintenance() {
    print_step "Full System Maintenance"
    echo ""
    print_info "This will:"
    echo "  1. Update all packages"
    echo "  2. Remove orphaned packages (with your approval)"
    echo "  3. Update Flatpak apps (if installed)"
    echo "  4. Clean package cache"
    echo ""

    echo -e "Do you want to proceed with ${CYAN}full maintenance${NC}? (y/n)"
    read -r response </dev/tty || response="n"
    if [[ ! "$response" =~ ^[Yy](es)?$ ]]; then
        print_step "Maintenance cancelled"
        return
    fi

    echo ""

    # Step 1: Update packages
    print_step "Step 1/4: Updating packages..."
    if sudo dnf upgrade --refresh -y; then
        print_success "Packages updated"
    else
        print_warning "Some packages may have failed to update"
    fi
    echo ""

    # Step 2: Check for orphaned packages
    print_step "Step 2/4: Checking for orphaned packages..."
    local orphans
    orphans=$(dnf repoquery --unneeded 2>/dev/null | head -50)

    if [ -n "$orphans" ]; then
        local orphan_count
        orphan_count=$(echo "$orphans" | wc -l)
        print_info "Found ${YELLOW}$orphan_count${NC} orphaned packages:"
        echo "$orphans" | head -20
        if [ "$orphan_count" -gt 20 ]; then
            echo -e "  ${YELLOW}... and $((orphan_count - 20)) more${NC}"
        fi
        echo ""

        echo -e "Do you want to ${RED}remove orphaned packages${NC}? (y/n)"
        read -r confirm </dev/tty || confirm="n"

        if [[ "$confirm" =~ ^[Yy](es)?$ ]]; then
            print_step "Removing orphaned packages..."
            if sudo dnf autoremove -y; then
                print_success "Orphaned packages removed"
            else
                print_warning "Some packages could not be removed"
            fi
        else
            print_step "Keeping orphaned packages"
        fi
    else
        print_success "No orphaned packages found"
    fi
    echo ""

    # Step 3: Update Flatpak apps
    if command -v flatpak &>/dev/null; then
        local flatpak_count
        flatpak_count=$(flatpak list --app --columns=application 2>/dev/null | wc -l)
        if [ "$flatpak_count" -gt 0 ]; then
            print_step "Step 3/4: Updating Flatpak apps ($flatpak_count installed)..."
            if flatpak update -y 2>/dev/null; then
                print_success "Flatpak apps updated"
            else
                print_warning "Some Flatpak apps may have failed to update"
            fi
        else
            print_step "Step 3/4: No Flatpak apps installed, skipping"
        fi
    else
        print_step "Step 3/4: Flatpak not installed, skipping"
    fi
    echo ""

    # Step 4: Clean cache
    print_step "Step 4/4: Cleaning package cache..."
    if sudo dnf clean all >/dev/null 2>&1; then
        print_success "Package cache cleaned"
    else
        print_warning "Could not clean cache"
    fi

    # Clean unused Flatpak runtimes
    if command -v flatpak &>/dev/null; then
        flatpak uninstall --unused -y 2>/dev/null || true
    fi

    echo ""
    print_success "Full maintenance complete!"

    check_reboot_needed
}

# Fedora version upgrade
do_version_upgrade() {
    print_header
    echo -e "  ${BOLD}${CYAN}Fedora Version Upgrade${NC}"
    echo ""

    get_system_info
    print_info "Current Fedora version: ${GREEN}$FEDORA_VERSION${NC}"
    echo ""

    # Important warnings
    echo -e "  ${YELLOW}╔═══════════════════════════════════════════════════════════╗${NC}"
    echo -e "  ${YELLOW}║  ${BOLD}IMPORTANT: Read before proceeding${NC}                        ${YELLOW}║${NC}"
    echo -e "  ${YELLOW}╠═══════════════════════════════════════════════════════════╣${NC}"
    echo -e "  ${YELLOW}║${NC}  • Back up your important data before upgrading           ${YELLOW}║${NC}"
    echo -e "  ${YELLOW}║${NC}  • Ensure you have a stable internet connection           ${YELLOW}║${NC}"
    echo -e "  ${YELLOW}║${NC}  • The system will reboot during the upgrade              ${YELLOW}║${NC}"
    echo -e "  ${YELLOW}║${NC}  • Upgrade may take 30-60 minutes depending on system     ${YELLOW}║${NC}"
    echo -e "  ${YELLOW}║${NC}  • Fedora recommends upgrading one version at a time      ${YELLOW}║${NC}"
    echo -e "  ${YELLOW}╚═══════════════════════════════════════════════════════════╝${NC}"
    echo ""

    # Calculate available versions
    local current_ver=$FEDORA_VERSION
    local next_ver=$((current_ver + 1))
    local next_next_ver=$((current_ver + 2))

    echo -e "Available upgrade targets:"
    echo -e "  ${GREEN}1)${NC} Fedora ${next_ver} (recommended - next version)"
    echo -e "  ${YELLOW}2)${NC} Fedora ${next_next_ver} (skip one version - not officially recommended)"
    echo -e "  ${CYAN}3)${NC} Enter custom version number"
    echo -e "  ${RED}4)${NC} Cancel"
    echo ""
    echo -n "Select option [1-4]: "
    read -r choice </dev/tty || choice="4"

    local target_version=""
    case "$choice" in
        1) target_version=$next_ver ;;
        2) target_version=$next_next_ver ;;
        3)
            echo -n "Enter target Fedora version number: "
            read -r target_version </dev/tty || target_version=""
            if ! [[ "$target_version" =~ ^[0-9]+$ ]]; then
                print_error "Invalid version number"
                return
            fi
            ;;
        *)
            print_step "Version upgrade cancelled"
            return
            ;;
    esac

    if [ -z "$target_version" ]; then
        print_error "No version selected"
        return
    fi

    if [ "$target_version" -le "$current_ver" ]; then
        print_error "Target version must be higher than current version ($current_ver)"
        return
    fi

    echo ""
    print_info "Target version: Fedora ${GREEN}$target_version${NC}"
    echo ""

    # Final confirmation
    echo -e "${RED}${BOLD}WARNING:${NC} You are about to upgrade from Fedora $current_ver to Fedora $target_version"
    echo -e "This process will ${YELLOW}download packages${NC} and then ${YELLOW}reboot your system${NC}."
    echo ""
    echo -e "Type ${GREEN}yes${NC} to confirm upgrade (anything else cancels): "
    read -r confirm </dev/tty || confirm="no"

    if [ "$confirm" != "yes" ]; then
        print_step "Upgrade cancelled"
        return
    fi

    echo ""

    # Step 1: Update current system first (required)
    print_step "Step 1/4: Updating current system (required before upgrade)..."
    if ! sudo dnf upgrade --refresh -y; then
        print_error "Failed to update current system. Please fix any issues and try again."
        return
    fi
    print_success "Current system updated"
    echo ""

    # Check if reboot needed before upgrade
    if dnf needs-restarting -r >/dev/null 2>&1; then
        : # No reboot needed, continue
    else
        print_warning "A reboot is recommended before upgrading."
        echo -e "Do you want to ${YELLOW}reboot now${NC} and run the upgrade after? (y/n)"
        read -r reboot_first </dev/tty || reboot_first="n"
        if [[ "$reboot_first" =~ ^[Yy](es)?$ ]]; then
            print_info "Please run 'update-draphyOS' again after reboot to continue the upgrade."
            echo ""
            echo -e "Press Enter to reboot..."
            read -r </dev/tty
            sudo systemctl reboot
            exit 0
        fi
    fi

    # Step 2: Download upgrade packages
    print_step "Step 2/4: Downloading upgrade packages for Fedora $target_version..."
    print_info "This may take a while depending on your internet connection."
    echo ""

    local download_output
    download_output=$(mktemp)
    chmod 600 "$download_output"  # Restrict permissions on temp file

    # Use subshell with pipefail to capture dnf exit code correctly
    local dnf_exit=0
    (
        set -o pipefail
        sudo dnf system-upgrade download --releasever="$target_version" -y 2>&1 | tee "$download_output"
    ) || dnf_exit=$?
    if [ "$dnf_exit" -ne 0 ]; then
        # Check if version doesn't exist (catch 404 errors, missing repos, etc.)
        if grep -qiE "unable to find.*releasever|no such.*version|not found|status code: 404|Failed to download metadata|Cannot prepare internal mirrorlist" "$download_output" 2>/dev/null; then
            print_error "Fedora $target_version is not available yet."
            print_info "This version may not be released or may not be supported for upgrade."
            echo ""

            # Suggest the latest available version
            local latest_stable=$((FEDORA_VERSION + 1))
            print_info "Try Fedora $latest_stable if it's available, or wait for the release."
            rm -f "$download_output"
            return
        fi

        print_error "Failed to download upgrade packages."
        print_info "This may be due to package conflicts that need resolution."
        echo ""
        echo -e "  ${RED}╔═══════════════════════════════════════════════════════════╗${NC}"
        echo -e "  ${RED}║  ${BOLD}WARNING: --allowerasing option${NC}                            ${RED}║${NC}"
        echo -e "  ${RED}╠═══════════════════════════════════════════════════════════╣${NC}"
        echo -e "  ${RED}║${NC}  This option will REMOVE conflicting packages to proceed.  ${RED}║${NC}"
        echo -e "  ${RED}║${NC}  Some applications you use may be removed!                 ${RED}║${NC}"
        echo -e "  ${RED}║${NC}  Review the list carefully before confirming.              ${RED}║${NC}"
        echo -e "  ${RED}╚═══════════════════════════════════════════════════════════╝${NC}"
        echo ""
        echo -e "Do you want to retry with ${YELLOW}--allowerasing${NC}? (y/n)"
        read -r retry </dev/tty || retry="n"

        if [[ "$retry" =~ ^[Yy](es)?$ ]]; then
            print_warning "Using --allowerasing - review removed packages carefully!"
            # Use subshell with pipefail to capture dnf exit code correctly
            local dnf_exit=0
            (
                set -o pipefail
                sudo dnf system-upgrade download --releasever="$target_version" --allowerasing -y 2>&1 | tee "$download_output"
            ) || dnf_exit=$?
            if [ "$dnf_exit" -ne 0 ]; then
                if grep -qiE "unable to find.*releasever|no such.*version|not found|status code: 404|Failed to download metadata|Cannot prepare internal mirrorlist" "$download_output" 2>/dev/null; then
                    print_error "Fedora $target_version is not available."
                    rm -f "$download_output"
                    return
                fi
                print_error "Upgrade download failed. Please check the errors above."
                rm -f "$download_output"
                return
            fi
        else
            print_step "Upgrade cancelled"
            rm -f "$download_output"
            return
        fi
    fi

    rm -f "$download_output"

    print_success "Upgrade packages downloaded"
    echo ""

    # Step 3: Inform user about the upgrade process
    print_step "Step 3/4: Ready to upgrade"
    echo ""
    echo -e "  ${CYAN}╔═══════════════════════════════════════════════════════════╗${NC}"
    echo -e "  ${CYAN}║  ${BOLD}What happens next:${NC}                                       ${CYAN}║${NC}"
    echo -e "  ${CYAN}╠═══════════════════════════════════════════════════════════╣${NC}"
    echo -e "  ${CYAN}║${NC}  1. Your system will reboot                               ${CYAN}║${NC}"
    echo -e "  ${CYAN}║${NC}  2. Upgrade will run (you'll see progress on screen)      ${CYAN}║${NC}"
    echo -e "  ${CYAN}║${NC}  3. System will reboot again when complete                ${CYAN}║${NC}"
    echo -e "  ${CYAN}║${NC}  4. You'll be running Fedora $target_version!                         ${CYAN}║${NC}"
    echo -e "  ${CYAN}╚═══════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo -e "  ${GREEN}╔═══════════════════════════════════════════════════════════╗${NC}"
    echo -e "  ${GREEN}║  ${BOLD}After upgrade completes, run:${NC}                            ${GREEN}║${NC}"
    echo -e "  ${GREEN}╠═══════════════════════════════════════════════════════════╣${NC}"
    echo -e "  ${GREEN}║${NC}                                                           ${GREEN}║${NC}"
    echo -e "  ${GREEN}║${NC}    ${BOLD}update-draphyOS --resume${NC}                              ${GREEN}║${NC}"
    echo -e "  ${GREEN}║${NC}                                                           ${GREEN}║${NC}"
    echo -e "  ${GREEN}║${NC}  This will clean up upgrade cache and remove orphans.     ${GREEN}║${NC}"
    echo -e "  ${GREEN}╚═══════════════════════════════════════════════════════════╝${NC}"
    echo ""

    # Step 4: Trigger upgrade reboot
    print_step "Step 4/4: Starting upgrade..."
    echo ""
    echo -e "Do you want to ${GREEN}reboot and start the upgrade now${NC}? (y/n)"
    read -r start_upgrade </dev/tty || start_upgrade="n"

    if [[ "$start_upgrade" =~ ^[Yy](es)?$ ]]; then
        # Save state for resume after reboot
        save_upgrade_state "fedora-upgrade" "$target_version"

        print_info "Starting upgrade reboot..."
        echo ""
        echo -e "${GREEN}The system will now reboot to perform the upgrade.${NC}"
        echo -e "${YELLOW}After upgrade completes, run: ${BOLD}update-draphyOS --resume${NC}"
        echo ""
        echo -e "Press Enter to continue..."
        read -r </dev/tty

        # Use dnf5 offline reboot for Fedora 41+ (DNF5)
        if command -v dnf5 &>/dev/null; then
            sudo dnf5 offline reboot
        else
            sudo dnf system-upgrade reboot
        fi
    else
        # Save state anyway in case user reboots manually
        save_upgrade_state "fedora-upgrade" "$target_version"

        print_info "Upgrade packages are downloaded and ready."
        echo ""
        echo -e "To start the upgrade, run:"
        echo -e "  ${CYAN}sudo dnf system-upgrade reboot${NC}"
        echo ""
        echo -e "After upgrade completes, run:"
        echo -e "  ${GREEN}update-draphyOS --resume${NC}"
        echo ""
        echo -e "To cancel the pending upgrade:"
        echo -e "  ${CYAN}sudo dnf system-upgrade clean${NC}"
    fi
}

# Save upgrade state for resume after reboot
save_upgrade_state() {
    local operation="$1"
    local target_version="$2"
    echo "OPERATION=$operation" > "$UPGRADE_STATE_FILE"
    echo "TARGET_VERSION=$target_version" >> "$UPGRADE_STATE_FILE"
    echo "TIMESTAMP=$(date +%s)" >> "$UPGRADE_STATE_FILE"
}

# Clear upgrade state
clear_upgrade_state() {
    rm -f "$UPGRADE_STATE_FILE"
}

# Check for pending upgrade and resume
check_and_resume_upgrade() {
    if [ ! -f "$UPGRADE_STATE_FILE" ]; then
        return 1
    fi

    local operation target_version
    operation=$(grep "^OPERATION=" "$UPGRADE_STATE_FILE" 2>/dev/null | cut -d= -f2)
    target_version=$(grep "^TARGET_VERSION=" "$UPGRADE_STATE_FILE" 2>/dev/null | cut -d= -f2)

    if [ -z "$operation" ]; then
        clear_upgrade_state
        return 1
    fi

    print_header
    echo -e "  ${CYAN}╔═══════════════════════════════════════════════════════════╗${NC}"
    echo -e "  ${CYAN}║  ${BOLD}Pending Operation Detected${NC}                               ${CYAN}║${NC}"
    echo -e "  ${CYAN}╚═══════════════════════════════════════════════════════════╝${NC}"
    echo ""

    case "$operation" in
        "fedora-upgrade")
            print_info "Fedora upgrade to version $target_version was in progress"
            echo ""
            do_post_upgrade_cleanup "$target_version"
            ;;
        *)
            print_warning "Unknown operation: $operation"
            clear_upgrade_state
            ;;
    esac

    return 0
}

# Post-upgrade cleanup after Fedora version upgrade
do_post_upgrade_cleanup() {
    local target_version="$1"

    get_system_info
    print_info "Current Fedora version: ${GREEN}$FEDORA_VERSION${NC}"
    echo ""

    if [ "$FEDORA_VERSION" = "$target_version" ]; then
        print_success "Fedora $target_version upgrade completed successfully!"
        echo ""
    else
        print_warning "Expected version $target_version, but running $FEDORA_VERSION"
        echo ""
    fi

    print_step "Running post-upgrade cleanup..."
    echo ""

    # Clean up upgrade cache
    print_info "Cleaning system-upgrade cache..."
    if sudo dnf system-upgrade clean 2>/dev/null || sudo dnf5 offline clean 2>/dev/null; then
        print_success "Upgrade cache cleaned"
    else
        print_info "No upgrade cache to clean"
    fi
    echo ""

    # Offer to remove orphaned packages
    local orphans
    orphans=$(dnf repoquery --unneeded 2>/dev/null | head -20)

    if [ -n "$orphans" ]; then
        print_info "Orphaned packages found after upgrade:"
        echo "$orphans" | head -10
        echo ""
        echo -e "Do you want to ${YELLOW}remove orphaned packages${NC}? (y/n)"
        read -r confirm </dev/tty || confirm="n"
        if [[ "$confirm" =~ ^[Yy](es)?$ ]]; then
            sudo dnf autoremove -y
            print_success "Orphaned packages removed"
        fi
    fi

    clear_upgrade_state

    echo ""
    print_success "Post-upgrade cleanup complete!"
    echo ""
    print_info "Your system is now running Fedora $FEDORA_VERSION"
}

# Check if reboot is needed
check_reboot_needed() {
    echo ""
    print_step "Checking if reboot is required..."

    # Check for kernel updates or other reboot-requiring changes
    # dnf needs-restarting -r returns 1 if reboot needed, 0 if not
    if ! dnf needs-restarting -r >/dev/null 2>&1; then
        echo ""
        echo -e "  ${YELLOW}╔═══════════════════════════════════════════════════════════╗${NC}"
        echo -e "  ${YELLOW}║  ${BOLD}Reboot Required${NC}                                          ${YELLOW}║${NC}"
        echo -e "  ${YELLOW}╠═══════════════════════════════════════════════════════════╣${NC}"
        echo -e "  ${YELLOW}║${NC}  A system reboot is needed to apply all changes.          ${YELLOW}║${NC}"
        echo -e "  ${YELLOW}║${NC}  This is usually due to kernel or core library updates.   ${YELLOW}║${NC}"
        echo -e "  ${YELLOW}╚═══════════════════════════════════════════════════════════╝${NC}"
        echo ""

        echo -e "Do you want to ${CYAN}reboot now${NC}? (y/n)"
        read -r reboot_now </dev/tty || reboot_now="n"

        if [[ "$reboot_now" =~ ^[Yy](es)?$ ]]; then
            print_info "Rebooting system..."
            sudo systemctl reboot
        else
            print_info "Please remember to reboot your system soon."
        fi
    else
        # Check for services that need restart
        local services
        services=$(dnf needs-restarting -s 2>/dev/null | head -10)

        if [ -n "$services" ]; then
            print_info "Some services may need to be restarted:"
            echo "$services"
            echo ""
            print_info "You can restart them manually or reboot when convenient."
        else
            print_success "No reboot required. System is up to date!"
        fi
    fi
}

# ============================================================
# draphyOS Config Update Functions
# ============================================================

# Get installed draphyOS version
get_draphyos_version() {
    if [ -f "$MARKER_FILE" ]; then
        local version
        version=$(grep "^DRAPHYOS_VERSION=" "$MARKER_FILE" 2>/dev/null | cut -d= -f2)
        # Return 0 if version line doesn't exist (old installations)
        echo "${version:-0}"
    else
        echo "0"
    fi
}

# Get latest draphyOS version from repo
get_latest_version() {
    if [ -f "$INSTALL_DIR/VERSION" ]; then
        tr -d '[:space:]' < "$INSTALL_DIR/VERSION"
    else
        echo "0"
    fi
}

# Check if a config file has been modified by user
# Compares user's current file against the stored checksum of what was installed
is_config_modified() {
    local config_name="$1"  # e.g., "i3/config"
    local user_file="$2"    # Full path to user's config

    # If user file doesn't exist, not modified
    [ ! -f "$user_file" ] && return 1

    # If it's a symlink to draphyOS, not modified
    if [ -L "$user_file" ]; then
        local target
        target=$(readlink -f "$user_file" 2>/dev/null)
        if [[ "$target" == *"draphyOS"* ]]; then
            return 1  # Not modified (still a symlink)
        fi
    fi

    # Check against stored checksum (format: config_name:checksum)
    if [ -f "$CHECKSUM_FILE" ]; then
        local stored_sum
        stored_sum=$(grep "^${config_name}:" "$CHECKSUM_FILE" 2>/dev/null | cut -d: -f2)
        if [ -n "$stored_sum" ]; then
            local current_sum
            current_sum=$(md5sum "$user_file" 2>/dev/null | cut -d' ' -f1)
            if [ "$stored_sum" != "$current_sum" ]; then
                return 0  # Modified
            fi
            return 1  # Not modified (matches what we installed)
        fi
    fi

    # Default: assume modified if it's a real file (not symlink)
    return 0
}

# Three-way merge a config file
# Returns: 0=success (auto-merged or no changes), 1=conflict needs user input, 2=error
three_way_merge() {
    local config_name="$1"   # e.g., "i3/config"
    local user_file="$2"     # User's current file
    local new_file="$3"      # New version from repo
    local original_file="$4" # Original file (pre-update repo version, contains user's changes if any)

    # If user file doesn't exist, just use new version
    if [ ! -f "$user_file" ] && [ ! -L "$user_file" ]; then
        return 0
    fi

    # If original doesn't exist or is empty, can't do three-way merge
    if [ -z "$original_file" ] || [ ! -f "$original_file" ]; then
        return 1
    fi

    # Handle symlinks: compare original (user's version before git reset) vs new
    # If user_file is a symlink, the user was editing the repo file directly
    # original_file contains what the user had, new_file has upstream changes
    if [ -L "$user_file" ]; then
        # Check if user modified the symlinked file (original vs new)
        if diff -q "$original_file" "$new_file" >/dev/null 2>&1; then
            # No difference - user didn't modify or changes are same as upstream
            return 0
        fi
        # User had changes that differ from upstream - conflict!
        return 1
    fi

    # For regular files: check if user modified
    if ! is_config_modified "$config_name" "$user_file"; then
        # Not modified, safe to replace
        return 0
    fi

    # Check if new version is different from original
    if diff -q "$original_file" "$new_file" >/dev/null 2>&1; then
        # No upstream changes, keep user's version
        print_info "  Keeping your changes: $config_name (no upstream changes)"
        return 0
    fi

    # Both modified - try three-way merge
    local merge_output
    merge_output=$(mktemp)

    if diff3 -m "$user_file" "$original_file" "$new_file" > "$merge_output" 2>/dev/null; then
        # Clean merge - no conflicts
        print_success "  Auto-merged: $config_name"
        cp "$merge_output" "$user_file"
        rm -f "$merge_output"
        return 0
    else
        # Conflicts exist
        rm -f "$merge_output"
        return 1
    fi
}

# Handle a single config file conflict
handle_config_conflict() {
    local config_name="$1"
    local user_file="$2"
    local new_file="$3"
    local original_file="$4"  # User's version saved before git reset

    # Use original_file for comparison (user_file may be symlink pointing to new version)
    local user_version="$original_file"
    [ ! -f "$user_version" ] && user_version="$user_file"

    # Use loop instead of recursion to avoid stack depth issues
    while true; do
        echo ""
        # Dynamic box width based on config name length
        local label="Conflict in: $config_name"
        local box_width=${#label}
        [ $box_width -lt 55 ] && box_width=55  # Minimum width
        box_width=$((box_width + 4))  # Add padding
        local border=""
        for ((i=0; i<box_width; i++)); do border+="═"; done
        echo -e "  ${YELLOW}╔${border}╗${NC}"
        printf "  ${YELLOW}║${NC}  Conflict in: ${BOLD}%-$((box_width - 18))s${NC}  ${YELLOW}║${NC}\n" "$config_name"
        echo -e "  ${YELLOW}╚${border}╝${NC}"
        echo ""

        # Show diff summary (your version vs new version)
        echo -e "  ${CYAN}Changes summary (your version → new version):${NC}"
        diff --color=auto -u "$user_version" "$new_file" 2>/dev/null | head -30
        local diff_lines
        diff_lines=$(diff -u "$user_version" "$new_file" 2>/dev/null | wc -l)
        if [ "$diff_lines" -gt 30 ]; then
            echo -e "  ${YELLOW}... ($((diff_lines - 30)) more lines)${NC}"
        fi
        echo ""

        echo -e "  How do you want to resolve this?"
        echo -e "    ${GREEN}1)${NC} Keep my version (ignore upstream changes)"
        echo -e "    ${CYAN}2)${NC} Use new version (discard my changes)"
        echo -e "    ${YELLOW}3)${NC} Backup mine and use new version"
        echo -e "    ${BLUE}4)${NC} Show full diff"
        echo ""
        echo -n "  Choice [1-4]: "
        read -r choice </dev/tty || choice="1"

        case "$choice" in
            1)
                # Restore user's version (convert symlink to regular file if needed)
                rm -f "$user_file" 2>/dev/null
                cp "$user_version" "$user_file"
                print_step "  Keeping your version of $config_name"
                return 0
                ;;
            2)
                rm -f "$user_file" 2>/dev/null
                cp "$new_file" "$user_file"
                print_success "  Updated to new version: $config_name"
                return 0
                ;;
            3)
                local backup_file="${user_file}.backup-$(date +%Y%m%d%H%M%S)"
                cp "$user_version" "$backup_file"
                rm -f "$user_file" 2>/dev/null
                cp "$new_file" "$user_file"
                print_success "  Updated $config_name (backup: $backup_file)"
                return 0
                ;;
            4)
                diff --color=auto -u "$user_version" "$new_file" 2>/dev/null | less
                # Loop continues to show menu again
                ;;
            *)
                # Restore user's version
                rm -f "$user_file" 2>/dev/null
                cp "$user_version" "$user_file"
                print_step "  Keeping your version of $config_name"
                return 0
                ;;
        esac
    done
}

# Run migration scripts sequentially
run_migrations() {
    local from_version="$1"
    local to_version="$2"

    # Validate inputs (backward compatibility for old installations)
    # Ensure versions are valid non-negative integers
    if [ -z "$from_version" ] || ! [[ "$from_version" =~ ^[0-9]+$ ]]; then
        from_version=0
    fi
    if [ -z "$to_version" ] || ! [[ "$to_version" =~ ^[0-9]+$ ]]; then
        to_version=1
    fi

    print_step "Checking for migrations from v$from_version to v$to_version..."

    local migrations_dir="$INSTALL_DIR/migrations"
    local ran_any=false

    if [ ! -d "$migrations_dir" ]; then
        print_info "No migrations directory found"
        return 0
    fi

    # Skip if versions are equal or invalid
    if [ "$from_version" -ge "$to_version" ]; then
        print_info "No migrations needed"
        return 0
    fi

    # Run migrations sequentially
    for v in $(seq $((from_version + 1)) "$to_version"); do
        local prev=$((v - 1))
        local migration_script="$migrations_dir/v${prev}_to_v${v}.sh"

        if [ -f "$migration_script" ]; then
            print_step "Running migration: v$prev → v$v"

            # Source and run the migration
            if source "$migration_script" && run_migration; then
                print_success "Migration v$prev → v$v complete"
                ran_any=true
            else
                print_error "Migration v$prev → v$v failed!"
                echo -e "Do you want to ${YELLOW}continue anyway${NC}? (y/n)"
                read -r cont </dev/tty || cont="n"
                if [[ ! "$cont" =~ ^[Yy](es)?$ ]]; then
                    return 1
                fi
            fi
        fi
    done

    if [ "$ran_any" = false ]; then
        print_info "No migrations needed"
    fi

    return 0
}

# Update packages based on PACKAGES manifest
update_packages() {
    local packages_file="$INSTALL_DIR/PACKAGES"

    if [ ! -f "$packages_file" ]; then
        print_warning "No package manifest found"
        return 0
    fi

    print_step "Checking for new packages..."

    # Parse PACKAGES file and get current version requirements
    local current_version
    current_version=$(get_latest_version)

    local new_packages=()
    local removed_packages=()

    while IFS=: read -r pkg min_ver max_ver || [ -n "$pkg" ]; do
        # Skip comments and empty lines
        [[ "$pkg" =~ ^#.*$ ]] && continue
        [[ -z "$pkg" ]] && continue

        # Clean up whitespace
        pkg=$(echo "$pkg" | tr -d '[:space:]')
        min_ver=$(echo "$min_ver" | tr -d '[:space:]')
        max_ver=$(echo "$max_ver" | tr -d '[:space:]')

        [ -z "$pkg" ] && continue

        # Check if package is required for current version
        local required=true

        # If min_version set and current < min, not required yet
        if [ -n "$min_ver" ] && [ "$current_version" -lt "$min_ver" ]; then
            required=false
        fi

        # If max_version set and current > max, no longer required
        if [ -n "$max_ver" ] && [ "$current_version" -gt "$max_ver" ]; then
            required=false
            # Check if it's installed - might need removal
            if rpm -q "$pkg" &>/dev/null; then
                removed_packages+=("$pkg")
            fi
        fi

        # If required and not installed, add to new packages
        if [ "$required" = true ] && ! rpm -q "$pkg" &>/dev/null; then
            new_packages+=("$pkg")
        fi
    done < "$packages_file"

    # Install new packages
    if [ ${#new_packages[@]} -gt 0 ]; then
        echo ""
        print_info "New packages needed:"
        for pkg in "${new_packages[@]}"; do
            echo -e "  ${GREEN}+${NC} $pkg"
        done
        echo ""

        echo -e "Do you want to ${GREEN}install these packages${NC}? (y/n)"
        read -r install_confirm </dev/tty || install_confirm="n"

        if [[ "$install_confirm" =~ ^[Yy](es)?$ ]]; then
            print_step "Installing new packages..."
            if sudo dnf install -y "${new_packages[@]}" --skip-unavailable; then
                print_success "New packages installed"
            else
                print_warning "Some packages may have failed to install"
            fi
        fi
    else
        print_success "All required packages are installed"
    fi

    # Handle removed packages
    if [ ${#removed_packages[@]} -gt 0 ]; then
        echo ""
        print_info "Packages no longer required:"
        for pkg in "${removed_packages[@]}"; do
            echo -e "  ${RED}-${NC} $pkg"
        done
        echo ""

        echo -e "Do you want to ${RED}remove these packages${NC}? (y/n)"
        read -r remove_confirm </dev/tty || remove_confirm="n"

        if [[ "$remove_confirm" =~ ^[Yy](es)?$ ]]; then
            print_step "Removing old packages..."
            sudo dnf remove -y --noautoremove "${removed_packages[@]}" 2>/dev/null || true
            print_success "Old packages removed"
        fi
    fi
}

# Discover all config files in the repo's configs/ directory
# Returns relative paths (e.g., "i3/config", "fish/config.fish")
# This ensures new configs added in future versions are automatically handled
get_all_config_files() {
    local configs_dir="$INSTALL_DIR/configs"
    if [ ! -d "$configs_dir" ]; then
        return
    fi
    find "$configs_dir" -type f -printf '%P\n' 2>/dev/null | sort
}

# Map a config name (relative path) to the user's file path
# Most configs live in ~/.config/, with a few exceptions
get_user_config_path() {
    local config_name="$1"
    case "$config_name" in
        xprofile)
            echo "$HOME/.xprofile"
            ;;
        *)
            echo "$HOME/.config/$config_name"
            ;;
    esac
}

# Update config checksums after successful update
# Stores checksums of REPO files so we can detect user modifications from baseline
update_config_checksums() {
    print_step "Updating config checksums..."

    > "$CHECKSUM_FILE"

    local config_name
    while IFS= read -r config_name; do
        [ -z "$config_name" ] && continue
        local repo_file="$INSTALL_DIR/configs/$config_name"
        if [ -f "$repo_file" ]; then
            local sum
            sum=$(md5sum "$repo_file" 2>/dev/null | cut -d' ' -f1)
            echo "$config_name:$sum" >> "$CHECKSUM_FILE"
        fi
    done < <(get_all_config_files)

    print_success "Config checksums updated"
}

# Main draphyOS update function
do_draphyos_update() {
    print_header
    echo -e "  ${BOLD}${MINT}draphyOS Config Update${NC}"
    echo ""

    # Check if draphyOS is installed
    if [ ! -d "$INSTALL_DIR" ]; then
        print_error "draphyOS is not installed."
        print_info "Run the installer first: git clone https://github.com/draphy/draphyOS ~/.draphyOS && ~/.draphyOS/install.sh"
        return 1
    fi

    # Prevent concurrent updates with lock file
    local lock_file="/tmp/draphyOS-update-$UID.lock"
    if [ -f "$lock_file" ]; then
        local lock_pid
        lock_pid=$(cat "$lock_file" 2>/dev/null)
        if [ -n "$lock_pid" ] && kill -0 "$lock_pid" 2>/dev/null; then
            print_error "Another draphyOS update is already running (PID: $lock_pid)"
            return 1
        fi
        # Stale lock, remove it
        rm -f "$lock_file"
    fi
    echo $$ > "$lock_file"
    # Clean up lock on function return and script exit (belt-and-suspenders)
    _cleanup_draphyos_lock() { rm -f "$lock_file" 2>/dev/null || true; }
    trap '_cleanup_draphyos_lock' RETURN
    trap '_cleanup_draphyos_lock; cleanup_sudo' EXIT INT TERM

    # Check for diff3 (required for three-way merge)
    if ! command -v diff3 &>/dev/null; then
        print_warning "diff3 not found. Installing diffutils..."
        if ! sudo dnf install -y diffutils &>/dev/null; then
            print_warning "Could not install diffutils. Three-way merge will be limited."
        fi
    fi

    local installed_version
    installed_version=$(get_draphyos_version)
    # Ensure version is a valid number (backward compatibility)
    if ! [[ "$installed_version" =~ ^[0-9]+$ ]]; then
        installed_version=0
    fi
    print_info "Installed draphyOS version: ${GREEN}$installed_version${NC}"
    echo ""

    # Step 1: Backup current configs
    print_step "Step 1/6: Creating backup..."
    local backup_dir="$HOME/.config-backup-update-$(date +%Y%m%d-%H%M%S)"
    mkdir -p "$backup_dir"

    # Backup all draphyOS-managed configs (auto-discovered from repo)
    local config_name
    while IFS= read -r config_name; do
        [ -z "$config_name" ] && continue
        local user_file
        user_file=$(get_user_config_path "$config_name")
        if [ -f "$user_file" ] && [ ! -L "$user_file" ]; then
            local rel_path="${user_file#$HOME/}"
            mkdir -p "$backup_dir/$(dirname "$rel_path")"
            cp "$user_file" "$backup_dir/$rel_path" 2>/dev/null || true
        fi
    done < <(get_all_config_files)
    print_success "Backup created: $backup_dir"
    echo ""

    # Step 2: Fetch latest from GitHub
    print_step "Step 2/6: Fetching latest draphyOS..."
    if ! cd "$INSTALL_DIR"; then
        print_error "Failed to access install directory: $INSTALL_DIR"
        return 1
    fi

    # Save original configs BEFORE git pull (for three-way merge)
    local original_configs_dir
    original_configs_dir=$(mktemp -d)
    if [ -z "$original_configs_dir" ] || [ ! -d "$original_configs_dir" ]; then
        print_error "Failed to create temp directory for config backup"
        return 1
    fi
    if [ -d "$INSTALL_DIR/configs" ]; then
        cp -r "$INSTALL_DIR/configs"/* "$original_configs_dir/" 2>/dev/null || true
    fi

    if [ -d ".git" ]; then
        # Fetch with timeout to prevent hanging on slow/stalled connections
        if ! timeout 60 git fetch --quiet origin 2>/dev/null; then
            print_error "Failed to fetch updates (timed out or network error)."
            print_info "Check your internet connection and try again."
            rm -rf "$original_configs_dir"
            return 1
        fi

        # Check if there are updates
        local local_head remote_head
        local_head=$(git rev-parse HEAD 2>/dev/null)
        remote_head=$(git rev-parse origin/latest 2>/dev/null || git rev-parse origin/main 2>/dev/null)

        if [ "$local_head" = "$remote_head" ]; then
            print_success "draphyOS is already up to date!"
            echo ""
            echo -e "Do you want to ${CYAN}re-apply configs anyway${NC}? (y/n)"
            read -r reapply </dev/tty || reapply="n"
            if [[ ! "$reapply" =~ ^[Yy](es)?$ ]]; then
                rm -rf "$original_configs_dir"
                return 0
            fi
        fi

        # Pull updates with timeout
        if ! timeout 30 git reset --hard origin/latest --quiet 2>/dev/null; then
            if ! timeout 30 git reset --hard origin/main --quiet 2>/dev/null; then
                print_error "Failed to update repository."
                rm -rf "$original_configs_dir"
                return 1
            fi
        fi
        print_success "Latest draphyOS fetched"
    else
        print_warning "Not a git repository, cannot auto-update"
        rm -rf "$original_configs_dir"
        return 1
    fi
    echo ""

    # Get new version
    local new_version
    new_version=$(get_latest_version)
    # Ensure version is a valid number
    if ! [[ "$new_version" =~ ^[0-9]+$ ]]; then
        new_version=1
    fi
    print_info "New draphyOS version: ${GREEN}$new_version${NC}"
    echo ""

    # Step 3: Run migrations if version changed
    print_step "Step 3/6: Running migrations..."
    if [ "$new_version" -gt "$installed_version" ] 2>/dev/null; then
        if ! run_migrations "$installed_version" "$new_version"; then
            print_error "Migration failed!"
            echo -e "Do you want to ${YELLOW}continue anyway${NC}? (y/n)"
            read -r cont </dev/tty || cont="n"
            if [[ ! "$cont" =~ ^[Yy](es)?$ ]]; then
                return 1
            fi
        fi
    else
        print_info "No version change, skipping migrations"
    fi
    echo ""

    # Step 4: Handle config files with three-way merge
    # Auto-discovers all files from configs/ — no hardcoded list needed
    print_step "Step 4/6: Updating config files..."

    local conflicts=()

    local config_name
    while IFS= read -r config_name; do
        [ -z "$config_name" ] && continue

        local user_file
        user_file=$(get_user_config_path "$config_name")
        local new_file="$INSTALL_DIR/configs/$config_name"
        local original_file="$original_configs_dir/$config_name"

        # Skip if new file doesn't exist
        [ ! -f "$new_file" ] && continue

        # Try three-way merge
        local merge_result
        three_way_merge "$config_name" "$user_file" "$new_file" "$original_file"
        merge_result=$?

        if [ $merge_result -eq 1 ]; then
            # Conflict - add to list for user resolution
            conflicts+=("$config_name")
        elif [ $merge_result -eq 0 ]; then
            if [ -L "$user_file" ] || [ ! -f "$user_file" ]; then
                # Symlink or missing — create/update symlink
                rm -f "$user_file" 2>/dev/null
                mkdir -p "$(dirname "$user_file")"
                ln -sf "$new_file" "$user_file"
                print_success "  Updated: $config_name"
            elif [ -f "$user_file" ] && ! is_config_modified "$config_name" "$user_file"; then
                # Regular file (e.g., cheatsheets) not modified by user — update the copy
                if [ -f "$original_file" ] && ! diff -q "$original_file" "$new_file" >/dev/null 2>&1; then
                    cp "$new_file" "$user_file"
                    print_success "  Updated (copy): $config_name"
                fi
            fi
        fi
    done < <(get_all_config_files)

    # Handle conflicts interactively
    if [ ${#conflicts[@]} -gt 0 ]; then
        echo ""
        print_warning "Found ${#conflicts[@]} config(s) with conflicts:"
        for cfg in "${conflicts[@]}"; do
            echo -e "  - $cfg"
        done
        echo ""

        for config_name in "${conflicts[@]}"; do
            local user_file
            user_file=$(get_user_config_path "$config_name")
            local new_file="$INSTALL_DIR/configs/$config_name"
            local original_file="$original_configs_dir/$config_name"
            handle_config_conflict "$config_name" "$user_file" "$new_file" "$original_file"
        done
    fi
    echo ""

    # Step 5: Update packages
    print_step "Step 5/6: Checking packages..."
    update_packages
    echo ""

    # Step 6: Post-update tasks
    print_step "Step 6/6: Finalizing..."

    # Update version in marker file
    if [ -f "$MARKER_FILE" ]; then
        sed -i "s/^DRAPHYOS_VERSION=.*/DRAPHYOS_VERSION=$new_version/" "$MARKER_FILE"
        # Add if not exists
        if ! grep -q "^DRAPHYOS_VERSION=" "$MARKER_FILE"; then
            echo "DRAPHYOS_VERSION=$new_version" >> "$MARKER_FILE"
        fi
    fi

    # Update config checksums
    update_config_checksums

    # Update the update script itself
    if [ -f "$INSTALL_DIR/update.sh" ]; then
        print_step "Updating system update script..."
        if sudo cp "$INSTALL_DIR/update.sh" /usr/local/bin/update-draphyOS && \
           sudo chmod +x /usr/local/bin/update-draphyOS; then
            print_success "Update script updated"
        fi
    fi

    # Update wallpaper script
    if [ -f "$INSTALL_DIR/scripts/set-wallpaper" ]; then
        print_step "Updating wallpaper script..."
        if sudo cp "$INSTALL_DIR/scripts/set-wallpaper" /usr/local/bin/set-wallpaper && \
           sudo chmod +x /usr/local/bin/set-wallpaper; then
            print_success "Wallpaper script updated"
        fi
    fi

    # Make scripts executable
    chmod +x "$HOME/.config/i3/lock.sh" 2>/dev/null || true
    chmod +x "$HOME/.config/i3/scripts/cheatsheet.sh" 2>/dev/null || true
    chmod +x "$HOME/.config/polybar/launch.sh" 2>/dev/null || true

    # Re-apply default app associations (in case environment.d changed)
    if [ -f "$HOME/.config/environment.d/defaults.conf" ]; then
        print_step "Updating default app associations..."
        # shellcheck source=/dev/null
        . "$HOME/.config/environment.d/defaults.conf"

        if command -v xdg-mime &>/dev/null; then
            # Browser
            for mime in x-scheme-handler/http x-scheme-handler/https x-scheme-handler/about x-scheme-handler/unknown text/html; do
                xdg-mime default "${BROWSER:-brave-browser}.desktop" "$mime" 2>/dev/null || true
            done
            # File manager
            xdg-mime default "${FILE_MANAGER:-thunar}.desktop" inode/directory 2>/dev/null || true
            # Video player (mpv)
            for mime in video/mp4 video/x-matroska video/webm video/x-msvideo video/quicktime video/x-flv video/ogg video/3gpp video/x-m4v; do
                xdg-mime default "${PLAYER:-mpv}.desktop" "$mime" 2>/dev/null || true
            done
            # Audio player (audacious)
            for mime in audio/mpeg audio/mp3 audio/flac audio/ogg audio/wav audio/x-wav audio/aac audio/mp4 audio/x-m4a audio/x-flac audio/x-vorbis+ogg audio/x-opus+ogg audio/x-ape audio/x-wavpack; do
                xdg-mime default "audacious.desktop" "$mime" 2>/dev/null || true
            done
            # Image viewer (feh)
            for mime in image/png image/jpeg image/gif image/webp image/bmp image/tiff image/x-icon image/svg+xml; do
                xdg-mime default "feh.desktop" "$mime" 2>/dev/null || true
            done
            # PDF & e-books (zathura)
            for mime in application/pdf application/epub+zip application/x-mobipocket-ebook; do
                xdg-mime default "org.pwmt.zathura.desktop" "$mime" 2>/dev/null || true
            done
            # Text/code editor (VS Code)
            for mime in text/plain application/json text/x-python text/x-shellscript text/x-script.python application/x-shellscript application/javascript text/javascript text/x-java text/x-c text/x-c++ text/x-rust text/x-go text/markdown text/x-markdown application/xml text/xml application/x-yaml text/yaml; do
                xdg-mime default "${IDE:-code}.desktop" "$mime" 2>/dev/null || true
            done
        fi
        print_success "Default app associations updated"
    fi

    # Cleanup temp directory
    rm -rf "$original_configs_dir"

    echo ""
    echo -e "${MINT}  ╔═══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${MINT}  ║              draphyOS Update Complete!                        ║${NC}"
    echo -e "${MINT}  ╚═══════════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo -e "  Version: ${GREEN}$installed_version${NC} → ${GREEN}$new_version${NC}"
    echo -e "  Backup: ${YELLOW}$backup_dir${NC}"
    echo ""
    echo -e "  ${YELLOW}Please reload i3 (Mod+Shift+r) to apply changes.${NC}"
    echo ""
}

# Enable/disable VM mode (picom xrender backend)
do_vm_mode() {
    local mode="${1:-enable}"
    local picom_config="$HOME/.config/picom/picom.conf"

    print_header
    echo -e "  ${BOLD}${MINT}VM Mode Configuration${NC}"
    echo ""

    if [ ! -f "$picom_config" ] && [ ! -L "$picom_config" ]; then
        print_error "Picom config not found: $picom_config"
        return 1
    fi

    # If it's a symlink, convert to real file for editing
    if [ -L "$picom_config" ]; then
        local target
        target=$(readlink -f "$picom_config")
        if [ ! -f "$target" ]; then
            print_error "Picom config symlink target not found: $target"
            return 1
        fi
        rm -f "$picom_config"
        cp "$target" "$picom_config"
        print_step "Converted picom config from symlink to file"
    fi

    if [ "$mode" = "enable" ]; then
        print_step "Enabling VM mode..."

        # Change backend from glx to xrender
        sed -i 's/^backend = "glx";/backend = "xrender";/' "$picom_config"
        # Disable corner radius
        sed -i 's/^corner-radius = [0-9]*;/corner-radius = 0;/' "$picom_config"
        # Disable vsync
        sed -i 's/^vsync = true;/vsync = false;/' "$picom_config"
        # Disable use-damage
        sed -i 's/^use-damage = true;/use-damage = false;/' "$picom_config"

        # Update marker file
        if [ -f "$MARKER_FILE" ]; then
            if ! grep -q "VM_MODE=" "$MARKER_FILE"; then
                echo "VM_MODE=true" >> "$MARKER_FILE"
            else
                sed -i 's/VM_MODE=.*/VM_MODE=true/' "$MARKER_FILE"
            fi
        fi

        print_success "VM mode enabled (xrender backend, no rounded corners)"
    else
        print_step "Disabling VM mode..."

        # Change backend from xrender to glx
        sed -i 's/^backend = "xrender";/backend = "glx";/' "$picom_config"
        # Enable corner radius
        sed -i 's/^corner-radius = 0;/corner-radius = 12;/' "$picom_config"
        # Enable vsync
        sed -i 's/^vsync = false;/vsync = true;/' "$picom_config"
        # Enable use-damage
        sed -i 's/^use-damage = false;/use-damage = true;/' "$picom_config"

        # Update marker file
        if [ -f "$MARKER_FILE" ]; then
            sed -i 's/VM_MODE=.*/VM_MODE=false/' "$MARKER_FILE"
        fi

        print_success "VM mode disabled (glx backend, rounded corners)"
    fi

    echo ""
    print_info "Restart picom or reload i3 (Mod+Shift+r) to apply changes."
}

# Show help
show_help() {
    echo "Usage: update-draphyOS [OPTION]"
    echo ""
    echo "A safe, interactive system update tool for Fedora and draphyOS."
    echo ""
    echo "Options:"
    echo "  (no option)     Show interactive menu (default)"
    echo "  -u, --update    Update all system packages (DNF + Flatpak)"
    echo "  -s, --security  Install security updates only"
    echo "  -c, --check     Check for updates without installing"
    echo "  -f, --full      Full maintenance (update + clean + Flatpak + autoremove)"
    echo "  -v, --upgrade   Fedora version upgrade"
    echo "  -r, --resume    Resume/cleanup after Fedora upgrade reboot"
    echo "  -d, --draphyos  Update draphyOS configs from GitHub"
    echo "  --vm            Enable VM mode (xrender backend for picom)"
    echo "  --no-vm         Disable VM mode (glx backend for picom)"
    echo "  -h, --help      Show this help message"
    echo ""
    echo "Examples:"
    echo "  update-draphyOS          # Interactive menu"
    echo "  update-draphyOS -u       # Quick system update"
    echo "  update-draphyOS -d       # Update draphyOS configs"
    echo "  update-draphyOS -v       # Fedora version upgrade"
    echo "  update-draphyOS --resume # Cleanup after Fedora upgrade"
    echo "  update-draphyOS --vm     # Enable VM mode"
    echo ""
}

# Interactive menu
show_menu() {
    while true; do
        print_header
        show_system_status

        echo -e "  ${BOLD}Select an option:${NC}"
        echo ""
        echo -e "    ${GREEN}1)${NC}  Update system packages (DNF + Flatpak)"
        echo -e "    ${CYAN}2)${NC}  Security updates only"
        echo -e "    ${BLUE}3)${NC}  Check for updates"
        echo -e "    ${YELLOW}4)${NC}  Full maintenance (update + clean + Flatpak + autoremove)"
        echo -e "    ${MINT}5)${NC}  Fedora version upgrade"
        echo -e "    ${MINT}6)${NC}  Update draphyOS configs"
        echo -e "    ${RED}7)${NC}  Exit"
        echo ""
        echo -n "  Enter choice [1-7]: "
        read -r choice </dev/tty || choice="7"

        echo ""

        case "$choice" in
            1)
                do_update
                echo ""
                echo -e "Press Enter to continue..."
                read -r </dev/tty
                ;;
            2)
                do_security_update
                echo ""
                echo -e "Press Enter to continue..."
                read -r </dev/tty
                ;;
            3)
                check_updates
                echo ""
                echo -e "Press Enter to continue..."
                read -r </dev/tty
                ;;
            4)
                do_full_maintenance
                echo ""
                echo -e "Press Enter to continue..."
                read -r </dev/tty
                ;;
            5)
                do_version_upgrade
                echo ""
                echo -e "Press Enter to continue..."
                read -r </dev/tty
                ;;
            6)
                do_draphyos_update
                echo ""
                echo -e "Press Enter to continue..."
                read -r </dev/tty
                ;;
            7)
                echo -e "${MINT}Goodbye!${NC}"
                exit 0
                ;;
            *)
                print_warning "Invalid option. Please select 1-7."
                sleep 1
                ;;
        esac
    done
}

# Main
main() {
    # Parse command
    case "${1:-}" in
        -h|--help)
            show_help
            exit 0
            ;;
        -u|--update)
            print_header
            check_not_root
            check_sudo
            do_update
            ;;
        -s|--security)
            print_header
            check_not_root
            check_sudo
            do_security_update
            ;;
        -c|--check)
            print_header
            check_updates
            ;;
        -f|--full)
            print_header
            check_not_root
            check_sudo
            do_full_maintenance
            ;;
        -v|--upgrade)
            check_not_root
            check_sudo
            do_version_upgrade
            ;;
        -r|--resume)
            check_not_root
            check_sudo
            if ! check_and_resume_upgrade; then
                print_info "No pending operations to resume."
            fi
            ;;
        -d|--draphyos)
            check_not_root
            check_sudo
            do_draphyos_update
            ;;
        --vm)
            check_not_root
            do_vm_mode "enable"
            ;;
        --no-vm)
            check_not_root
            do_vm_mode "disable"
            ;;
        "")
            # No arguments - show interactive menu
            check_not_root
            check_sudo
            show_menu
            ;;
        *)
            print_error "Unknown option: $1"
            echo ""
            show_help
            exit 1
            ;;
    esac
}

main "$@"