#!/usr/bin/env bash
#
# zfs-snapshot.sh — Interactive (menuconfig-style) ZFS storage manager for Proxmox VE
#
# A whiptail-based TUI to manage ZFS snapshots, scheduled auto-snapshots,
# datasets, LXC container mountpoints, dataset replication, and pool scrubs
# on a Proxmox VE host.
#
#   * Snapshots   — list / create / delete / rollback ZFS snapshots
#   * Schedules   — cron-driven auto-snapshots with auto-pruning (keep last N)
#   * Datasets    — view all ZFS datasets with usage info
#   * LXC Mounts  — view / add / remove bind-mountpoints on LXC containers (pct)
#   * Replication — zfs send | recv to another dataset (full or incremental)
#   * Scrub       — start / stop / schedule pool scrubs
#
# All cron entries managed by this tool live in a single owner file
# (/etc/cron.d/pve-zfs-tui); your other cron jobs are never touched.
#
# Usage:  ./zfs-snapshot.sh        (re-execs itself with sudo automatically)
#
set -euo pipefail

APP_TITLE="Proxmox ZFS Manager"

# ---------------------------------------------------------------------------
# Paths owned by this tool
# ---------------------------------------------------------------------------
CRON_FILE="/etc/cron.d/pve-zfs-tui"
SNAPSHOT_SCRIPT="/usr/local/bin/pve-zfs-snapshot.sh"
REPLICATION_SCRIPT="/usr/local/bin/pve-zfs-replicate.sh"
REPLICATION_CONFIG="/etc/pve-zfs-tui/replication.tsv"

# Single-owner cron header. Entries below are rewritten in full on every change.
CRON_HEADER='# Auto-generated by zfs-snapshot.sh — do not edit manually.
# Manage entries via the Schedules / Scrub menus.
SHELL=/bin/bash
PATH=/usr/sbin:/usr/bin:/sbin:/bin
MAILTO=""'

# ---------------------------------------------------------------------------
# Privilege handling: re-exec under sudo so zfs/zpool/pct run as root.
# ---------------------------------------------------------------------------
if [[ ${EUID} -ne 0 ]]; then
    echo "Root privileges are required. Re-running with sudo..."
    exec sudo "$0" "$@"
fi

# ---------------------------------------------------------------------------
# Ensure whiptail is available (part of the 'whiptail' package on Debian).
# ---------------------------------------------------------------------------
if ! command -v whiptail >/dev/null 2>&1; then
    echo "Installing whiptail (required for the TUI)..."
    apt-get update -qq || true
    apt-get install -y whiptail
fi

# ---------------------------------------------------------------------------
# Generic helpers
# ---------------------------------------------------------------------------
msg()  { whiptail --title "$APP_TITLE" --msgbox "$1" "${2:-12}" "${3:-72}"; }
yesno(){ whiptail --title "$APP_TITLE" --yesno "$1" "${2:-12}" "${3:-72}"; }

# Display the stdout of a producer command/function in a scrollable box.
#
# The producer runs with `set -e` disabled and its exit status discarded, so a
# non-zero result from zfs / awk / pipefail can never abort the whole TUI — this
# is what made "View all snapshots" crash back to the shell. The captured report
# is shown via --textbox (which reads from a file) instead of --msgbox (which
# takes the entire report as one argv string and chokes on large snapshot lists).
show() {
    local title="$1"; shift
    local out tmp
    set +e
    out="$("$@" 2>&1)"
    set -e
    tmp="$(mktemp)"
    printf '%s\n' "${out:-(no output)}" >"$tmp"
    whiptail --title "$title" --scrolltext --textbox "$tmp" 24 90 || true
    rm -f "$tmp"
}

# Run a command, capturing output, and show the result in a scrollable box.
run_and_show() {
    local title="$1"; shift
    local out rc
    set +e
    out="$("$@" 2>&1)"; rc=$?
    set -e
    whiptail --title "$title" --scrolltext --msgbox \
        "$(printf 'Command: %s\nExit code: %s\n\n%s' "$*" "$rc" "${out:-(no output)}")" \
        22 90
    return $rc
}

