#!/bin/bash
# ZFS replication: snapshot source and send to destination (local or remote via SSH).
# Usage: omv-zfs-replicate <dataset> <prefix> <dest> <type> <remotehost>
#                          <certref> <sshuser> <sshport> <recursive> <raw>
#                          [-r retention] [--bookmarks] [--no-snapshot]
#                          [-e always|onerror] [-c comment]
#
#   dataset       - source ZFS dataset (e.g. tank/data)
#   prefix        - snapshot name prefix (e.g. replicate)
#   dest          - destination dataset path (e.g. backup/data)
#   type          - "local" or "remote"
#   remotehost    - remote hostname/IP (ignored when type=local)
#   certref       - OMV SSH certificate UUID (ignored when type=local, empty = default key)
#   sshuser       - SSH user on remote (ignored when type=local)
#   sshport       - SSH port on remote (ignored when type=local)
#   recursive     - 1 to snapshot/send recursively, 0 otherwise
#   raw           - 1 to use zfs send -w (raw/encrypted), 0 otherwise
#   -r N          - keep N replication snapshots on source (0 = keep all); ignored with --no-snapshot
#   --bookmarks   - use ZFS bookmarks as incremental base; allows aggressive source pruning
#   --no-snapshot - do not create a new snapshot; replicate the most recent existing snapshot
#                   matching the prefix. Use when a separate snapshot job manages the snapshots.
#   -e            - email mode: "always" or "onerror"
#   -c            - email subject comment

set -euo pipefail

LOGFILE="/var/log/omv-zfs-replicate.log"
FULL_COMMAND="$0 $*"

DATASET="${1:?dataset required}"
PREFIX="${2:?prefix required}"
DEST="${3:?dest required}"
TYPE="${4:-local}"
REMOTEHOST="${5:-}"
CERTREF="${6:-}"
SSHUSER="${7:-root}"
SSHPORT="${8:-22}"
RECURSIVE="${9:-0}"
RAW="${10:-0}"
shift 10 || true

# Parse optional flags
email_mode=""
email_comment=""
RETENTION=0
USEBOOKMARKS=0
NOSNAPSHOT=0
while [[ $# -gt 0 ]]; do
    case "$1" in
        -r)            RETENTION="$2";  shift 2 ;;
        --bookmarks)   USEBOOKMARKS=1;  shift   ;;
        --no-snapshot) NOSNAPSHOT=1;    shift   ;;
        -e)            email_mode="$2"; shift 2 ;;
        -c)            email_comment="$2"; shift 2 ;;
        *)             shift ;;
    esac
done

# ---------------------------------------------------------------------------
# Output capture + email on exit (compose-style)
# ---------------------------------------------------------------------------
tmpfile=""
if [ -n "${email_mode}" ]; then
    tmpfile=$(mktemp)
    exec > "${tmpfile}" 2>&1
else
    exec > >(tee -a "${LOGFILE}") 2>&1
fi

_send_email() {
    local rc="$1"
    # Always append to logfile when using email mode
    if [ -f "${tmpfile:-}" ]; then
        cat "${tmpfile}" >> "${LOGFILE}"
    fi
    if [ -n "${email_mode}" ] && [ -f "${tmpfile:-}" ]; then
        if [ "${email_mode}" = "always" ] || \
           { [ "${email_mode}" = "onerror" ] && [ "${rc}" -ne 0 ]; }; then
            mail -E -s "ZFS Replication${email_comment:+ - ${email_comment}}" \
                 -a "From: Cron Daemon <root>" root < "${tmpfile}" >/dev/null 2>&1
        fi
        rm -f "${tmpfile}"
    fi
}
trap '_send_email $?' EXIT

# ---------------------------------------------------------------------------
# SSH helpers
# ---------------------------------------------------------------------------
SSH_KEY_DIR="${OMV_SSH_KEYS_DIR:-/etc/ssh}"
SSH_KEY_PREFIX="${OMV_SSH_KEY_PREFIX:-openmediavault-}"

SSH_CMD=""
if [ "${TYPE}" = "remote" ]; then
    SSH_OPTS="-o BatchMode=yes -o StrictHostKeyChecking=accept-new -p ${SSHPORT}"
    if [ -n "${CERTREF}" ]; then
        SSH_OPTS="${SSH_OPTS} -i ${SSH_KEY_DIR}/${SSH_KEY_PREFIX}${CERTREF}"
    fi
    SSH_CMD="ssh ${SSH_OPTS} ${SSHUSER}@${REMOTEHOST}"
