#!/bin/sh
# freebsd-automerge.sh
#
# Unified script — manages the complete FreeBSD upgrade workflow.
# Operates in two modes:
#
#   MAIN MODE (normal usage):
#     sh freebsd-automerge.sh [major|release]
#     Examples:
#       sh freebsd-automerge.sh 15            # auto-detect latest 15.x-RELEASE
#       sh freebsd-automerge.sh 15.1-RELEASE
#       sh freebsd-automerge.sh               # merge conflicts after manual upgrade
#
#   EDITOR MODE (invoked automatically by freebsd-update):
#     freebsd-update calls this script as EDITOR when it finds a conflicted
#     file. The script performs an automatic 3-way merge instead of opening
#     nano. Falls back to nano if conflicts remain. Never invoke manually.
#
#   Additional option (any mode):
#   --auto-confirm
#       Automatically answer 'y' to all interactive freebsd-update prompts.
#       Requires: pkg install expect
#       Example: sh freebsd-automerge.sh 15 --auto-confirm
#
# FULL CYCLE (major upgrade e.g. 14.x -> 15.x):
#   1. Find latest RELEASE for target branch
#   2. Download base.txz OLD and NEW for 3-way merge
#   3. EDITOR=<this script> freebsd-update -r 15.x-RELEASE upgrade
#      -> conflicts resolved automatically during upgrade
#   4. Resolve any remaining <<<<<<< markers in /etc
#   5. freebsd-update install (kernel) + reboot notice
#   6. [after reboot] freebsd-update install (userland) via rc.d + notice
#
# DEPENDENCIES: merge(1) — included in FreeBSD base system

# ---------------------------------------------------------------------------
# DISPATCH: detect which mode we were invoked in
# ---------------------------------------------------------------------------
# freebsd-update calls EDITOR with a path containing "//"
# e.g.: /var/db/freebsd-update/merge/new//etc/blacklistd.conf
# If $1 contains "//" -> EDITOR mode, otherwise -> MAIN mode.

case "$1" in
    *//*) _MODE="EDITOR" ;;
    *)    _MODE="MAIN"   ;;
esac

# ---------------------------------------------------------------------------
# Shared constants
# ---------------------------------------------------------------------------

STATE_DIR="/var/db/automerge"
STATE_FILE="${STATE_DIR}/state"
RCD_SCRIPT="/etc/rc.d/automerge_resume"
WORKDIR="/tmp/automerge"
ARCHIVES="base.txz src.txz"

# ---------------------------------------------------------------------------
# ██████████████████████████████████████████████████████████████████████████
# EDITOR MODE
# Called by freebsd-update as: <script> <fbu_filepath>
# ██████████████████████████████████████████████████████████████████████████
# ---------------------------------------------------------------------------