# Free-text input box. Echoes the entered value, or empty on cancel.
ask() {
    local prompt="$1" default="${2:-}"
    whiptail --title "$APP_TITLE" --inputbox "$prompt" 11 78 "$default" 3>&1 1>&2 2>&3
}

# Pick one item from a newline-separated list (passed on stdin). The visible
# tag IS the value. Echoes the chosen value, or empty on cancel / empty list.
#   pick "Prompt text" < <(some_command_producing_lines)
pick() {
    local prompt="$1"
    local -a items=()
    local line
    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        items+=("$line" "")
    done
    if [[ ${#items[@]} -eq 0 ]]; then
        msg "Nothing to choose from.\n\n$prompt" 10 70
        return 1
    fi
    whiptail --title "$APP_TITLE" --notags --menu "$prompt" 22 86 12 \
        "${items[@]}" 3>&1 1>&2 2>&3
}

# Timestamp used for snapshot names.
ts() { date +%Y%m%d-%H%M%S; }

# ---------------------------------------------------------------------------
# ZFS primitives
# ---------------------------------------------------------------------------
list_pools()    { zpool list -H -o name 2>/dev/null; }
list_datasets() { zfs list -H -o name,used,avail,mountpoint 2>/dev/null; }
list_dataset_names() { zfs list -H -o name 2>/dev/null; }

# All snapshots, one per line: name<TAB>used<TAB>refer<TAB>creation
list_snapshots() {
    zfs list -t snapshot -H -o name,used,refer,creation 2>/dev/null
}

# Short snapshot names (without the dataset@ prefix) for one exact dataset,
# oldest first.
dataset_snapshots() {
    local ds="$1"
    zfs list -H -t snapshot -o name -s creation "$ds" 2>/dev/null \
        | awk -F'@' -v d="$ds" '$1==d {print $2}'
}

latest_snapshot() { dataset_snapshots "$1" | tail -n1; }

dataset_exists() { zfs list -H "$1" >/dev/null 2>&1; }

# Latest snapshot short-name present on BOTH source and target.
latest_common_snapshot() {
    local src="$1" tgt="$2" s
    local tgt_snaps; tgt_snaps="$(dataset_snapshots "$tgt")"
    while IFS= read -r s; do
        [[ -z "$s" ]] && continue
        if grep -qxF "$s" <<<"$tgt_snaps"; then echo "$s"; fi
    done < <(dataset_snapshots "$src") | tail -n1
}

# ---------------------------------------------------------------------------
# Snapshots menu
# ---------------------------------------------------------------------------
# Report producers print to stdout; show() captures them with `set -e` off.
_report_snapshots() {
    local rows; rows="$(list_snapshots)"
    [[ -z "$rows" ]] && { echo "(no snapshots found)"; return 0; }
    printf '%s\n' "$rows" | awk -F'\t' \
        'BEGIN{printf "%-48s %8s %8s  %s\n","SNAPSHOT","USED","REFER","CREATED"}
         {printf "%-48s %8s %8s  %s\n",$1,$2,$3,$4}'
}
view_snapshots() { show "ZFS Snapshots" _report_snapshots; }

create_snapshot() {
    local ds name
    ds="$(pick "Select a dataset to snapshot:" < <(list_dataset_names))" || return 0
    [[ -z "$ds" ]] && return 0
    name="$(ask "Snapshot name for $ds:" "manual-$(ts)")" || return 0
    [[ -z "$name" ]] && { msg "No name entered." 8 50; return 0; }
    run_and_show "zfs snapshot $ds@$name" zfs snapshot "$ds@$name"
}

delete_snapshot() {
    local snap
    snap="$(pick "Select a snapshot to DESTROY:" < <(list_snapshots | cut -f1))" || return 0
    [[ -z "$snap" ]] && return 0
    yesno "Destroy snapshot:\n\n  $snap\n\nThis cannot be undone." || return 0
    run_and_show "zfs destroy $snap" zfs destroy "$snap"
}

rollback_snapshot() {
    local snap
    snap="$(pick "Select a snapshot to ROLL BACK to:" < <(list_snapshots | cut -f1))" || return 0
    [[ -z "$snap" ]] && return 0
    yesno "ROLL BACK to:\n\n  $snap\n\nThis DESTROYS all newer snapshots and ALL changes\nmade since then. Continue?" 13 72 || return 0
    run_and_show "zfs rollback -r $snap" zfs rollback -r "$snap"
}

menu_snapshots() {
    local choice
    while true; do
        choice=$(whiptail --title "$APP_TITLE — Snapshots" --notags \
            --menu "Manage ZFS snapshots:" 18 72 7 \
            view     "1) View all snapshots" \
            create   "2) Create a snapshot" \
            delete   "3) Delete a snapshot" \
            rollback "4) Roll back to a snapshot" \
            back     "5) Back" \
            3>&1 1>&2 2>&3) || break
        case "$choice" in
            view)     view_snapshots    ;;
            create)   create_snapshot   ;;
            delete)   delete_snapshot   ;;
            rollback) rollback_snapshot ;;
            back)     break ;;
        esac
    done
}

