#!/bin/bash
#
# ZFS Snapshot Management Script
# Author: (C) 2024 Thorsten Spille <thorsten@spille-edv.de>

__version__=v1.0.0

# Function to display usage information
usage() {
    echo "Usage: $0 -d <dataset> -k <keep> -l <label> [-f <dateformat>] [-v] [-h <host>] [-p <port>]"
    echo "  -d: Name of the ZFS dataset"
    echo "  -k: Number of snapshots to keep"
    echo "  -l: Label for the snapshot"
    echo "  -f: (Optional) Date format to append to the snapshot name (default: '%Y-%m-%d-%H%M')"
    echo "  -v: (Optional) Enable verbose/debug logging"
    echo "  -h: (Optional) Remote host for SSH management"
    echo "  -p: (Optional) SSH port (default: 22)"
    exit 1
}

# Logging function
function log() {
    echo -e "$(date +'%b %d %T') $1"
}

# Debug logging function
function debug() {
    if [[ "$verbose" == true ]]; then
        echo -e "$(date +'%b %d %T') [DEBUG] $1"
    fi
}

# Default date format
default_dateformat="%Y-%m-%d-%H%M"
verbose=false
ssh_host=""
ssh_port=22

# Parse command-line arguments
while getopts "d:k:l:f:vh:p:" opt; do
    case "$opt" in
        d) dataset="$OPTARG" ;;
        k) keep="$OPTARG" ;;
        l) label="$OPTARG" ;;
        f) dateformat="$OPTARG" ;;
        v) verbose=true ;;
        h) ssh_host="$OPTARG" ;;
        p) ssh_port="$OPTARG" ;;
        *) usage ;;
    esac
done

# Validate required parameters
if [[ -z "$dataset" || -z "$keep" || -z "$label" ]]; then
    usage
fi

# Set dateformat to default if not provided
if [[ -z "$dateformat" ]]; then
    dateformat="$default_dateformat"
fi

# Validate "keep" parameter is a number
if ! [[ "$keep" =~ ^[0-9]+$ ]]; then
    log "Error: 'keep' must be a positive integer."
    exit 1
fi

# Base command for ZFS operations
zfs_cmd="zfs"
if [[ -n "$ssh_host" ]]; then
    # Test SSH connection before proceeding
    debug "Testing SSH connection to $ssh_host on port $ssh_port"
    ssh -q -o PasswordAuthentication=no -p $ssh_port $ssh_host exit
    if [[ $? -ne 0 ]]; then
        log "Error: Unable to connect to $ssh_host on port $ssh_port. Ensure SSH access is configured correctly."
        exit 1
    fi
    zfs_cmd="ssh -p $ssh_port $ssh_host zfs"
    debug "Using remote host: $ssh_host (port $ssh_port)"
fi

debug "Dataset: $dataset"
debug "Keep: $keep"
debug "Label: $label"
debug "Date format: $dateformat"
debug "Verbose mode: $verbose"

# Generate the snapshot name
snapshot_name="${dataset}@${label}-$(date -u +"$dateformat")"
debug "Generated snapshot name: $snapshot_name"

# Create the ZFS snapshot
log "Creating snapshot: $snapshot_name"
$zfs_cmd snapshot "$snapshot_name"
if [[ $? -ne 0 ]]; then
    log "Error: Failed to create snapshot."
    exit 1
fi

# List existing snapshots for the dataset
debug "Listing existing snapshots for dataset: $dataset"
existing_snapshots=$($zfs_cmd list -t snapshot -o name -s createtxg "$dataset" | grep "^${dataset}@${label}-")
debug "Existing snapshots:\n$existing_snapshots"

# Count the total number of snapshots matching the label
snapshot_count=$(echo "$existing_snapshots" | wc -l)
debug "Snapshot count: $snapshot_count"

# Check if pruning is necessary
if [[ "$snapshot_count" -gt "$keep" ]]; then
    # Calculate the number of snapshots to delete
    delete_count=$((snapshot_count - keep))
    debug "Number of snapshots to delete: $delete_count"

    # Get the snapshots to delete (oldest first)
    snapshots_to_delete=$(echo "$existing_snapshots" | head -n "$delete_count")
    debug "Snapshots to delete:\n$snapshots_to_delete"

    # Delete the old snapshots
    log "Deleting $delete_count old snapshot(s):"
    for snapshot in $snapshots_to_delete; do
        log "Deleting snapshot: $snapshot"
        $zfs_cmd destroy "$snapshot"
        if [[ $? -ne 0 ]]; then
            log "Error: Failed to delete snapshot $snapshot."
        fi
    done
else
    log "No snapshots to delete. Total snapshots ($snapshot_count) <= keep ($keep)."
fi

log "Snapshot management completed."
exit 0
