#!/usr/bin/env bash
#
# ep-zpool-balancer.sh
#
# Rebalances data across all vdevs of a ZFS pool (raidz2/raidz3 or any other
# vdev layout) after new disks/vdevs have been added.
#
# Background: ZFS only spreads *new* writes across all vdevs in a pool.
# Blocks written before new vdevs were added stay put, so old data keeps
# living on the original (now more full / more used) vdevs while new vdevs
# stay empty. This script forces a rebalance by rewriting every file
# in place: copy -> verify checksum -> atomic rename. The rewritten data
# is allocated fresh by ZFS across the whole pool, including new vdevs.
#
# Usage:
#   ep-zpool-balancer.sh -p /pool/dataset [options]
#   ep-zpool-balancer.sh -P poolname [options]     # rebalance every mounted dataset in the pool
#
# See --help for full option list.

set -uo pipefail

# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------
SCRIPT_NAME="$(basename "$0")"
TARGET_PATHS=()
POOL_NAME=""
DRY_RUN=0
VERBOSE=0
INCLUDE_HARDLINKS=0
MIN_FREE_PERCENT=10
STATE_FILE=""
RESET_STATE=0
EXCLUDES=()
FORCE_CLEAN_TMP=0
LOG_FILE=""

TOTAL_FILES=0
TOTAL_BYTES=0
DONE_FILES=0
DONE_BYTES=0
SKIPPED_FILES=0
ERROR_FILES=0
START_TIME=$(date +%s)

TMP_SUFFIX=".ep-zpool-balancer.tmp"

# ---------------------------------------------------------------------------
# Logging helpers
# ---------------------------------------------------------------------------
log() {
    local level="$1"; shift
    local msg="$*"
    local ts
    ts="$(date '+%Y-%m-%d %H:%M:%S')"
    local line="[$ts] [$level] $msg"
    echo "$line"
    if [[ -n "$LOG_FILE" ]]; then
        echo "$line" >> "$LOG_FILE"
    fi
}
info()  { log "INFO"  "$@"; }
warn()  { log "WARN"  "$@"; }
err()   { log "ERROR" "$@" >&2; }
debug() { [[ "$VERBOSE" -eq 1 ]] && log "DEBUG" "$@" || true; }

die() {
    err "$@"
    exit 1
}

usage() {
    cat <<EOF
$SCRIPT_NAME - rebalance a ZFS raidz2/raidz3 pool after adding disks

USAGE:
  $SCRIPT_NAME -p PATH [PATH ...] [options]
  $SCRIPT_NAME -P POOLNAME [options]

TARGET SELECTION (one required):
  -p, --path PATH          Directory (inside a ZFS dataset) to rebalance.
                            Can be given multiple times.
  -P, --pool POOLNAME      Rebalance every mounted dataset of this pool.

OPTIONS:
  -m, --min-free PERCENT   Abort/skip if pool free space would drop below
                            this percentage while rewriting a file.
                            Default: $MIN_FREE_PERCENT
  -s, --state-file PATH    File used to track already-rebalanced files so
                            the run can be safely resumed. Default:
                            <target>/.ep-zpool-balancer-state
  -r, --reset-state        Ignore/clear existing state file and start fresh.
  -e, --exclude PATTERN    Shell glob pattern to exclude (matched against
                            the full file path). Can be given multiple times.
  -H, --include-hardlinks  Also rewrite files with more than one hardlink.
                            Default: such files are skipped, since rewriting
                            them in place would break the hardlink group.
  -l, --log-file PATH      Also append log output to this file.
  -n, --dry-run            Show what would be done without changing anything.
  -c, --clean-tmp          Remove leftover *.ep-zpool-balancer.tmp files from
                            a previous interrupted run before starting.
  -v, --verbose            Verbose/debug output.
  -h, --help               Show this help.

EXAMPLES:
  # Rebalance one dataset, dry-run first
  $SCRIPT_NAME -p /tank/data -n
  $SCRIPT_NAME -p /tank/data

  # Rebalance every dataset in pool "tank", keep 15% free
  $SCRIPT_NAME -P tank -m 15

  # Resume an interrupted run
  $SCRIPT_NAME -p /tank/data

NOTES:
  * The script rewrites files in place (copy, verify checksum, atomic
    rename), so it needs headroom equal to roughly the largest file size
    on the target dataset. Use -m/--min-free to keep a safety margin.
  * Progress is tracked in a state file so the script can be interrupted
    (Ctrl-C) and safely resumed later.
  * Files with multiple hardlinks are skipped by default (see -H).
EOF
}

# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------
while [[ $# -gt 0 ]]; do
    case "$1" in
        -p|--path)
            [[ $# -ge 2 ]] || die "$1 requires an argument"
            TARGET_PATHS+=("$2"); shift 2 ;;
        -P|--pool)
            [[ $# -ge 2 ]] || die "$1 requires an argument"
            POOL_NAME="$2"; shift 2 ;;
        -m|--min-free)
            [[ $# -ge 2 ]] || die "$1 requires an argument"
            MIN_FREE_PERCENT="$2"; shift 2 ;;
        -s|--state-file)
            [[ $# -ge 2 ]] || die "$1 requires an argument"
            STATE_FILE="$2"; shift 2 ;;
        -r|--reset-state)
            RESET_STATE=1; shift ;;
        -e|--exclude)
            [[ $# -ge 2 ]] || die "$1 requires an argument"
            EXCLUDES+=("$2"); shift 2 ;;
        -H|--include-hardlinks)
            INCLUDE_HARDLINKS=1; shift ;;
        -l|--log-file)
            [[ $# -ge 2 ]] || die "$1 requires an argument"
            LOG_FILE="$2"; shift 2 ;;
        -n|--dry-run)
            DRY_RUN=1; shift ;;
        -c|--clean-tmp)
            FORCE_CLEAN_TMP=1; shift ;;
        -v|--verbose)
            VERBOSE=1; shift ;;
        -h|--help)
            usage; exit 0 ;;
        *)
            die "Unknown option: $1 (see --help)" ;;
    esac
done

if [[ ${#TARGET_PATHS[@]} -eq 0 && -z "$POOL_NAME" ]]; then
    usage
    die "Either -p/--path or -P/--pool is required"
fi

if ! [[ "$MIN_FREE_PERCENT" =~ ^[0-9]+$ ]]; then
    die "--min-free must be an integer percentage"
fi

for cmd in zfs zpool find stat sha256sum awk df numfmt; do
    command -v "$cmd" >/dev/null 2>&1 || die "Required command not found: $cmd"
done

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

# Resolve the ZFS pool name that a given directory lives on.
path_to_pool() {
    local path="$1"
    local fstype src
    fstype="$(df -PT "$path" 2>/dev/null | awk 'NR==2{print $2}')"
    [[ "$fstype" == "zfs" ]] || return 1
    src="$(df -P "$path" 2>/dev/null | awk 'NR==2{print $1}')"
    echo "${src%%/*}"
}

# Return available bytes and pool capacity percent used, for a pool.
pool_free_percent() {
    local pool="$1"
    zfs list -Hp -o avail,used "$pool" 2>/dev/null | awk '{
        avail=$1; used=$2; total=avail+used;
        if (total <= 0) { print 0; exit }
        printf "%d", (avail / total) * 100
    }'
}

pool_avail_bytes() {
    local pool="$1"
    zfs list -Hp -o avail "$pool" 2>/dev/null
}

is_excluded() {
    local path="$1" pattern
    for pattern in "${EXCLUDES[@]:-}"; do
        [[ -n "$pattern" ]] || continue
        # shellcheck disable=SC2053
        [[ "$path" == $pattern ]] && return 0
    done
    return 1
}

human() {
    numfmt --to=iec-i --suffix=B "$1" 2>/dev/null || echo "${1}B"
}

state_has() {
    local key="$1"
    [[ -n "$STATE_FILE" && -f "$STATE_FILE" ]] || return 1
    grep -qxF "$key" "$STATE_FILE" 2>/dev/null
}

state_add() {
    local key="$1"
    [[ "$DRY_RUN" -eq 1 || -z "$STATE_FILE" ]] && return 0
    echo "$key" >> "$STATE_FILE"
}

interrupted=0
on_interrupt() {
    interrupted=1
    warn "Interrupted. Finishing current file safely, then stopping..."
}
trap on_interrupt INT TERM

print_summary() {
    local elapsed=$(( $(date +%s) - START_TIME ))
    info "----------------------------------------------------------------"
    info "Summary: processed=$DONE_FILES skipped=$SKIPPED_FILES errors=$ERROR_FILES"
    info "Data rewritten: $(human "$DONE_BYTES") of $(human "$TOTAL_BYTES") total"
    info "Elapsed: ${elapsed}s"
    info "----------------------------------------------------------------"
}

# ---------------------------------------------------------------------------
# Core: rebalance a single file
# ---------------------------------------------------------------------------
rebalance_file() {
    local file="$1"
    local pool="$2"

    if is_excluded "$file"; then
        debug "Excluded: $file"
        return 0
    fi

    if state_has "$file"; then
        debug "Already done (state file): $file"
        SKIPPED_FILES=$((SKIPPED_FILES + 1))
        return 0
    fi

    if [[ ! -f "$file" || -L "$file" ]]; then
        debug "Skipping non-regular/symlink: $file"
        return 0
    fi

    local nlink size
    nlink="$(stat -c '%h' "$file" 2>/dev/null || echo 1)"
    size="$(stat -c '%s' "$file" 2>/dev/null || echo 0)"

    if [[ "$INCLUDE_HARDLINKS" -eq 0 && "$nlink" -gt 1 ]]; then
        warn "Skipping (hardlinked x$nlink, use -H to include): $file"
        SKIPPED_FILES=$((SKIPPED_FILES + 1))
        return 0
    fi

    local avail
    avail="$(pool_avail_bytes "$pool")"
    if [[ -z "$avail" ]]; then
        err "Could not determine free space for pool '$pool'; skipping: $file"
        ERROR_FILES=$((ERROR_FILES + 1))
        return 0
    fi

    # Require headroom for the copy itself plus the configured free-space floor.
    local needed=$(( size + (size / 10) + 1 ))
    local free_pct
    free_pct="$(pool_free_percent "$pool")"
    if (( avail < needed )) || (( free_pct < MIN_FREE_PERCENT )); then
        warn "Skipping (insufficient free space, ${free_pct}% free < ${MIN_FREE_PERCENT}%): $file"
        SKIPPED_FILES=$((SKIPPED_FILES + 1))
        return 0
    fi

    if [[ "$DRY_RUN" -eq 1 ]]; then
        info "[dry-run] Would rebalance: $file ($(human "$size"))"
        DONE_FILES=$((DONE_FILES + 1))
        DONE_BYTES=$((DONE_BYTES + size))
        return 0
    fi

    local tmp="${file}${TMP_SUFFIX}.$$"
    debug "Rewriting: $file -> $tmp"

    if ! cp --preserve=all -- "$file" "$tmp" 2>/dev/null; then
        err "Copy failed, leaving original untouched: $file"
        rm -f -- "$tmp"
        ERROR_FILES=$((ERROR_FILES + 1))
        return 0
    fi

    local sum_orig sum_tmp
    sum_orig="$(sha256sum -- "$file" | awk '{print $1}')"
    sum_tmp="$(sha256sum -- "$tmp" | awk '{print $1}')"

    if [[ "$sum_orig" != "$sum_tmp" || -z "$sum_orig" ]]; then
        err "Checksum mismatch, aborting this file (original left intact): $file"
        rm -f -- "$tmp"
        ERROR_FILES=$((ERROR_FILES + 1))
        return 0
    fi

    if ! mv -f -- "$tmp" "$file"; then
        err "Atomic replace failed, cleaning up tmp file: $file"
        rm -f -- "$tmp"
        ERROR_FILES=$((ERROR_FILES + 1))
        return 0
    fi

    state_add "$file"
    DONE_FILES=$((DONE_FILES + 1))
    DONE_BYTES=$((DONE_BYTES + size))
    debug "Rebalanced: $file ($(human "$size"))"
    return 0
}

clean_stale_tmp_files() {
    local target="$1"
    local found
    found="$(find "$target" -xdev -type f -name "*${TMP_SUFFIX}.*" 2>/dev/null)"
    [[ -z "$found" ]] && return 0

    if [[ "$FORCE_CLEAN_TMP" -eq 1 ]]; then
        warn "Removing leftover temp files under $target"
        find "$target" -xdev -type f -name "*${TMP_SUFFIX}.*" -print -delete 2>/dev/null \
            | while read -r f; do debug "Removed stale tmp: $f"; done
    else
        warn "Found leftover temp files from a previous interrupted run under $target."
        warn "Re-run with --clean-tmp to remove them automatically, or delete manually."
    fi
}

# ---------------------------------------------------------------------------
# Rebalance one target directory
# ---------------------------------------------------------------------------
rebalance_path() {
    local target
    target="$(realpath -e "$1" 2>/dev/null)" || die "Path does not exist: $1"
    [[ -d "$target" ]] || die "Not a directory: $target"

    local pool
    pool="$(path_to_pool "$target")" || die "Not on a ZFS filesystem: $target"
    info "Target: $target (pool: $pool)"

    local free_pct
    free_pct="$(pool_free_percent "$pool")"
    info "Pool '$pool' free space: ${free_pct}%"
    if (( free_pct < MIN_FREE_PERCENT )); then
        die "Pool '$pool' free space (${free_pct}%) is already below --min-free (${MIN_FREE_PERCENT}%); aborting."
    fi

    local state="$STATE_FILE"
    if [[ -z "$state" ]]; then
        state="$target/.ep-zpool-balancer-state"
    fi
    if [[ "$RESET_STATE" -eq 1 && -f "$state" ]]; then
        info "Resetting state file: $state"
        rm -f -- "$state"
    fi
    STATE_FILE="$state"
    [[ "$DRY_RUN" -eq 0 ]] && : > "${STATE_FILE}.lock" 2>/dev/null && rm -f "${STATE_FILE}.lock"

    clean_stale_tmp_files "$target"

    info "Scanning files under $target ..."
    local list
    list="$(mktemp)"
    trap 'rm -f "$list"' RETURN
    find "$target" -xdev -type f ! -name "$(basename "$state")" ! -name "*${TMP_SUFFIX}.*" -print0 > "$list.nul" 2>/dev/null
    mv "$list.nul" "$list"

    local file_count=0
    local total_bytes=0
    while IFS= read -r -d '' f; do
        file_count=$((file_count + 1))
        total_bytes=$(( total_bytes + $(stat -c '%s' "$f" 2>/dev/null || echo 0) ))
    done < "$list"

    TOTAL_FILES=$((TOTAL_FILES + file_count))
    TOTAL_BYTES=$((TOTAL_BYTES + total_bytes))
    info "Found $file_count files ($(human "$total_bytes")) to consider."

    local processed=0
    while IFS= read -r -d '' f; do
        if [[ "$interrupted" -eq 1 ]]; then
            warn "Stopping early due to interrupt."
            break
        fi
        processed=$((processed + 1))
        rebalance_file "$f" "$pool"
        if (( processed % 100 == 0 )); then
            info "Progress: $processed/$file_count files in this target ($(human "$DONE_BYTES") rewritten so far overall)"
        fi
    done < "$list"

    rm -f "$list"
    trap - RETURN
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
main() {
    if [[ -n "$POOL_NAME" ]]; then
        zpool list -H "$POOL_NAME" >/dev/null 2>&1 || die "Pool not found: $POOL_NAME"
        local mountpoints
        mountpoints="$(zfs list -Hr -o mountpoint "$POOL_NAME" 2>/dev/null | grep -v -E '^(none|legacy|-)$')"
        [[ -n "$mountpoints" ]] || die "No mounted datasets found for pool: $POOL_NAME"
        while IFS= read -r mp; do
            [[ -n "$mp" ]] || continue
            TARGET_PATHS+=("$mp")
        done <<< "$mountpoints"
    fi

    info "Starting ep-zpool-balancer (dry-run=$DRY_RUN, min-free=${MIN_FREE_PERCENT}%)"
    for p in "${TARGET_PATHS[@]}"; do
        [[ "$interrupted" -eq 1 ]] && break
        STATE_FILE_SAVED="$STATE_FILE"
        rebalance_path "$p"
        # Reset per-target state file override so next path uses its own default.
        STATE_FILE="$STATE_FILE_SAVED"
    done

    print_summary
    [[ "$ERROR_FILES" -gt 0 ]] && exit 2
    [[ "$interrupted" -eq 1 ]] && exit 130
    exit 0
}

main "$@"