#!/usr/bin/env bash
set -euo pipefail

CURRENT_VERSION="0.3.0"
STATUS_FILE="/var/cache/fedora-update/status.json"

BOLD="\033[1m"
BLUE="\033[1;34m"
GREEN="\033[1;32m"
YELLOW="\033[1;33m"
CYAN="\033[1;36m"
RESET="\033[0m"

# ------------------------------------------------------------------------------
# I18N INITIALIZATION
# ------------------------------------------------------------------------------
I18N_PY="/usr/libexec/fedora-update/fedora_update_i18n.py"
if [ ! -f "$I18N_PY" ]; then
    I18N_PY="$(dirname "$(readlink -f "$0")")/../lib/fedora_update_i18n.py"
fi

_() {
    python3 "$I18N_PY" "$@"
}

cleanup_on_exit() {
    echo ""
    read -r -p "$(_ PRESS_ENTER_EXIT)" || true
}
trap cleanup_on_exit EXIT

# ------------------------------------------------------------------------------
# MAIN EXECUTION FLOW
# ------------------------------------------------------------------------------
main() {
    echo -e "${CYAN}${BOLD}$(_ CLI_TITLE version="$CURRENT_VERSION")${RESET}\n"

    # Refresh status cache if missing
    if [ ! -f "$STATUS_FILE" ]; then
        echo -e "${CYAN}$(_ CHECKING_FOR_UPDATES)${RESET}"
        /usr/libexec/fedora-update/check_backend.sh 2>/dev/null || true
    fi

    # Display pending updates if status file is present
    if [ -f "$STATUS_FILE" ]; then
        echo -e "${CYAN}${BOLD}$(_ CLI_UPDATES_HEADER)${RESET}"
        python3 -c '
import json, sys, os
sys.path.insert(0, "/usr/libexec/fedora-update")
dev_lib = os.path.abspath(os.path.join(os.path.dirname(os.path.realpath("'"$0"'")), "../lib"))
sys.path.insert(0, dev_lib)
try:
    from fedora_update_i18n import _
except ImportError:
    def _(k, **kw): return k

status_file = "'"$STATUS_FILE"'"
if os.path.exists(status_file):
    try:
        with open(status_file, "r", encoding="utf-8") as f:
            d = json.load(f)
            sources = [
                (_("CAT_DNF"), d.get("dnf_updates", [])),
                (_("CAT_FLATPAK"), d.get("flatpak_updates", []))
            ]
            has_updates = False
            for name, pkgs in sources:
                if pkgs:
                    has_updates = True
                    print(f"\033[1;32m{name}:\033[0m")
                    for pkg in pkgs:
                        print(f"  • {pkg}")
            if not has_updates:
                print(_("NO_PENDING_UPDATES"))
    except Exception:
        pass
'
        echo ""
    fi

    # Overall confirmation prompt
    read -r -p "$(_ PROMPT_CONTINUE_UPDATES)" REPLY || REPLY="Y"
    REPLY=${REPLY:-Y}
    if [[ ! "$REPLY" =~ ^[YyJjSsOo]$ ]]; then
        echo -e "\n${YELLOW}$(_ CANCELLED_BY_USER)${RESET}\n"
        exit 0
    fi

    echo -e "\n${CYAN}$(_ VALIDATING_ADMIN)${RESET}"
    sudo -v

    # Snapshot installed packages and kernels before upgrades for accurate delta comparison
    KERNELS_BEFORE=$(rpm -q --qf "%{VERSION}-%{RELEASE}.%{ARCH}\n" kernel-core 2>/dev/null | sort -u || true)
    NVIDIA_BEFORE=$(rpm -q --qf "%{NAME}-%{VERSION}-%{RELEASE}\n" akmod-nvidia xorg-x11-drv-nvidia nvidia-driver 2>/dev/null | sort -u || true)
    CRITICAL_BEFORE=$(rpm -q --qf "%{NAME}-%{VERSION}-%{RELEASE}\n" glibc systemd-udev kwin plasma-workspace 2>/dev/null | sort -u || true)

    # --------------------------------------------------------------------------
    # 1. DNF UPGRADE STEP (Direct TTY attachment preserves live download rates)
    # --------------------------------------------------------------------------
    read -r -p "$(_ PROMPT_UPGRADE_DNF)" DNF_REPLY || DNF_REPLY="Y"
    DNF_REPLY=${DNF_REPLY:-Y}

    if [[ "$DNF_REPLY" =~ ^[YyJjSsOo]$ ]]; then
        echo -e "\n${BLUE}${BOLD}$(_ UPGRADING_DNF)${RESET}"
        if command -v dnf5 >/dev/null 2>&1; then
            sudo dnf5 upgrade --refresh
        else
            sudo dnf upgrade --refresh
        fi
    fi

    # --------------------------------------------------------------------------
    # 2. FLATPAK UPGRADE STEP (Step-by-step confirmation)
    # --------------------------------------------------------------------------
    if command -v flatpak >/dev/null 2>&1; then
        echo -e ""
        read -r -p "$(_ PROMPT_UPGRADE_FLATPAK)" FLATPAK_REPLY || FLATPAK_REPLY="Y"
        FLATPAK_REPLY=${FLATPAK_REPLY:-Y}
        if [[ "$FLATPAK_REPLY" =~ ^[YyJjSsOo]$ ]]; then
            echo -e "\n${BLUE}${BOLD}$(_ UPGRADING_FLATPAK)${RESET}"
            flatpak update
        fi
    fi

    # --------------------------------------------------------------------------
    # 3. ENVIRONMENT DISCOVERY & HARDWARE DETECTION
    # --------------------------------------------------------------------------
    NVIDIA_PRESENT=0
    if lspci -k 2>/dev/null | grep -i "vga\|3d" | grep -iq "nvidia" || lsmod 2>/dev/null | grep -iq "nvidia"; then
        NVIDIA_PRESENT=1
    fi

    SB_ACTIVE=0
    if command -v mokutil >/dev/null 2>&1 && mokutil --sb-state 2>/dev/null | grep -qi "enabled"; then
        SB_ACTIVE=1
    fi

    CUSTOM_KERNEL=0
    KERNEL_VENDOR=$(rpm -q --qf "%{VENDOR}\n" kernel-core 2>/dev/null | head -n1 || true)
    if [ -z "$KERNEL_VENDOR" ] || ! echo "$KERNEL_VENDOR" | grep -Ei -q "(Fedora|Red Hat)"; then
        CUSTOM_KERNEL=1
    fi

    # --------------------------------------------------------------------------
    # 4. WATERPROOF SECURE BOOT & MOK VERIFICATION
    # --------------------------------------------------------------------------
    if [ "$SB_ACTIVE" -eq 1 ] && { [ "$NVIDIA_PRESENT" -eq 1 ] || [ "$CUSTOM_KERNEL" -eq 1 ]; }; then
        MOK_PUB="/etc/pki/akmods/certs/public_key.der"
        MOK_PRIV="/etc/pki/akmods/private/private_key.priv"
        MOK_READY=0
        MOK_ENROLLED_LIST=$(mokutil --list-enrolled 2>/dev/null || true)

        # Method 1: Check if any installed or loaded NVIDIA module is signed with an enrolled MOK key
        for mod in $(find /lib/modules/ -name "nvidia.ko*" 2>/dev/null); do
            s=$(modinfo -F signer "$mod" 2>/dev/null || true)
            if [ -n "$s" ] && echo "$MOK_ENROLLED_LIST" | grep -qi "$s"; then
                MOK_READY=1
                break
            fi
        done

        # Method 2: Check if standard akmods keypair exists on disk and its fingerprint is in MOK
        if [ "$MOK_READY" -eq 0 ]; then
            if sudo test -f "$MOK_PUB" && sudo test -f "$MOK_PRIV"; then
                if command -v openssl >/dev/null 2>&1; then
                    PUB_FP=$(sudo openssl x509 -in "$MOK_PUB" -inform DER -fingerprint -noout 2>/dev/null | cut -d= -f2 | tr -d ':' | tr '[:upper:]' '[:lower:]' || true)
                    MOK_FPS=$(echo "$MOK_ENROLLED_LIST" | grep -i "Fingerprint" | awk '{print $NF}' | tr -d ':' | tr '[:upper:]' '[:lower:]' || true)
                    if [ -n "$PUB_FP" ] && echo "$MOK_FPS" | grep -qi "$PUB_FP"; then
                        MOK_READY=1
                    fi
                else
                    MOK_READY=1
                fi
            fi
        fi

        # Prompt ONLY if signing key is truly missing or unenrolled
        if [ "$MOK_READY" -eq 0 ]; then
            echo -e "\n${YELLOW}${BOLD}$(_ MOK_REQUIRED_TITLE)${RESET}"
            read -r -p "$(_ PROMPT_MOK_CREATE)" MOK_REPLY || MOK_REPLY="Y"
            MOK_REPLY=${MOK_REPLY:-Y}
            if [[ "$MOK_REPLY" =~ ^[YyJjSsOo]$ ]]; then
                if ! sudo test -f "$MOK_PUB" && command -v kmodgenca >/dev/null 2>&1; then
                    sudo kmodgenca -a || true
                fi
                if sudo test -f "$MOK_PUB"; then
                    sudo mokutil --import "$MOK_PUB" || true
                    echo -e "${CYAN}--------------------------------------------------${RESET}"
                    echo -e "${YELLOW}${BOLD}$(_ MOK_INSTRUCTIONS_HEADER)${RESET}"
                    echo -e "${YELLOW}$(_ MOK_INSTRUCTIONS_BODY)${RESET}"
                    echo -e "${CYAN}--------------------------------------------------${RESET}"
                fi
            fi
        fi
    fi

    # --------------------------------------------------------------------------
    # 5. DIFFERENTIAL NVIDIA MODULE & INITRAMFS BAKING
    # --------------------------------------------------------------------------
    KERNELS_AFTER=$(rpm -q --qf "%{VERSION}-%{RELEASE}.%{ARCH}\n" kernel-core 2>/dev/null | sort -u || true)
    NVIDIA_AFTER=$(rpm -q --qf "%{NAME}-%{VERSION}-%{RELEASE}\n" akmod-nvidia xorg-x11-drv-nvidia nvidia-driver 2>/dev/null | sort -u || true)
    CRITICAL_AFTER=$(rpm -q --qf "%{NAME}-%{VERSION}-%{RELEASE}\n" glibc systemd-udev kwin plasma-workspace 2>/dev/null | sort -u || true)

    NEW_KERNELS=$(comm -13 <(echo "$KERNELS_BEFORE") <(echo "$KERNELS_AFTER") 2>/dev/null | grep -v '^$' || true)
    NVIDIA_UPDATED=$(comm -3 <(echo "$NVIDIA_BEFORE") <(echo "$NVIDIA_AFTER") 2>/dev/null | grep -v '^$' || true)
    CRITICAL_UPDATED=$(comm -3 <(echo "$CRITICAL_BEFORE") <(echo "$CRITICAL_AFTER") 2>/dev/null | grep -v '^$' || true)

    if [ "$NVIDIA_PRESENT" -eq 1 ]; then
        if [ -n "$NVIDIA_UPDATED" ]; then
            # Driver updated -> Bake into ALL installed kernels (including any new kernel)
            echo -e "\n${CYAN}$(_ BAKING_NVIDIA_MODULES)${RESET}"
            sudo akmods --force --kernels all >/dev/null 2>&1 || true
            sudo dracut --regenerate-all --force >/dev/null 2>&1 || true
            echo -e "${GREEN}$(_ BAKED_NVIDIA_MODULES_SUCCESS)${RESET}"
        elif [ -n "$NEW_KERNELS" ]; then
            # Kernel updated -> Bake ONLY into newly installed kernel(s)
            echo -e "\n${CYAN}$(_ BAKING_NVIDIA_MODULES)${RESET}"
            for k in $NEW_KERNELS; do
                [ -z "$k" ] && continue
                sudo akmods --force --kernel "$k" >/dev/null 2>&1 || true
                sudo depmod -a "$k" >/dev/null 2>&1 || true
                sudo dracut --force --kver "$k" >/dev/null 2>&1 || true
            done
            echo -e "${GREEN}$(_ BAKED_NVIDIA_MODULES_SUCCESS)${RESET}"
        fi
    fi

    # --------------------------------------------------------------------------
    # 6. REFRESH CACHE & REBOOT EVALUATION
    # --------------------------------------------------------------------------
    NEEDS_REBOOT=0
    if [ -n "$NEW_KERNELS" ] || [ -n "$NVIDIA_UPDATED" ] || [ -n "$CRITICAL_UPDATED" ] || [ -f /run/reboot-required ]; then
        NEEDS_REBOOT=1
    fi

    echo -e "\n${CYAN}$(_ REFRESHING_CACHE)${RESET}"
    sudo systemctl start --no-block fedora-update-check.service 2>/dev/null || true

    echo -e "\n${GREEN}${BOLD}$(_ UPGRADE_COMPLETE_SYSTEM)${RESET}\n"

    if [ "$NEEDS_REBOOT" -eq 1 ]; then
        read -r -p "$(_ PROMPT_REBOOT_NOW)" REBOOT_CHOICE || REBOOT_CHOICE="Y"
        REBOOT_CHOICE=${REBOOT_CHOICE:-Y}
        if [[ "$REBOOT_CHOICE" =~ ^[YyJjSsOo]$ ]]; then
            echo -e "\n${YELLOW}$(_ REBOOTING_IN_5S)${RESET}"
            sleep 5
            sudo reboot
        fi
    else
        echo -e "${GREEN}$(_ NO_REBOOT_REQUIRED)${RESET}"
    fi

    exit 0
}

main "$@"
