#!/usr/bin/env zsh

# Exit code 0 = Reboot required
# Exit code 1 = System is up to date / No reboot needed

# 1. Branch based on detected distro
if [[ -f /etc/os-release ]]; then
    source /etc/os-release
else
    echo "Cannot determine OS."
    exit 1
fi

# ==========================================
# DEBIAN LOGIC
# ==========================================
if [[ "$ID" == "debian" || "$ID_LIKE" == *"debian"* ]]; then
    # Debian automatically creates this file when a kernel, microcode, 
    # or core library update requires a reboot.
    if [[ -f /var/run/reboot-required ]]; then
        exit 0
    else
        exit 1
    fi

# ==========================================
# ARCH LINUX LOGIC
# ==========================================
elif [[ "$ID" == "arch" || "$ID_LIKE" == *"arch"* ]]; then
    
    # --- KERNEL CHECK ---
    # Get the currently running kernel version
    current_kernel=$(uname -r)
    
    # Determine the installed kernel package based on the running kernel's name
    # e.g., "6.1.60-1-lts" -> "linux-lts", "6.5.9-arch1-1" -> "linux"
    if [[ "$current_kernel" == *"-lts" ]]; then
        kernel_pkg="linux-lts"
    elif [[ "$current_kernel" == *"-zen" ]]; then
        kernel_pkg="linux-zen"
    elif [[ "$current_kernel" == *"-hardened" ]]; then
        kernel_pkg="linux-hardened"
    else
        kernel_pkg="linux"
    fi

    # Get the installed version from pacman.
    # Note: pacman uses periods (6.5.9.arch1-1) while uname uses dashes (6.5.9-arch1-1).
    # We strip the package name and format it to match `uname -r` for a clean comparison.
    installed_kernel=$(pacman -Q "$kernel_pkg" | awk '{print $2}' | sed 's/\./-/g' | sed 's/-\([^-]*\)$/.\1/')

    # In some UKI setups or specific Arch versions, the sed replacement above might not 
    # perfectly match `uname -r` format. A bulletproof fallback is checking if the 
    # modules directory for the currently running kernel has been removed by a pacman update.
    if [[ ! -d "/usr/lib/modules/$current_kernel" ]]; then
        exit 0
    fi

    # --- MICROCODE CHECK ---
    # Check CPU vendor to determine which microcode package to look for
    vendor=$(grep -m 1 'vendor_id' /proc/cpuinfo | awk '{print $3}')
    if [[ "$vendor" == "GenuineIntel" ]]; then
        ucode_pkg="intel-ucode"
    elif [[ "$vendor" == "AuthenticAMD" ]]; then
        ucode_pkg="amd-ucode"
    else
        ucode_pkg=""
    fi

    if [[ -n "$ucode_pkg" ]]; then
        # Check if the microcode package was updated recently by looking at the 
        # build date of the installed package vs the boot time of the system.
        
        # Get system boot time in seconds since epoch
        boot_time=$(awk '{print int($1)}' /proc/uptime)
        boot_epoch=$(( $(date +%s) - boot_time ))
        
        # Get the install time of the microcode package
        ucode_install_epoch=$(expac '%time' "$ucode_pkg" 2>/dev/null || echo 0)

        # If the microcode was installed AFTER the system booted, reboot needed
        if (( ucode_install_epoch > boot_epoch )); then
            exit 0
        fi
    fi

    # If we made it this far, neither kernel nor microcode triggered a reboot
    exit 1

else
    # Fallback for unsupported OS
    exit 1
fi