# ---------------------------------------------------------------------------
# Cron-backed schedules (snapshots + scrubs share one owner file)
# ---------------------------------------------------------------------------

# Emit existing snapshot schedules as: dataset<TAB>cron<TAB>keep
parse_snapshot_schedules() {
    [[ -f "$CRON_FILE" ]] || return 0
    awk -v s="$SNAPSHOT_SCRIPT" '
        /^[[:space:]]*#/ {next} /^[[:space:]]*$/ {next} $1 ~ /=/ {next}
        { idx=0; for (i=1;i<=NF;i++) if ($i==s) {idx=i; break}
          if (idx) printf "%s\t%s %s %s %s %s\t%s\n", $(idx+1),$1,$2,$3,$4,$5,$(idx+2) }
    ' "$CRON_FILE"
}

# Emit existing scrub schedules as: pool<TAB>cron
parse_scrub_schedules() {
    [[ -f "$CRON_FILE" ]] || return 0
    awk '
        /^[[:space:]]*#/ {next} /^[[:space:]]*$/ {next} $1 ~ /=/ {next}
        { idx=0; for (i=1;i<=NF;i++) if ($i=="scrub") {idx=i; break}
          if (idx && $(idx-1) ~ /zpool/) printf "%s\t%s %s %s %s %s\n", $(idx+1),$1,$2,$3,$4,$5 }
    ' "$CRON_FILE"
}

# Rewrite the whole cron file from the two TSV blobs ($1 snaps, $2 scrubs).
write_cron_file() {
    local snaps="$1" scrubs="$2" ds cron keep pool
    mkdir -p "$(dirname "$CRON_FILE")"
    {
        printf '%s\n\n' "$CRON_HEADER"
        if [[ -n "$snaps" ]]; then
            echo "# Snapshot schedules (auto-prune keeps last N)"
            while IFS=$'\t' read -r ds cron keep; do
                [[ -z "$ds" ]] && continue
                printf '%s root %s %s %s\n' "$cron" "$SNAPSHOT_SCRIPT" "$ds" "$keep"
            done <<<"$snaps"
            echo
        fi
        if [[ -n "$scrubs" ]]; then
            echo "# Scrub schedules"
            while IFS=$'\t' read -r pool cron; do
                [[ -z "$pool" ]] && continue
                printf '%s root /usr/sbin/zpool scrub %s\n' "$cron" "$pool"
            done <<<"$scrubs"
            echo
        fi
    } >"$CRON_FILE"
    chmod 644 "$CRON_FILE"
}