editor_mode() {
    FBUFILE="$1"
    LOGFILE="${WORKDIR}/automerge_editor.log"
    OLD_EXTRACTDIR="${WORKDIR}/editor_old"
    NEW_EXTRACTDIR="${WORKDIR}/editor_new"

    mkdir -p "${OLD_EXTRACTDIR}" "${NEW_EXTRACTDIR}"

    elog() {
        echo "[$(date '+%H:%M:%S')] [editor] $*" >> "${LOGFILE}"
        echo "[$(date '+%H:%M:%S')] [editor] $*" >&2
    }

    fallback_nano() {
        elog "Falling back to editor for: ${FBUFILE}"
        open_editor "${FBUFILE}"
        exit 0
    }

    # Read OLD_RELEASE and NEW_RELEASE from state file
    if [ ! -f "${STATE_FILE}" ]; then
        elog "ERROR: state file not found: ${STATE_FILE}"
        fallback_nano
    fi
    . "${STATE_FILE}"

    # Derive system path from freebsd-update path
    # /var/db/freebsd-update/merge/new//etc/foo.conf -> /etc/foo.conf
    SYSPATH=$(echo "${FBUFILE}" | sed 's|.*//||')
    SYSPATH="/${SYSPATH}"

    elog "File : ${SYSPATH}"
    elog "OLD  : ${OLD_RELEASE}  NEW: ${NEW_RELEASE}"

    # Extract OLD and NEW base versions from txz archives
    editor_extract() {
        release="$1"
        syspath="$2"
        extractdir="$3"
        relpath="${syspath#/}"
        for archive in ${ARCHIVES}; do
            txzpath="${WORKDIR}/txz/${release}/${archive}"
            [ -f "${txzpath}" ] || continue
            if tar -tf "${txzpath}" "./${relpath}" >/dev/null 2>&1; then
                tar -xf "${txzpath}" -C "${extractdir}" "./${relpath}" 2>/dev/null && {
                    echo "${extractdir}/${relpath}"
                    return 0
                }
            fi
        done
        return 1
    }

    base_old=$(editor_extract "${OLD_RELEASE}" "${SYSPATH}" "${OLD_EXTRACTDIR}") || {
        elog "WARN: ${SYSPATH} not found in base ${OLD_RELEASE}"
        fallback_nano
    }

    base_new=$(editor_extract "${NEW_RELEASE}" "${SYSPATH}" "${NEW_EXTRACTDIR}") || {
        elog "WARN: ${SYSPATH} not found in base ${NEW_RELEASE}"
        fallback_nano
    }

    # Use current system file (without conflict markers) as merge base
    if [ -f "${SYSPATH}" ]; then
        cp -p "${SYSPATH}" "${FBUFILE}"
        elog "Base: current system file"
    else
        elog "WARN: ${SYSPATH} does not exist on system, using fbu file as-is"
    fi

    # 3-way merge in-place
    merge_exit=0
    merge_output=$(merge "${FBUFILE}" "${base_old}" "${base_new}" 2>&1) || merge_exit=$?

    case "${merge_exit}" in
        0)
            elog "OK: clean merge"
            exit 0
            ;;
        1)
            elog "Residual conflicts — opening nano"
            echo ""
            echo "  +======================================================+"
            printf "  |  Conflicts in: %-38s|\n" "${SYSPATH}"
            echo "  |  Resolve the <<<<<<< markers, then save and exit.   |"
            echo "  +======================================================+"
            echo ""
            open_editor "${FBUFILE}"
            exit 0
            ;;
        *)
            elog "ERROR: merge failed (exit ${merge_exit}): ${merge_output}"
            fallback_nano
            ;;
    esac
}

# ---------------------------------------------------------------------------
# ██████████████████████████████████████████████████████████████████████████
# MAIN MODE
# ██████████████████████████████████████████████████████████████████████████
# ---------------------------------------------------------------------------

main_mode() {

# ---------------------------------------------------------------------------
# Architecture detection
# ---------------------------------------------------------------------------

# FreeBSD mirror URL format: releases/$machine/$arch
# uname -m = machine hardware name (e.g. amd64, aarch64)
# uname -p = processor type       (e.g. amd64, aarch64)
MACH="$(uname -m)"
PROC="$(uname -p)"

FETCH_BASE="https://download.freebsd.org/releases/${MACH}/${PROC}"

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

TIMESTAMP=$(date '+%Y%m%d_%H%M%S')
OLD_EXTRACTDIR="${WORKDIR}/old"
NEW_EXTRACTDIR="${WORKDIR}/new"
LOGFILE="${WORKDIR}/automerge.log"
REPORT="${WORKDIR}/report.txt"
BACKUPDIR="${WORKDIR}/backup_${TIMESTAMP}"
BACKUP_MANIFEST="${BACKUPDIR}/MANIFEST.txt"
RESTORE_SCRIPT="${BACKUPDIR}/restore.sh"

OLD_RELEASE=""
NEW_RELEASE=""
DO_UPGRADE=0
AUTO_CONFIRM=0
USERLAND_ONLY=0

# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

# log() writes to stderr — does NOT pollute $() subshells
log() {
    echo "[$(date '+%H:%M:%S')] $*" | tee -a "${LOGFILE}" >&2
}

die() {
    printf "ERROR: %s\n" "$*" >&2
    exit 1
}

# Open the best available editor — respect user preference,
# fall back to ee then vi (both always present in FreeBSD base)
open_editor() {
    filepath="$1"
    if [ -n "${VISUAL}" ] && command -v "${VISUAL}" >/dev/null 2>&1; then
        "${VISUAL}" "${filepath}" < /dev/tty > /dev/tty 2>&1
    elif [ -n "${EDITOR}" ] && command -v "${EDITOR}" >/dev/null 2>&1; then
        "${EDITOR}" "${filepath}" < /dev/tty > /dev/tty 2>&1
    elif command -v ee >/dev/null 2>&1; then
        ee "${filepath}" < /dev/tty > /dev/tty 2>&1
    else
        vi "${filepath}" < /dev/tty > /dev/tty 2>&1
    fi
}

major_of() {
    echo "$1" | cut -d. -f1
}

is_major_upgrade() {
    [ "$(major_of "${OLD_RELEASE}")" != "$(major_of "${NEW_RELEASE}")" ]
}

# Returns: MAJOR, MINOR, STABLE, or CURRENT
upgrade_type() {
    case "${NEW_RELEASE}" in
        *-STABLE)  echo "STABLE"  ;;
        *-CURRENT) echo "CURRENT" ;;
        *)
            if is_major_upgrade; then
                echo "MAJOR"
            else
                echo "MINOR"
            fi
            ;;
    esac
}

