#!/bin/bash

# Interactive system update script for apt, flatpak, and brew

set -euo pipefail

# Color definitions
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Add help option
if [[ "${1:-}" =~ ^-+h ]]; then
    echo "Usage: $0"
    echo "Interactive system update script for apt, flatpak, and brew"
    exit 0
fi

# Handle SIGINT (Ctrl+C)
trap 'echo; echo "${YELLOW}[WARN]${NC} Script interrupted by user."; exit 130' INT

# Enhanced logging with timestamps, colors and severity levels
log() {
    local level="${1:-INFO}"
    shift
    local color=""

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

    echo -e "[$(date '+%Y-%m-%d %H:%M:%S')] [${color}$level${NC}] $*" | tee -a /tmp/update_script.log >&2
}

# Request sudo upfront
if ! sudo -v; then
    log "ERROR" "sudo access required to run this script."
    exit 1
fi

echo

# Track actions
updates_performed=false
updates_failed=false
dkms_ran=false
dkms_success=false
reboot_required=false


# Ensure required commands are available
ensure_command() {
    if ! command -v "$1" > /dev/null; then
        log "WARN" "$1 is required but not installed."
        return 1
    fi
}

# Update APT packages interactively
if ensure_command apt; then
    log "INFO" "Updating APT package list..."
    sudo apt update

    log "INFO" "Checking for upgrades..."

    # Skip header line from apt list
    if apt list --upgradable 2>/dev/null | tail -n +2 | grep -q "."; then
        echo
        echo "The following packages can be upgraded:"
        apt list --upgradable
        echo
        read -p "Proceed with interactive upgrade? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            if sudo apt full-upgrade; then
                updates_performed=true
            else
                updates_failed=true
            fi
        else
            log "WARN" "APT upgrade skipped."
        fi
    else
        log "INFO" "No APT packages to upgrade."
    fi

    # Clean up
    sudo apt autoremove
fi

# Update Flatpak interactively
if ensure_command flatpak; then
    log "INFO" "Checking for Flatpak updates..."

    if flatpak remote-ls --updates > /dev/null 2>&1 && [ $(flatpak remote-ls --updates 2>/dev/null | wc -l) -gt 0 ]; then
        echo
        read -p "Proceed with Flatpak update? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            if flatpak update && flatpak uninstall --unused; then
                updates_performed=true
            else
                updates_failed=true
            fi
        else
            log "WARN" "Flatpak update skipped."
        fi
    else
        log "INFO" "No Flatpak updates available."
    fi
fi

# Update Homebrew interactively
if command -v brew > /dev/null; then
    log "INFO" "Updating Homebrew..."
    brew update > /dev/null

    # Check outdated formulae
    if brew outdated | grep -q "."; then
        echo
        echo "Homebrew formulae to upgrade:"
        brew outdated
        echo
        read -p "Upgrade all outdated formulae? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            if brew upgrade; then
                updates_performed=true
            else
                updates_failed=true
            fi
        else
           log "WARN" "Brew formulae upgrade skipped."
        fi
    else
        log "INFO" "No outdated Homebrew formulae."
    fi

    # Cleanup
    brew cleanup
else
    log "WARN" "Homebrew not installed. Skipping."
fi

# Handle DKMS (non-interactive)
if command -v dkms > /dev/null; then
    log "INFO" "Rebuilding DKMS modules..."
    # dkms autoinstall will handle missing headers and failed rebuilds
    if sudo dkms autoinstall; then
        log "INFO" "DKMS modules rebuilt successfully."
    else
        log "WARN" "DKMS module rebuild failed."
    fi
fi

# Reboot warning
if [ -f /var/run/reboot-required ]; then
    log "WARN" "System reboot is required."
    cat /var/run/reboot-required.pkgs | xargs echo "Packages requiring reboot:"
    reboot_required=true
fi

# Final status output
echo
echo "==============================="
log "INFO" "Update process completed."

if [[ "$updates_performed" == "true" ]]; then
    log "INFO" "Package updates were applied successfully."
elif [[ "$updates_failed" == "true" ]]; then
    log "WARN" "Some updates failed to apply."
else
    log "INFO" "No package updates were applied."
fi

if [[ "$dkms_ran" == "true" ]]; then
    if [[ "$dkms_success" == "true" ]]; then
        log "INFO" "DKMS modules rebuilt successfully."
    else
        log "WARN" "DKMS module rebuild failed."
    fi
elif ! command -v dkms > /dev/null; then
    log "INFO" "DKMS not installed."
else
    log "INFO" "No DKMS modules to rebuild."
fi

if [[ "$reboot_required" == "true" ]]; then
    log "WARN" "REBOOT REQUIRED"
else
    log "INFO" "System does not require reboot."
fi
echo "==============================="

echo

# Return appropriate exit code
if [[ "$reboot_required" == "true" ]]; then
    exit 2
elif [[ "$updates_failed" == "true" ]]; then
    exit 1
else
    exit 0
fi