# Install the standalone snapshot+prune helper invoked by cron.
install_snapshot_script() {
    cat >"$SNAPSHOT_SCRIPT" <<'EOF'
#!/bin/bash
# Auto-generated by zfs-snapshot.sh. Creates an auto-* snapshot and prunes
# the oldest auto-* snapshots beyond the keep count.  Usage: <dataset> [keep]
DATASET="$1"; KEEP="${2:-7}"; PREFIX="auto"
[ -z "$DATASET" ] && { echo "Usage: $0 <dataset> [keep_count]"; exit 1; }

TS=$(date +%Y%m%d-%H%M%S)
/usr/sbin/zfs snapshot "${DATASET}@${PREFIX}-${TS}"

SNAPS=$(/usr/sbin/zfs list -H -t snapshot -o name -s creation "$DATASET" | grep "@${PREFIX}-")
COUNT=$(echo "$SNAPS" | grep -c . )
if [ "$COUNT" -gt "$KEEP" ]; then
  DEL=$((COUNT - KEEP))
  echo "$SNAPS" | head -n "$DEL" | while read -r S; do
    /usr/sbin/zfs destroy "$S"
  done
fi
EOF
    chmod 755 "$SNAPSHOT_SCRIPT"
}

# Cron-frequency presets shared by snapshot/scrub scheduling. Echoes a chosen
# 5-field cron expression, or empty on cancel.
pick_cron() {
    local kind="$1" choice custom
    if [[ "$kind" == "scrub" ]]; then
        choice=$(whiptail --title "$APP_TITLE" --notags \
            --menu "How often should the scrub run?" 16 72 5 \
            "0 3 1 * *"    "Monthly (1st @ 3am)" \
            "0 3 1,15 * *" "Bi-weekly (1st & 15th @ 3am)" \
            "0 3 * * 0"    "Weekly (Sunday 3am)" \
            "custom"       "Custom cron expression..." \
            3>&1 1>&2 2>&3) || return 1
    else
        choice=$(whiptail --title "$APP_TITLE" --notags \
            --menu "How often should the snapshot run?" 17 72 6 \
            "0 * * * *"   "Hourly" \
            "0 */6 * * *" "Every 6 hours" \
            "0 0 * * *"   "Daily at midnight" \
            "0 3 * * *"   "Daily at 3am" \
            "0 2 * * 0"   "Weekly (Sunday 2am)" \
            "custom"      "Custom cron expression..." \
            3>&1 1>&2 2>&3) || return 1
    fi
    if [[ "$choice" == "custom" ]]; then
        custom="$(ask "Enter a 5-field cron expression (min hour dom mon dow):" "0 3 * * *")" || return 1
        [[ -z "$custom" ]] && return 1
        echo "$custom"
    else
        echo "$choice"
    fi
}

_report_schedules() {
    local rows; rows="$(parse_snapshot_schedules)"
    [[ -z "$rows" ]] && { echo "(no snapshot schedules configured)"; return 0; }
    printf '%s\n' "$rows" | awk -F'\t' \
        'BEGIN{printf "%-40s %-16s %s\n","DATASET","CRON","KEEP"}
         {printf "%-40s %-16s %s\n",$1,$2,$3}'
}
view_schedules() { show "Snapshot Schedules" _report_schedules; }

add_schedule() {
    local ds cron keep snaps scrubs
    ds="$(pick "Schedule auto-snapshots for which dataset:" < <(list_dataset_names))" || return 0
    [[ -z "$ds" ]] && return 0
    cron="$(pick_cron snapshot)" || return 0
    keep="$(ask "Keep how many auto snapshots for $ds?" "7")" || return 0
    [[ "$keep" =~ ^[0-9]+$ ]] || keep=7

    install_snapshot_script
    snaps="$(parse_snapshot_schedules | grep -vP "^${ds}\t" || true)"
    snaps="$(printf '%s\n%s\t%s\t%s' "$snaps" "$ds" "$cron" "$keep")"
    scrubs="$(parse_scrub_schedules)"
    write_cron_file "$snaps" "$scrubs"
    msg "Schedule saved for:\n  $ds\n\nCron: $cron   Keep: $keep\n\nWritten to $CRON_FILE" 13 72
}

snapshot_now() {
    local row ds keep
    row="$(pick "Take an auto-snapshot now for which scheduled dataset:" \
            < <(parse_snapshot_schedules | cut -f1))" || return 0
    [[ -z "$row" ]] && return 0
    ds="$row"
    keep="$(parse_snapshot_schedules | awk -F'\t' -v d="$ds" '$1==d {print $3; exit}')"
    [[ "$keep" =~ ^[0-9]+$ ]] || keep=7
    install_snapshot_script
    run_and_show "snapshot now: $ds (keep $keep)" "$SNAPSHOT_SCRIPT" "$ds" "$keep"
}