# ---------------------------------------------------------------------------
# Reboot notice (no automatic reboot)
# ---------------------------------------------------------------------------

warn_reboot() {
    msg="$1"
    echo ""
    echo "+======================================================+"
    echo "|                   *** WARNING ***                    |"
    echo "+======================================================+"
    printf  "|  %-52s  |\n" "${msg}"
    echo "|                                                      |"
    echo "|  Run manually: reboot                                |"
    echo "+======================================================+"
    echo ""
    log "Reboot notice: ${msg}"
}

# ---------------------------------------------------------------------------
# Persistent state (survives reboot)
# ---------------------------------------------------------------------------

write_state() {
    mkdir -p "${STATE_DIR}"
    cat > "${STATE_FILE}" << EOF
OLD_RELEASE="${OLD_RELEASE}"
NEW_RELEASE="${NEW_RELEASE}"
PHASE="$1"
BACKUPDIR="${BACKUPDIR}"
LOGFILE="${LOGFILE}"
REPORT="${REPORT}"
RESTORE_SCRIPT="${RESTORE_SCRIPT}"
EOF
    log "State saved: PHASE=$1"
}

# ---------------------------------------------------------------------------
# rc.d script for second phase after reboot (major upgrade)
# ---------------------------------------------------------------------------

install_rcd_resume() {
    cat > "${RCD_SCRIPT}" << 'RCDEOF'
#!/bin/sh
# PROVIDE: automerge_resume
# REQUIRE: LOGIN
# KEYWORD: shutdown

. /etc/rc.subr

name="automerge_resume"
rcvar="automerge_resume_enable"
start_cmd="automerge_resume_start"
stop_cmd=":"

automerge_resume_start()
{
    /usr/sbin/daemon -f /bin/sh /var/db/automerge/resume.sh
}

load_rc_config $name
: ${automerge_resume_enable:="NO"}
run_rc_command "$1"
RCDEOF
    chmod 555 "${RCD_SCRIPT}"

    cat > "${STATE_DIR}/resume.sh" << RESUMEOF
#!/bin/sh
LOGFILE="${LOGFILE}"
RCD_SCRIPT="${RCD_SCRIPT}"

log() {
    echo "[\$(date '+%H:%M:%S')] \$*" | tee -a "\${LOGFILE}"
}

sleep 10
log "=== automerge_resume: second install phase ==="
freebsd-update install || log "WARN: freebsd-update install returned an error"

sysrc -x automerge_resume_enable 2>/dev/null || true
rm -f "\${RCD_SCRIPT}"
rm -f "${STATE_DIR}/resume.sh"
rm -f "${STATE_FILE}"

log "=== Second phase complete ==="
echo ""
echo "+======================================================+"
echo "|                   *** WARNING ***                    |"
echo "+======================================================+"
echo "|  Second phase complete. Userland updated.            |"
echo "|                                                      |"
echo "|  Run manually: reboot                                |"
echo "+======================================================+"
echo ""
RESUMEOF

    chmod 700 "${STATE_DIR}/resume.sh"
    sysrc automerge_resume_enable="YES"
    log "rc.d resume installed: ${RCD_SCRIPT}"
}

# ---------------------------------------------------------------------------
# Query mirror and return list of available versions
# ---------------------------------------------------------------------------