fi

dest_exec() {
    if [ -n "${SSH_CMD}" ]; then
        ${SSH_CMD} "$@"
    else
        "$@"
    fi
}

run() {
    log "+ $*"
    "$@"
}

# ---------------------------------------------------------------------------
# Job tag — prefixed on every log line so concurrent jobs are distinguishable
# ---------------------------------------------------------------------------
if [ "${TYPE}" = "remote" ]; then
    JOB_TAG="[${DATASET} -> ${SSHUSER}@${REMOTEHOST}:${DEST}]"
else
    JOB_TAG="[${DATASET} -> ${DEST}]"
fi

log() { echo "$(date -Iseconds): ${JOB_TAG} $*"; }

# ---------------------------------------------------------------------------
# Log the full command for troubleshooting
# ---------------------------------------------------------------------------
log "Command: ${FULL_COMMAND}"

# ---------------------------------------------------------------------------
# Create new snapshot on source, or use the most recent existing one
# ---------------------------------------------------------------------------
if [ "${NOSNAPSHOT}" = "1" ]; then
    NEW_SNAP=$(zfs list -H -t snapshot -o name -s creation -r "${DATASET}" \
        | grep "^${DATASET}@${PREFIX}-" | tail -1 || true)
    if [ -z "${NEW_SNAP}" ]; then
        log "ERROR: No snapshots found matching ${DATASET}@${PREFIX}-*. Cannot replicate. Check that the snapshot prefix in this replication job matches the prefix used by the snapshot job managing this dataset."
        exit 1
    fi
    SNAPNAME="${NEW_SNAP##*@}"
    log "Using existing snapshot ${NEW_SNAP} (--no-snapshot mode)"
    if dest_exec zfs list -H -t snapshot "${DEST}@${SNAPNAME}" >/dev/null 2>&1; then
        log "Destination already has ${SNAPNAME}, replication is up to date."
        exit 0
    fi
else
    SNAPNAME="${PREFIX}-$(date +%Y%m%d-%H%M%S)"
    NEW_SNAP="${DATASET}@${SNAPNAME}"
    snap_err=$(mktemp)
    if [ "${RECURSIVE}" = "1" ]; then
        log "Creating recursive snapshot ${NEW_SNAP}"
        snap_cmd="zfs snapshot -r ${NEW_SNAP}"
    else
        log "Creating snapshot ${NEW_SNAP}"
        snap_cmd="zfs snapshot ${NEW_SNAP}"
    fi
    log "+ ${snap_cmd}"
    if ! ${snap_cmd} 2>"${snap_err}"; then
        if grep -q "dataset already exists" "${snap_err}"; then
            log "Snapshot ${NEW_SNAP} already exists (created by another job), using it for replication."
        else
            cat "${snap_err}" >&2
            rm -f "${snap_err}"
            exit 1
        fi
    fi
    rm -f "${snap_err}"
fi

# ---------------------------------------------------------------------------
# Determine send strategy: bookmark-based incremental, snapshot-based, or full
# ---------------------------------------------------------------------------
SEND_FLAGS=""
[ "${RAW}"       = "1" ] && SEND_FLAGS="${SEND_FLAGS} -w"
[ "${RECURSIVE}" = "1" ] && SEND_FLAGS="${SEND_FLAGS} -R"

_incremental_send() {
    local base="$1"
    log "Incremental send ${base} -> ${NEW_SNAP} to ${DEST}"
    log "+ zfs send${SEND_FLAGS:+ ${SEND_FLAGS}} -i ${base} ${NEW_SNAP} | zfs receive -F ${DEST}"
    zfs send ${SEND_FLAGS} -i "${base}" "${NEW_SNAP}" \
        | dest_exec zfs receive -F "${DEST}"
}

sent=0