remove_schedule() {
    local ds snaps scrubs
    ds="$(pick "Remove the snapshot schedule for which dataset:" \
            < <(parse_snapshot_schedules | cut -f1))" || return 0
    [[ -z "$ds" ]] && return 0
    yesno "Remove auto-snapshot schedule for:\n  $ds ?" || return 0
    snaps="$(parse_snapshot_schedules | grep -vP "^${ds}\t" || true)"
    scrubs="$(parse_scrub_schedules)"
    write_cron_file "$snaps" "$scrubs"
    msg "Removed schedule for $ds." 8 60
}

menu_schedules() {
    local choice
    while true; do
        choice=$(whiptail --title "$APP_TITLE — Schedules" --notags \
            --menu "Auto-snapshot schedules (cron):" 18 74 7 \
            view   "1) View schedules" \
            add    "2) Add / update a schedule" \
            now    "3) Snapshot a scheduled dataset NOW" \
            remove "4) Remove a schedule" \
            back   "5) Back" \
            3>&1 1>&2 2>&3) || break
        case "$choice" in
            view)   view_schedules  ;;
            add)    add_schedule    ;;
            now)    snapshot_now    ;;
            remove) remove_schedule ;;
            back)   break ;;
        esac
    done
}

# ---------------------------------------------------------------------------
# Datasets menu
# ---------------------------------------------------------------------------
_report_datasets() {
    local rows; rows="$(list_datasets)"
    [[ -z "$rows" ]] && { echo "(no datasets found — is 'zfs list' working as root?)"; return 0; }
    printf '%s\n' "$rows" | awk -F'\t' \
        'BEGIN{printf "%-40s %8s %8s  %s\n","NAME","USED","AVAIL","MOUNTPOINT"}
         {printf "%-40s %8s %8s  %s\n",$1,$2,$3,$4}'
}
view_datasets() { show "ZFS Datasets" _report_datasets; }

# ---------------------------------------------------------------------------
# LXC mountpoints menu (Proxmox pct)
# ---------------------------------------------------------------------------
have_pct() { command -v pct >/dev/null 2>&1; }

# Container IDs, one per line.
list_lxc_ids() { pct list 2>/dev/null | awk 'NR>1 {print $1}'; }

# Mountpoints for one container: slot<TAB>hostpath<TAB>ctpath
lxc_mountpoints() {
    local ctid="$1"
    pct config "$ctid" 2>/dev/null | awk '
        /^mp[0-9]+:/ {
            key=$1; sub(":","",key)
            line=$0; sub(/^mp[0-9]+:[[:space:]]*/,"",line)
            n=split(line,a,","); host=a[1]; target=""
            for (i=2;i<=n;i++) if (a[i] ~ /^mp=/) { target=substr(a[i],4) }
            printf "%s\t%s\t%s\n", key, host, target
        }'
}

_report_mounts() {
    local ctid slot host target found=0
    while IFS= read -r ctid; do
        [[ -z "$ctid" ]] && continue
        while IFS=$'\t' read -r slot host target; do
            [[ -z "$slot" ]] && continue
            printf 'CT %-6s %-6s %-30s -> %s\n' "$ctid" "$slot" "$host" "$target"
            found=1
        done < <(lxc_mountpoints "$ctid")
    done < <(list_lxc_ids)
    [[ "$found" -eq 0 ]] && echo "(no LXC mountpoints found)"
    return 0
}
view_mounts() { show "LXC Mountpoints" _report_mounts; }