query_mirror_versions() {
    pattern="$1"
    tmpindex="${WORKDIR}/mirror_index.html"
    mkdir -p "${WORKDIR}"
    fetch -qo "${tmpindex}" "${FETCH_BASE}/" 2>/dev/null || \
        die "Cannot reach mirror: ${FETCH_BASE}/"
    grep -oE "href=\"${pattern}" "${tmpindex}" \
        | grep -oE "${pattern}" \
        | sed 's|/||g' \
        | sort -t. -k1,1n -k2,2n
    rm -f "${tmpindex}"
}

# ---------------------------------------------------------------------------
# Interactive target selection menu
# ---------------------------------------------------------------------------

select_target_menu() {
    major="$1"

    echo ""
    echo "  Querying FreeBSD mirror for available versions..."
    echo ""

    # Fetch all available versions filtered by major if given
    if [ -n "${major}" ]; then
        _pat_rel="${major}\\.[0-9]+-RELEASE/"
        _pat_rc="${major}\\.[0-9]+-RC[0-9]*/"
        _pat_beta="${major}\\.[0-9]+-BETA[0-9]*/"
        _pat_stable="${major}-STABLE/"
    else
        _pat_rel="[0-9]+\\.[0-9]+-RELEASE/"
        _pat_rc="[0-9]+\\.[0-9]+-RC[0-9]*/"
        _pat_beta="[0-9]+\\.[0-9]+-BETA[0-9]*/"
        _pat_stable="[0-9]+-STABLE/"
    fi

    _releases=$(query_mirror_versions "${_pat_rel}" | sort -t. -k2 -rn | head -5)
    _rcs=$(query_mirror_versions "${_pat_rc}" | sort -t. -k2 -rn | head -3)
    _betas=$(query_mirror_versions "${_pat_beta}" | sort -t. -k2 -rn | head -3)
    _stables=$(query_mirror_versions "${_pat_stable}" | sort -rn | head -3)

    echo "  +======================================================+"
    echo "  |           SELECT UPGRADE TARGET                      |"
    echo "  +======================================================+"
    echo "  |                                                      |"

    _idx=0
    _menu=""

    # RELEASE (stable, recommended)
    for _v in ${_releases}; do
        _idx=$((_idx + 1))
        _menu="${_menu}${_idx}:${_v} "
        printf "  |  [%2d] %-30s %s  |\n" "${_idx}" "${_v}" "RELEASE   "
    done

    # STABLE (rolling) — warn user these may be unreliable
    for _v in ${_stables}; do
        _idx=$((_idx + 1))
        _menu="${_menu}${_idx}:${_v} "
        printf "  |  [%2d] %-30s %s  |\n" "${_idx}" "${_v}" "STABLE (*)"
    done

    # RC (pre-release)
    for _v in ${_rcs}; do
        _idx=$((_idx + 1))
        _menu="${_menu}${_idx}:${_v} "
        printf "  |  [%2d] %-30s %s  |\n" "${_idx}" "${_v}" "RC        "
    done

    # BETA (testing)
    for _v in ${_betas}; do
        _idx=$((_idx + 1))
        _menu="${_menu}${_idx}:${_v} "
        printf "  |  [%2d] %-30s %s  |\n" "${_idx}" "${_v}" "BETA      "
    done

    # Manual entry
    _idx=$((_idx + 1))
    _manual_idx="${_idx}"
    printf "  |  [%2d] %-30s %s  |\n" "${_idx}" "Enter version manually" "          "

    echo "  |                                                      |"
    echo "  +======================================================+"
    echo ""
    printf "  Your current version : %s\n" "${OLD_RELEASE}"
    echo "  (*) STABLE/CURRENT snapshots may be temporarily unavailable on the mirror."
    printf "  Choice [1-%d]: " "${_idx}"
    read -r _choice

    if [ "${_choice}" -eq "${_manual_idx}" ] 2>/dev/null; then
        printf "  Enter target (e.g. 14.4-RELEASE, 15-STABLE, 15.1-RC2): "
        read -r NEW_RELEASE
    else
        _picked=$(echo "${_menu}" | tr ' ' '\n' | grep "^${_choice}:" | cut -d: -f2)
        if [ -z "${_picked}" ]; then
            die "Invalid choice: ${_choice}"
        fi
        NEW_RELEASE="${_picked}"
    fi

    # Warn if downgrade or same version
    _old_major=$(echo "${OLD_RELEASE}" | cut -d. -f1)
    _new_major=$(echo "${NEW_RELEASE}" | cut -d. -f1 | cut -d- -f1)
    if [ "${NEW_RELEASE}" = "${OLD_RELEASE}" ]; then
        echo ""
        echo "  WARNING: target is the same as current version (${OLD_RELEASE})."
        printf "  Continue anyway? [y/n]: "
        read -r _ans
        case "${_ans}" in [Yy]*) ;; *) die "Aborted." ;; esac
    fi

    log "Target selected: ${NEW_RELEASE}"
}

