#!/bin/bash
# ==========================================================
# Proxmox VE Update Script — Safe / Non-Invasive Variant
# ==========================================================
# Author      : MacRimi
# Copyright   : (c) 2024 MacRimi
# License     : GPL-3.0
# ==========================================================
# Description:
# Update path intended for a Proxmox host ALREADY in
# production. Unlike scripts/global/update-pve8.sh and
# update-pve9_2.sh (invoked by post_install), this variant preserves
# operator-maintained sources unless an unsubscribed host explicitly chooses
# to switch. Otherwise, the operator's source configuration is preserved:
#
#   - Does NOT silently disable Enterprise / Ceph repositories
#   - Does NOT delete or deduplicate existing repo files
#   - Does NOT overwrite proxmox.sources / debian.sources
#   - Does NOT purge alternative NTP services
#   - Does NOT force-install zfsutils / chrony /
#     proxmox-backup-restore-image
#   - Does NOT write no-firmware-warnings.conf
#
# What it DOES:
#   1. Sanity checks (disk space)
#   2. ensure_repositories() — check subscription and request consent if needed
#   3. apt-get update, with automatic GPG key import when apt
#      reports NO_PUBKEY (any repo, user's or ours)
#   4. Detect pending upgrades + security count
#   5. Confirmation dialog
#   6. apt-get full-upgrade with --force-confdef / --force-confold
#      (never overwrites the operator's edited config files)
#   7. lvm_repair_check() — refreshes VG metadata when disks
#      passed through to guest VMs (DSM, TrueNAS, …) come back
#      with old PV headers
#   8. apt-get autoremove + autoclean
#
# Reboot detection is handled by the caller (utilities/proxmox_update.sh).
# ==========================================================

LOCAL_SCRIPTS="/usr/local/share/proxmenux/scripts"
BASE_DIR="/usr/local/share/proxmenux"
UTILS_FILE="$BASE_DIR/utils.sh"
APT_ENV="env DEBIAN_FRONTEND=noninteractive LC_ALL=C LANG=C"

if [[ -f "$UTILS_FILE" ]]; then
    source "$UTILS_FILE"
fi

load_language
initialize_cache

download_common_functions() {
    if ! source "$LOCAL_SCRIPTS/global/common-functions.sh"; then
        return 1
    fi
}

# ensure_repositories() lives with the install helpers.
source_install_functions() {
    local f="$LOCAL_SCRIPTS/global/utils-install-functions.sh"
    if [[ -f "$f" ]]; then
        source "$f"
    else
        return 1
    fi
}