add_mount() {
    local ctid host target used slot=0
    ctid="$(pick "Add a mountpoint to which container:" < <(list_lxc_ids))" || return 0
    [[ -z "$ctid" ]] && return 0
    host="$(ask "Host path on the Proxmox host (e.g. /RAID1_1TB/MyData):")" || return 0
    [[ -z "$host" ]] && return 0
    target="$(ask "Mount path inside the container (e.g. /data):")" || return 0
    [[ -z "$target" ]] && return 0

    used="$(lxc_mountpoints "$ctid" | sed -n 's/^mp\([0-9]\+\).*/\1/p')"
    while grep -qx "$slot" <<<"$used"; do slot=$((slot + 1)); done
    run_and_show "pct set $ctid -mp$slot" pct set "$ctid" "-mp${slot}" "${host},mp=${target}"
}

remove_mount() {
    local ctid slot
    ctid="$(pick "Remove a mountpoint from which container:" < <(list_lxc_ids))" || return 0
    [[ -z "$ctid" ]] && return 0
    slot="$(pick "Which mountpoint slot to remove from CT $ctid:" \
            < <(lxc_mountpoints "$ctid" | awk -F'\t' '{print $1" ("$2" -> "$3")"}' | sed 's/ .*//'))" || return 0
    [[ -z "$slot" ]] && return 0
    yesno "Remove $slot from container $ctid?\n\nThe container may need to be stopped for this to apply." || return 0
    run_and_show "pct set $ctid -delete $slot" pct set "$ctid" -delete "$slot"
}

menu_mounts() {
    if ! have_pct; then
        msg "'pct' not found — LXC mountpoint management needs Proxmox VE.\n\nThis menu is unavailable on this host." 10 72
        return 0
    fi
    local choice
    while true; do
        choice=$(whiptail --title "$APP_TITLE — LXC Mounts" --notags \
            --menu "LXC container bind-mountpoints (pct):" 17 74 6 \
            view   "1) View mountpoints" \
            add    "2) Add a mountpoint" \
            remove "3) Remove a mountpoint" \
            back   "4) Back" \
            3>&1 1>&2 2>&3) || break
        case "$choice" in
            view)   view_mounts  ;;
            add)    add_mount    ;;
            remove) remove_mount ;;
            back)   break ;;
        esac
    done
}

# ---------------------------------------------------------------------------
# Replication menu (zfs send | recv)
# ---------------------------------------------------------------------------
load_replication() { [[ -f "$REPLICATION_CONFIG" ]] && cat "$REPLICATION_CONFIG"; return 0; }

save_replication() {
    mkdir -p "$(dirname "$REPLICATION_CONFIG")"
    printf '%s' "$1" >"$REPLICATION_CONFIG"
}

# Standalone replication helper for cron / manual use.
install_replication_script() {
    cat >"$REPLICATION_SCRIPT" <<'EOF'
#!/bin/bash
# Auto-generated by zfs-snapshot.sh. Incrementally replicates a ZFS dataset.
# Usage: pve-zfs-replicate.sh <source-dataset> <target-dataset>
set -e
SOURCE="$1"; TARGET="$2"
[ -z "$SOURCE" ] || [ -z "$TARGET" ] && { echo "Usage: $0 <source> <target>"; exit 1; }

LATEST=$(/usr/sbin/zfs list -H -t snapshot -o name -s creation "$SOURCE" | tail -n1 | awk -F@ '{print $2}')
[ -z "$LATEST" ] && { echo "No snapshots on $SOURCE"; exit 1; }

if /usr/sbin/zfs list -H "$TARGET" >/dev/null 2>&1; then
  COMMON=""
  for s in $(/usr/sbin/zfs list -H -t snapshot -o name -s creation "$SOURCE" | awk -F@ '{print $2}' | tac); do
    if /usr/sbin/zfs list -H -t snapshot "${TARGET}@${s}" >/dev/null 2>&1; then COMMON="$s"; break; fi
  done
  [ -z "$COMMON" ] && { echo "No common snapshot between $SOURCE and $TARGET; aborting." >&2; exit 2; }
  if [ "$COMMON" = "$LATEST" ]; then echo "Up-to-date: $SOURCE@$COMMON"; exit 0; fi
  /usr/sbin/zfs send -R -I "${SOURCE}@${COMMON}" "${SOURCE}@${LATEST}" | /usr/sbin/zfs receive -F "$TARGET"
  echo "Incremental ${COMMON} -> ${LATEST}"
else
  /usr/sbin/zfs send -R "${SOURCE}@${LATEST}" | /usr/sbin/zfs receive "$TARGET"
  echo "Full -> ${LATEST}"
fi
EOF
    chmod 755 "$REPLICATION_SCRIPT"
}