# ---------------------------------------------------------------------------
# Detect OLD and NEW release versions
# ---------------------------------------------------------------------------

detect_versions() {
    # Argument parsing: [--auto-confirm] [major]
    arg=""
    for _a in "$@"; do
        case "${_a}" in
            --auto-confirm)
                AUTO_CONFIRM=1
                log "Auto-confirm mode enabled: y/n prompts will be answered automatically"
                ;;
            *)
                arg="${_a}"
                ;;
        esac
    done

    OLD_RELEASE=$(uname -r | sed 's/-p[0-9]*$//')

    if [ -z "${arg}" ]; then
        if [ -f /var/db/freebsd-update/tag ]; then
            NEW_RELEASE=$(awk -F'|' '{print $3}' /var/db/freebsd-update/tag 2>/dev/null || cut -d'/' -f1 /var/db/freebsd-update/tag)
            DO_UPGRADE=0
            log "Resuming existing upgrade to: ${NEW_RELEASE}"
        else
            # No argument and no tag: show full menu
            select_target_menu ""
            DO_UPGRADE=1
        fi
    elif echo "${arg}" | grep -qE '^[0-9]+$'; then
        # Major number: show filtered menu (e.g. 14, 15)
        select_target_menu "${arg}"
        DO_UPGRADE=1
    elif echo "${arg}" | grep -qE '^[0-9]+'; then
        # Explicit version: use directly, no menu
        NEW_RELEASE="${arg}"
        DO_UPGRADE=1
        log "Target specified directly: ${NEW_RELEASE}"
    else
        die "Unrecognized argument: '${arg}'\nUsage: $0 [--auto-confirm] [major]\nEx:    $0          (full menu)\nEx:    $0 14        (menu filtered to branch 14)\nEx:    $0 14.4-RELEASE  (direct, no menu)\nEx:    $0 --auto-confirm 14"
    fi
}

# ---------------------------------------------------------------------------
# Download and extraction
# ---------------------------------------------------------------------------

download_archive() {
    release="$1"; archive="$2"
    destdir="${WORKDIR}/txz/${release}"
    dest="${destdir}/${archive}"
    mkdir -p "${destdir}"
    if [ -f "${dest}" ]; then
        log "${release}/${archive} already present, skipping download."
        return 0
    fi
    log "Downloading ${FETCH_BASE}/${release}/${archive} ..."
    fetch -o "${dest}" "${FETCH_BASE}/${release}/${archive}" 2>/dev/null || {
        log "WARN: failed to download ${release}/${archive}"
        return 1
    }
}

extract_file() {
    release="$1"; filepath="$2"; extractdir="$3"
    relpath="${filepath#/}"
    for archive in ${ARCHIVES}; do
        txzpath="${WORKDIR}/txz/${release}/${archive}"
        [ -f "${txzpath}" ] || continue
        if tar -tf "${txzpath}" "./${relpath}" >/dev/null 2>&1; then
            tar -xf "${txzpath}" -C "${extractdir}" "./${relpath}" 2>/dev/null && {
                echo "${extractdir}/${relpath}"
                return 0
            }
        fi
    done
    return 1
}

find_conflicted_files() {
    grep -rl '^<<<<<<< ' /etc /usr/share 2>/dev/null || true
}

# ---------------------------------------------------------------------------
# Centralized backup
# ---------------------------------------------------------------------------