# Try bookmark-based incremental (preferred when --bookmarks is set).
# Bookmarks cannot be used as the base for recursive sends (zfs send -R -i
# bookmark snapshot fails with "multiple snapshots cannot be sent from a
# bookmark"), so skip this path when RECURSIVE=1.
if [ "${USEBOOKMARKS}" = "1" ] && [ "${RECURSIVE}" = "0" ]; then
    PREV_BK=$(zfs list -H -t bookmark -o name -r "${DATASET}" \
        | grep "^${DATASET}#${PREFIX}-" | tail -1 || true)
    if [ -n "${PREV_BK}" ]; then
        PREV_SNAPNAME="${PREV_BK##*#}"
        if dest_exec zfs list -H -t snapshot "${DEST}@${PREV_SNAPNAME}" \
                >/dev/null 2>&1; then
            _incremental_send "${PREV_BK}"
            sent=1
        else
            log "Destination missing base snapshot ${PREV_SNAPNAME}, trying snapshot-based incremental."
        fi
    fi
fi

# Try snapshot-based incremental.
# Collect all destination snapshots matching the prefix in a single call (one
# SSH round-trip for remote), then walk source snapshots newest-to-oldest to
# find the most recent one that exists on both sides.  This handles the case
# where snapshots from another job (e.g. an auto-snapshot job sharing the same
# prefix) exist on the source but were never replicated to the destination.
if [ "${sent}" = "0" ]; then
    DST_SNAPS=$(dest_exec zfs list -H -t snapshot -o name -r "${DEST}" \
        2>/dev/null | grep "^${DEST}@${PREFIX}-" | sed 's/.*@//' || true)
    COMMON_BASE=""
    while IFS= read -r snap; do
        [ -n "${snap}" ] || continue
        snapname="${snap##*@}"
        if [[ $'\n'"${DST_SNAPS}"$'\n' == *$'\n'"${snapname}"$'\n'* ]]; then
            COMMON_BASE="${snap}"
            break
        fi
    done < <(zfs list -H -t snapshot -o name -s creation -r "${DATASET}" \
        | grep "^${DATASET}@${PREFIX}-" \
        | grep -v "^${NEW_SNAP}$" \
        | tac \
        || true)
    if [ -n "${COMMON_BASE}" ]; then
        _incremental_send "${COMMON_BASE}"
        sent=1
    else
        log "No common base snapshot found on source and destination, falling back to full send."
    fi
fi

# Full send if no incremental base was found.
if [ "${sent}" = "0" ]; then
    log "Full send ${NEW_SNAP} -> ${DEST}"
    log "+ zfs send${SEND_FLAGS:+ ${SEND_FLAGS}} ${NEW_SNAP} | zfs receive -F ${DEST}"
    zfs send ${SEND_FLAGS} "${NEW_SNAP}" | dest_exec zfs receive -F "${DEST}"
fi

# ---------------------------------------------------------------------------
# Post-send: source snapshot pruning (skipped in --no-snapshot mode)
# ---------------------------------------------------------------------------
if [ "${RETENTION}" -gt 0 ] && [ "${NOSNAPSHOT}" = "0" ]; then
    ALL_SNAPS=$(zfs list -H -t snapshot -o name -s creation -r "${DATASET}" \
        | grep "^${DATASET}@${PREFIX}-" || true)
    SNAP_COUNT=$(printf '%s\n' "${ALL_SNAPS}" | grep -c . || true)
    if [ "${SNAP_COUNT}" -gt "${RETENTION}" ]; then
        TO_PRUNE=$(printf '%s\n' "${ALL_SNAPS}" | head -n $(( SNAP_COUNT - RETENTION )))
        while IFS= read -r snap; do
            [ -n "${snap}" ] || continue
            log "Pruning old source snapshot ${snap}"
            [ "${RECURSIVE}" = "1" ] && run zfs destroy -r "${snap}" || run zfs destroy "${snap}"
        done <<< "${TO_PRUNE}"
    fi
fi

# ---------------------------------------------------------------------------
# Post-send: bookmark management (non-recursive only — see send guard above)
# ---------------------------------------------------------------------------
if [ "${USEBOOKMARKS}" = "1" ] && [ "${RECURSIVE}" = "0" ]; then
    log "Creating bookmark ${DATASET}#${SNAPNAME}"
    run zfs bookmark "${NEW_SNAP}" "${DATASET}#${SNAPNAME}"
    OLD_BKS=$(zfs list -H -t bookmark -o name -r "${DATASET}" \
        | grep "^${DATASET}#${PREFIX}-" \
        | grep -v "^${DATASET}#${SNAPNAME}$" || true)
    while IFS= read -r bk; do
        [ -n "${bk}" ] || continue
        log "Removing old bookmark ${bk}"
        run zfs destroy "${bk}"
    done <<< "${OLD_BKS}"
fi

log "Replication completed successfully."