_report_replication() {
    local src tgt common latest status found=0
    while IFS=$'\t' read -r src tgt; do
        [[ -z "$src" ]] && continue
        found=1
        latest="$(latest_snapshot "$src")"; latest="${latest:--}"
        common="$(latest_common_snapshot "$src" "$tgt")"; common="${common:--}"
        if ! dataset_exists "$tgt"; then status="target missing (full send needed)"
        elif [[ "$common" == "$latest" && "$common" != "-" ]]; then status="up-to-date"
        elif [[ "$common" == "-" ]]; then status="no common snap (mismatch)"
        else status="incremental pending"; fi
        printf '%-30s -> %-30s [%s]\n' "$src" "$tgt" "$status"
    done < <(load_replication)
    [[ "$found" -eq 0 ]] && echo "(no replication targets configured)"
    return 0
}
view_replication() { show "Replication Targets" _report_replication; }

add_replication() {
    local src tgt cfg
    src="$(pick "Replicate which source dataset:" < <(list_dataset_names))" || return 0
    [[ -z "$src" ]] && return 0
    tgt="$(ask "Target dataset (must be on an imported pool):" "backup/${src}")" || return 0
    [[ -z "$tgt" ]] && return 0
    cfg="$(load_replication)"
    if grep -qxF "$(printf '%s\t%s' "$src" "$tgt")" <<<"$cfg"; then
        msg "That target is already configured." 8 60; return 0
    fi
    cfg="$(printf '%s%s\t%s\n' "${cfg:+$cfg$'\n'}" "$src" "$tgt")"
    save_replication "$cfg"
    install_replication_script
    msg "Added replication target:\n  $src -> $tgt" 9 72
}

run_replication() {
    local pair src tgt
    pair="$(pick "Replicate which target now:" \
            < <(load_replication | awk -F'\t' '{print $1" -> "$2}'))" || return 0
    [[ -z "$pair" ]] && return 0
    src="${pair%% -> *}"; tgt="${pair##* -> }"
    install_replication_script
    msg "Replicating $src -> $tgt ...\n\nThis may take a while; press OK to start." 9 72
    run_and_show "replicate $src -> $tgt" "$REPLICATION_SCRIPT" "$src" "$tgt"
}

remove_replication() {
    local pair src tgt cfg
    pair="$(pick "Remove which replication target:" \
            < <(load_replication | awk -F'\t' '{print $1" -> "$2}'))" || return 0
    [[ -z "$pair" ]] && return 0
    src="${pair%% -> *}"; tgt="${pair##* -> }"
    yesno "Remove replication target?\n  $src -> $tgt\n\n(The target dataset itself is NOT destroyed.)" || return 0
    cfg="$(load_replication | grep -vxF "$(printf '%s\t%s' "$src" "$tgt")" || true)"
    save_replication "$cfg"
    msg "Removed: $src -> $tgt" 8 60
}

menu_replication() {
    local choice
    while true; do
        choice=$(whiptail --title "$APP_TITLE — Replication" --notags \
            --menu "Replicate datasets via zfs send | recv:" 18 74 7 \
            view   "1) View targets & status" \
            add    "2) Add a target" \
            run    "3) Replicate a target NOW" \
            remove "4) Remove a target" \
            back   "5) Back" \
            3>&1 1>&2 2>&3) || break
        case "$choice" in
            view)   view_replication   ;;
            add)    add_replication    ;;
            run)    run_replication    ;;
            remove) remove_replication ;;
            back)   break ;;
        esac
    done
}