backup_all_files() {
    log "Creating backup in ${BACKUPDIR} ..."
    mkdir -p "${BACKUPDIR}"

    cat > "${BACKUP_MANIFEST}" << EOF
# freebsd-automerge backup
# Date     : $(date)
# OLD      : ${OLD_RELEASE}
# NEW      : ${NEW_RELEASE}
# Hostname : $(hostname)
# Format   : <sha256> <original_path>
EOF

    echo "${CONFLICTED}" | while read -r filepath; do
        mkdir -p "${BACKUPDIR}$(dirname "${filepath}")"
        cp -p "${filepath}" "${BACKUPDIR}${filepath}"
        sha256 -q "${filepath}" | awk -v p="${filepath}" '{print $1 "  " p}' \
            >> "${BACKUP_MANIFEST}"
        log "  Saved: ${filepath}"
    done

    generate_restore_script
    log "Backup complete. Restore script: ${RESTORE_SCRIPT}"
    echo ""
}

generate_restore_script() {
    cat > "${RESTORE_SCRIPT}" << RESTORE_EOF
#!/bin/sh
# restore.sh — restores files to their pre-merge state
# OLD: ${OLD_RELEASE}  NEW: ${NEW_RELEASE}
# Usage: sh ${RESTORE_SCRIPT}

set -e
BACKUPDIR="${BACKUPDIR}"
[ "\$(id -u)" -ne 0 ] && echo "ERROR: must run as root." && exit 1

echo "Verifying backup integrity..."
errors=0
while read -r checksum filepath; do
    case "\${checksum}" in '#'*|'') continue ;; esac
    actual=\$(sha256 -q "\${BACKUPDIR}\${filepath}" 2>/dev/null || echo MISSING)
    [ "\${actual}" != "\${checksum}" ] && echo "  ERROR: \${filepath}" && errors=\$((errors+1))
done < "${BACKUP_MANIFEST}"
[ "\${errors}" -gt 0 ] && echo "\${errors} corrupted files. Aborted." && exit 1
echo "  OK. Restoring..."
RESTORE_EOF

    echo "${CONFLICTED}" | while read -r filepath; do
        echo "cp -p \"\${BACKUPDIR}${filepath}\" \"${filepath}\" && echo \"  OK: ${filepath}\"" \
            >> "${RESTORE_SCRIPT}"
    done
    echo 'echo ""; echo "=== Restore complete ==="' >> "${RESTORE_SCRIPT}"
    chmod 700 "${RESTORE_SCRIPT}"
}

# ---------------------------------------------------------------------------
# expect installation helper
# ---------------------------------------------------------------------------

ensure_expect() {
    if command -v expect >/dev/null 2>&1; then
        return 0
    fi

    echo ""
    echo "  'expect' is required for --auto-confirm but is not installed."
    printf "  Install it now? (pkg install expect) [y/n]: "
    read -r answer
    case "${answer}" in
        [Yy]*)
            pkg install -y expect || die "Failed to install expect."
            log "expect installed successfully."
            ;;
        *)
            log "WARN: expect not installed. Proceeding without auto-confirm."
            AUTO_CONFIRM=0
            ;;
    esac
}

# ---------------------------------------------------------------------------
# freebsd-update upgrade with EDITOR=this script
# ---------------------------------------------------------------------------

run_freebsd_update_upgrade() {
    # Write state first so automerge editor mode can read OLD/NEW_RELEASE
    write_state "UPGRADING"

    SELF=$(realpath "$0")

    # Handle --auto-confirm: ensure expect is available
    if [ "${AUTO_CONFIRM}" -eq 1 ]; then
        ensure_expect
    fi

    log "Running: EDITOR=${SELF} freebsd-update -r ${NEW_RELEASE} upgrade"
    log "Conflicts will be resolved automatically (EDITOR mode)"
    echo ""
    echo "------------------------------------------------------"

    if [ "${AUTO_CONFIRM}" -eq 1 ]; then
        log "Auto-confirm active: answering 'y' automatically to freebsd-update prompts"
        expect -c "
            set timeout -1
            spawn env EDITOR=${SELF} freebsd-update -r ${NEW_RELEASE} upgrade
            expect {
                -re {Does this look reasonable.*\?} {
                    send \"y\r\"
                    exp_continue
                }
                -re {\(y/n\)} {
                    send \"y\r\"
                    exp_continue
                }
                -re {\[y\|n\]} {
                    send \"y\r\"
                    exp_continue
                }
                eof
            }
            catch wait result
            exit [lindex \$result 3]
        " || die "freebsd-update upgrade failed."
    else
        EDITOR="${SELF}" freebsd-update -r "${NEW_RELEASE}" upgrade || \
            die "freebsd-update upgrade failed."
    fi

    echo "------------------------------------------------------"
    echo ""
    [ -f /var/db/freebsd-update/tag ] || \
        die "Tag file not found. Upgrade did not complete."
}