update_pve_safe() {
    local pve_version
    pve_version=$(pveversion 2>/dev/null | grep -oP 'pve-manager/\K[0-9]+' | head -1)
    if [[ -z "$pve_version" ]]; then
        msg_error "$(translate "Unable to detect Proxmox version")"
        return 1
    fi

    local start_time
    start_time=$(date +%s)
    local log_file="/var/log/proxmox-update-$(date +%Y%m%d-%H%M%S).log"
    {
        echo "=== ProxMenux — Proxmox VE update ==="
        echo "Started: $(date -Iseconds)"
        echo "Host:    $(hostname)"
        echo "Running: $(pveversion 2>/dev/null | head -1)"
    } > "$log_file"
    # Screen capture: replay the pre-upgrade context lines after `clear`
    # so the operator keeps the visual history around the noisy apt run.
    local screen_capture="/tmp/proxmenux_screen_capture_$$.txt"
    : > "$screen_capture"

    if ! download_common_functions || ! source_install_functions; then
        msg_error "$(translate 'Required update helpers unavailable. Update stopped.')"
        rm -f "$screen_capture"
        return 1
    fi

    {
        msg_info2 "$(translate "Detected: Proxmox VE $pve_version — running safe update path")"
    } | tee -a "$screen_capture"

    # ── 1. Sanity checks ──
    local available_space
    available_space=$(df /var/cache/apt/archives | awk 'NR==2 {print int($4/1024)}')
    if [ "$available_space" -lt 1024 ]; then
        msg_error "$(translate "Insufficient disk space. Available: ${available_space}MB")"
        echo -e
        msg_success "$(translate "Press Enter to return to menu...")"
        read -r
        rm -f "$screen_capture"
        return 1
    fi

    # Enterprise hosts and custom mirrors need not reach the public CDN.
    # The configured sources, not that hostname, are checked by APT below.
    if ! declare -F ensure_repositories >/dev/null 2>&1 || ! ensure_repositories; then
        msg_error "$(translate 'Repository check failed. Update stopped.')"
        rm -f "$screen_capture"
        return 1
    fi

    # ── 3. apt-get update with automatic key recovery ──
    local update_output update_exit_code
    update_output=$(apt-get update 2>&1)
    update_exit_code=$?
    {
        echo
        echo "--- apt-get update (exit $update_exit_code) ---"
        printf '%s\n' "$update_output"
    } >> "$log_file"

    if [ $update_exit_code -eq 0 ]; then
        msg_ok "$(translate "Package lists updated successfully")" | tee -a "$screen_capture"
    else
        if echo "$update_output" | grep -Eq "NO_PUBKEY|GPG error"; then
            local key
            key=$(echo "$update_output" | sed -n 's/.*NO_PUBKEY \([0-9A-F]\{8,40\}\).*/\1/p' | head -1)
            if [ -n "$key" ]; then
                mkdir -p /etc/apt/keyrings
                if command -v gpg >/dev/null 2>&1; then
                    if gpg --batch --keyserver keyserver.ubuntu.com --recv-keys "$key" \
                    && gpg --batch --export "$key" | gpg --dearmor -o "/etc/apt/keyrings/${key}.gpg"; then
                        msg_ok "$(translate "Imported missing GPG key: $key")" | tee -a "$screen_capture"
                    else
                        msg_warn "$(translate "Keyrings method failed; trying apt-key fallback")"
                        apt-key adv --keyserver keyserver.ubuntu.com --recv-keys "$key" >/dev/null 2>&1 || true
                    fi
                else
                    msg_warn "$(translate "gpg not found; trying apt-key fallback")"
                    apt-key adv --keyserver keyserver.ubuntu.com --recv-keys "$key" >/dev/null 2>&1 || true
                fi
            fi
            if apt-get update >> "$log_file" 2>&1; then
                msg_ok "$(translate "Package lists updated after GPG fix")" | tee -a "$screen_capture"
            else
                msg_error "$(translate "Failed to update package lists. Check log: $log_file")"
                rm -f "$screen_capture"
                return 1
            fi
        elif echo "$update_output" | grep -Eq "404|Failed to fetch"; then
            msg_warn "$(translate "Some repositories are not available, continuing with available ones...")"
        else
            msg_error "$(translate "Failed to update package lists. Check log: $log_file")"
            echo "Error details: $update_output"
            rm -f "$screen_capture"
            return 1
        fi
    fi

    # No duplicate-source rewrite here: even a seemingly duplicate entry may
    # be operator-maintained. Only the consented switch above edits sources.

    # ── 5-6. Detect + confirm ──
    local current_pve_version available_pve_version upgradable security_updates
    local upgradable_raw upgradable_list
    current_pve_version=$(pveversion 2>/dev/null | grep -oP 'pve-manager/\K[0-9]+\.[0-9]+\.[0-9]+' | head -1)
    available_pve_version=$(apt-cache policy pve-manager 2>/dev/null | grep -oP 'Candidate: \K[0-9]+\.[0-9]+\.[0-9]+' | head -1)
    upgradable_raw=$($APT_ENV apt list --upgradable 2>/dev/null | sed '1d' | sed '/^\s*$/d')
    upgradable=$(printf '%s' "$upgradable_raw" | grep -c . )
    security_updates=$(printf '%s\n' "$upgradable_raw" | grep -ci '\-security')
    upgradable_list=$(printf '%s\n' "$upgradable_raw" | pmx_format_upgradable)

    {
        echo
        echo "--- Packages to upgrade ($upgradable) ---"
        [ "$upgradable" -gt 0 ] && printf '%s\n' "$upgradable_list"
    } >> "$log_file"

    local menu_text
    menu_text="$(translate "System Update Information")\n\n"
    menu_text+="$(translate "Current PVE Version"): $current_pve_version\n"
    if [ -n "$available_pve_version" ] && [ "$available_pve_version" != "$current_pve_version" ]; then
        menu_text+="$(translate "Available PVE Version"): $available_pve_version\n"
    fi
    menu_text+="\n$(translate "Package Updates Available"): $upgradable\n"
    menu_text+="$(translate "Security Updates"): $security_updates\n\n"

    if [ "$upgradable" -eq 0 ]; then
        menu_text+="$(translate "System is already up to date")"
        whiptail --title "$(translate "Update Status")" --msgbox "$menu_text" 15 70
        apt-get -y autoremove >/dev/null 2>&1 || true
        apt-get -y autoclean >/dev/null 2>&1 || true
        echo -e "\nSystem is already up to date." >> "$log_file"
        rm -f "$screen_capture"
        return 0
    fi

    # The package list rides in the same dialog as the summary, so the
    # decision is taken knowing what is about to be replaced. --scrolltext
    # keeps the buttons reachable however long the list is.
    menu_text+="$(translate "Packages to be upgraded"):\n$upgradable_list\n\n"
    menu_text+="$(translate "Do you want to proceed with the system update?")"
    if ! whiptail --title "$(translate "Proxmox Update")" --scrolltext --yesno "$menu_text" 24 78; then
        msg_info2 "$(translate "Update cancelled by user")"
        apt-get -y autoremove >/dev/null 2>&1 || true
        apt-get -y autoclean >/dev/null 2>&1 || true
        rm -f "$screen_capture"
        return 0
    fi

    # ── 7. Full upgrade — --force-confdef/confold preserves user-edited configs ──
    # Redraw the ProxMenux frame before apt starts printing so the operator
    # keeps the visual context around the noisy upgrade output.
    clear
    show_proxmenux_logo
    msg_title "$(translate "$SCRIPT_TITLE")"
    cat "$screen_capture"

    # dpkg's log is read back from here on, so the transaction can be
    # reported package by package without holding apt's output.
    local dpkg_mark
    dpkg_mark=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "\n--- apt full-upgrade ---" >> "$log_file"

    # apt's own progress bar (Progress: [ %]) prints on stderr and only
    # when stdout is a TTY. We pipe stderr through tee to keep a log copy
    # while letting apt keep its interactive stdout, so the native bar
    # keeps rendering at the bottom of the terminal as the user expects.
    DEBIAN_FRONTEND=noninteractive apt -y \
        -o Dpkg::Options::='--force-confdef' \
        -o Dpkg::Options::='--force-confold' \
        full-upgrade 2> >(tee -a "$log_file" >&2)
    local upgrade_exit_code=$?
    echo -e

    {
        echo
        echo "--- Packages changed (exit $upgrade_exit_code) ---"
        pmx_dpkg_changes_since "$dpkg_mark"
    } >> "$log_file"

    # Redraw once more so the wrap-up (LVM check, cleanup, summary) reads
    # cleanly instead of scrolling under half-a-screen of apt noise.
    clear
    show_proxmenux_logo
    msg_title "$(translate "$SCRIPT_TITLE")"
    cat "$screen_capture"

    if [ $upgrade_exit_code -ne 0 ]; then
        msg_error "$(translate "System upgrade failed. Check log: $log_file")"
        echo "Finished: $(date -Iseconds) — upgrade failed (exit $upgrade_exit_code)" >> "$log_file"
        rm -f "$screen_capture"
        return 1
    fi

    msg_ok "$(translate "System upgrade completed")"

    # ── 8. LVM header repair (only touches VGs actually flagged as stale) ──
    if declare -f lvm_repair_check >/dev/null 2>&1; then
        lvm_repair_check
    fi

    # ── 9. DKMS driver rebuild if a new kernel was staged ──
    if declare -f pmx_rebuild_dkms_after_kernel >/dev/null 2>&1; then
        pmx_rebuild_dkms_after_kernel
    fi

    # ── 10. Final cleanup ──
    msg_info "$(translate "Running cleanup")"
    apt-get -y autoremove >/dev/null 2>&1 || true
    apt-get -y autoclean >/dev/null 2>&1 || true
    msg_ok "$(translate "Cleanup finished")"

    local end_time duration minutes seconds
    end_time=$(date +%s)
    duration=$((end_time - start_time))
    minutes=$((duration / 60))
    seconds=$((duration % 60))

    echo -e "${TAB}${BGN}$(translate "====== PVE UPDATE COMPLETED ======")${CL}"
    echo -e "${TAB}${GN}⏱️  $(translate "Duration")${CL}: ${BL}${minutes}m ${seconds}s${CL}"
    echo -e "${TAB}${GN}📄 $(translate "Log file")${CL}: ${BL}$log_file${CL}"
    echo -e "${TAB}${GN}📦 $(translate "Packages upgraded")${CL}: ${BL}$upgradable${CL}"
    echo -e "${TAB}${GN}🖥️  $(translate "Proxmox VE")${CL}: ${BL}${available_pve_version:-$current_pve_version}${CL}"

    {
        echo
        echo "Finished: $(date -Iseconds) — ${minutes}m ${seconds}s"
        echo "Running: $(pveversion 2>/dev/null | head -1)"
    } >> "$log_file"

    msg_ok "$(translate "Proxmox VE safe update completed")"
    rm -f "$screen_capture"
}

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    update_pve_safe
fi