# ---------------------------------------------------------------------------
# Scrub menu
# ---------------------------------------------------------------------------
_report_scrub() {
    local pool sched found=0
    while IFS= read -r pool; do
        [[ -z "$pool" ]] && continue
        found=1
        sched="$(parse_scrub_schedules | awk -F'\t' -v p="$pool" '$1==p {print $2; exit}')"
        printf '===== %s  (schedule: %s) =====\n' "$pool" "${sched:-none}"
        zpool status "$pool" 2>&1 | grep -E '^[[:space:]]*(state|scan|errors):' || true
        echo
    done < <(list_pools)
    [[ "$found" -eq 0 ]] && echo "(no pools found)"
    return 0
}
view_scrub() { show "Pool Scrub Status" _report_scrub; }

scrub_start() {
    local pool
    pool="$(pick "Start a scrub on which pool:" < <(list_pools))" || return 0
    [[ -z "$pool" ]] && return 0
    run_and_show "zpool scrub $pool" zpool scrub "$pool"
}

scrub_stop() {
    local pool
    pool="$(pick "Stop the running scrub on which pool:" < <(list_pools))" || return 0
    [[ -z "$pool" ]] && return 0
    run_and_show "zpool scrub -s $pool" zpool scrub -s "$pool"
}

scrub_schedule() {
    local pool cron snaps scrubs
    pool="$(pick "Schedule scrubs for which pool:" < <(list_pools))" || return 0
    [[ -z "$pool" ]] && return 0
    cron="$(pick_cron scrub)" || return 0
    snaps="$(parse_snapshot_schedules)"
    scrubs="$(parse_scrub_schedules | grep -vP "^${pool}\t" || true)"
    scrubs="$(printf '%s\n%s\t%s' "$scrubs" "$pool" "$cron")"
    write_cron_file "$snaps" "$scrubs"
    msg "Scrub scheduled for $pool.\n\nCron: $cron\nWritten to $CRON_FILE" 11 72
}

scrub_unschedule() {
    local pool snaps scrubs
    pool="$(pick "Remove the scrub schedule for which pool:" \
            < <(parse_scrub_schedules | cut -f1))" || return 0
    [[ -z "$pool" ]] && return 0
    yesno "Remove scrub schedule for pool $pool?" || return 0
    snaps="$(parse_snapshot_schedules)"
    scrubs="$(parse_scrub_schedules | grep -vP "^${pool}\t" || true)"
    write_cron_file "$snaps" "$scrubs"
    msg "Removed scrub schedule for $pool." 8 60
}

menu_scrub() {
    local choice
    while true; do
        choice=$(whiptail --title "$APP_TITLE — Scrub" --notags \
            --menu "Pool scrub management:" 18 72 8 \
            view       "1) View scrub status" \
            start      "2) Start a scrub" \
            stop       "3) Stop a running scrub" \
            schedule   "4) Schedule periodic scrubs" \
            unschedule "5) Remove a scrub schedule" \
            back       "6) Back" \
            3>&1 1>&2 2>&3) || break
        case "$choice" in
            view)       view_scrub       ;;
            start)      scrub_start      ;;
            stop)       scrub_stop       ;;
            schedule)   scrub_schedule   ;;
            unschedule) scrub_unschedule ;;
            back)       break ;;
        esac
    done
}

# ---------------------------------------------------------------------------
# Main menu loop
# ---------------------------------------------------------------------------
main_menu() {
    local choice
    while true; do
        choice=$(whiptail --title "$APP_TITLE" --notags \
            --menu "Select an area  (host: $(hostname))" 20 74 9 \
            snapshots   "1) Snapshots — list / create / delete / rollback" \
            schedules   "2) Schedules — cron auto-snapshots (keep last N)" \
            datasets    "3) Datasets — view ZFS datasets & usage" \
            mounts      "4) LXC Mounts — manage container bind-mounts" \
            replication "5) Replication — zfs send | recv to another dataset" \
            scrub       "6) Scrub — start / stop / schedule pool scrubs" \
            quit        "7) Quit" \
            3>&1 1>&2 2>&3) || break

        case "$choice" in
            snapshots)   menu_snapshots   ;;
            schedules)   menu_schedules   ;;
            datasets)    view_datasets    ;;
            mounts)      menu_mounts      ;;
            replication) menu_replication ;;
            scrub)       menu_scrub       ;;
            quit)        break ;;
        esac
    done
    clear
    echo "Done."
}

main_menu