# ---------------------------------------------------------------------------
# Resolve remaining <<<<<<< conflict markers (post-upgrade)
# ---------------------------------------------------------------------------

resolve_file() {
    filepath="$1"
    log "--- ${filepath}"

    base_old=$(extract_file "${OLD_RELEASE}" "${filepath}" "${OLD_EXTRACTDIR}") || {
        log "  WARN: not found in base ${OLD_RELEASE} — skipping"
        echo "SKIPPED (no base old): ${filepath}" >> "${REPORT}"
        return
    }
    base_new=$(extract_file "${NEW_RELEASE}" "${filepath}" "${NEW_EXTRACTDIR}") || {
        log "  WARN: not found in base ${NEW_RELEASE} — skipping"
        echo "SKIPPED (no base new): ${filepath}" >> "${REPORT}"
        return
    }

    merge_exit=0
    merge_output=$(merge "${filepath}" "${base_old}" "${base_new}" 2>&1) || merge_exit=$?

    case "${merge_exit}" in
        0) log "  OK: clean merge"
           echo "OK:       ${filepath}" >> "${REPORT}" ;;
        1) log "  WARN: residual conflicts — manual review required"
           echo "CONFLICT: ${filepath}" >> "${REPORT}" ;;
        *) log "  ERROR: merge failed: ${merge_output}"
           echo "ERROR:    ${filepath}" >> "${REPORT}"
           cp -p "${BACKUPDIR}${filepath}" "${filepath}"
           log "  Restored from backup" ;;
    esac
}

# ---------------------------------------------------------------------------
# Report
# ---------------------------------------------------------------------------

print_report() {
    echo ""
    echo "======================================================"
    echo " MERGE REPORT"
    echo "======================================================"
    echo "  OLD : ${OLD_RELEASE}   NEW : ${NEW_RELEASE}"
    echo ""
    ok_count=0;       [ -f "${REPORT}" ] && ok_count=$(grep -c '^OK:'       "${REPORT}" 2>/dev/null) || true
    conflict_count=0; [ -f "${REPORT}" ] && conflict_count=$(grep -c '^CONFLICT:' "${REPORT}" 2>/dev/null) || true
    skip_count=0;     [ -f "${REPORT}" ] && skip_count=$(grep -c '^SKIPPED'   "${REPORT}" 2>/dev/null) || true
    error_count=0;    [ -f "${REPORT}" ] && error_count=$(grep -c '^ERROR:'    "${REPORT}" 2>/dev/null) || true
    echo "  Clean merges      : ${ok_count}"
    echo "  Residual conflicts: ${conflict_count}"
    echo "  Skipped           : ${skip_count}"
    echo "  Errors            : ${error_count}"
    echo ""
    [ "${conflict_count}" -gt 0 ] && \
        grep '^CONFLICT:' "${REPORT}" | sed 's/^CONFLICT:/  CONFLICT ->/' && echo ""
    [ "${error_count}" -gt 0 ] && \
        grep '^ERROR:' "${REPORT}" | sed 's/^ERROR:/  ERROR ->/' && echo ""
    echo "  Backup  : ${BACKUPDIR}"
    echo "  Restore : ${RESTORE_SCRIPT}"
    echo "  Log     : ${LOGFILE}"
    echo "  Report  : ${REPORT}"
    echo "======================================================"
    echo ""
}

# ---------------------------------------------------------------------------
# Ask user whether to upgrade kernel + userland or userland only
# ---------------------------------------------------------------------------

ask_upgrade_components() {
    echo ""
    echo "  +======================================================+"
    echo "  |           UPGRADE COMPONENT SELECTION                |"
    echo "  +======================================================+"
    echo "  |                                                      |"
    echo "  |  [1] Kernel + Userland  (standard upgrade)           |"
    echo "  |  [2] Userland only      (keep current custom kernel) |"
    echo "  |                                                      |"
    echo "  +======================================================+"
    echo ""
    printf "  Choice [1/2]: "
    read -r choice
    case "${choice}" in
        2)
            USERLAND_ONLY=1
            log "Userland-only upgrade selected — kernel will not be touched."
            ;;
        *)
            USERLAND_ONLY=0
            log "Full upgrade selected — kernel + userland."
            ;;
    esac
}

# ---------------------------------------------------------------------------
# Install + reboot notice
# ---------------------------------------------------------------------------

run_install_and_reboot_cycle() {
    conflict_count=0
    [ -f "${REPORT}" ] && conflict_count=$(grep -c '^CONFLICT:' "${REPORT}" 2>/dev/null) || true
    if [ "${conflict_count}" -gt 0 ]; then
        echo "  ${conflict_count} residual conflict(s) require manual resolution."
        echo "  Then run: freebsd-update install && reboot"
        return
    fi

    _uptype=$(upgrade_type)
    if [ "${USERLAND_ONLY}" -eq 1 ]; then
        # Userland-only: single install pass, single reboot
        log "Userland-only install — skipping kernel"
        echo ""
        echo "  Installing userland only (kernel unchanged)..."
        echo ""
        freebsd-update install || log "WARN: freebsd-update install returned an error"
        warn_reboot "Userland upgrade complete. Please reboot."
    elif [ "${_uptype}" = "MAJOR" ]; then
        log "Major upgrade — first pass (kernel)"
        echo ""
        echo "  Major upgrade: ${OLD_RELEASE} -> ${NEW_RELEASE}"
        echo "  Pass 1/2: installing new kernel..."
        echo ""
        freebsd-update install || log "WARN: freebsd-update install returned an error"
        write_state "AFTER_FIRST_REBOOT"
        install_rcd_resume
        warn_reboot "Kernel updated (1/2). Reboot now — userland will be installed automatically."
    else
        log "${_uptype} upgrade — installing"
        freebsd-update install || log "WARN: freebsd-update install returned an error"
        warn_reboot "Upgrade complete. Please reboot."
    fi
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

clear

[ "$(id -u)" -ne 0 ] && die "This script must run as root."
command -v merge >/dev/null 2>&1 || die "merge(1) not found. Install with: pkg install rcs"

mkdir -p "${WORKDIR}/txz" "${OLD_EXTRACTDIR}" "${NEW_EXTRACTDIR}" "${STATE_DIR}"
: > "${LOGFILE}"
: > "${REPORT}"

detect_versions "$@"

log "=== freebsd-automerge started ==="
log "Arch   : ${MACH}/${PROC}"
log "OLD    : ${OLD_RELEASE}"
log "NEW    : ${NEW_RELEASE}"
log "Mirror : ${FETCH_BASE}"

ask_upgrade_components

log "Components : $([ "${USERLAND_ONLY}" -eq 1 ] && echo 'userland only' || echo 'kernel + userland')"
log "Type       : $(upgrade_type)"

[ "${DO_UPGRADE}" -eq 1 ] && run_freebsd_update_upgrade

for archive in ${ARCHIVES}; do
    download_archive "${OLD_RELEASE}" "${archive}" || true
    download_archive "${NEW_RELEASE}" "${archive}" || true
done

log "Searching for remaining <<<<<<< conflict markers..."
CONFLICTED=$(find_conflicted_files)

if [ -z "${CONFLICTED}" ]; then
    log "No residual conflicts found."
    run_install_and_reboot_cycle
    exit 0
fi

COUNT=$(echo "${CONFLICTED}" | wc -l | tr -d ' ')
log "Found ${COUNT} file(s) with residual conflicts:"
echo "${CONFLICTED}" | while read -r f; do log "  ${f}"; done

backup_all_files

log "Starting 3-way merge on residual conflicts..."
echo "${CONFLICTED}" | while read -r filepath; do
    resolve_file "${filepath}"
done

print_report
run_install_and_reboot_cycle

} # end main_mode

# ---------------------------------------------------------------------------
# DISPATCH
# ---------------------------------------------------------------------------

case "${_MODE}" in
    EDITOR) editor_mode "$1" ;;
    MAIN)   main_mode   "$@" ;;
esac