#! /bin/bash

export PATH=$PATH:/sbin

FILTER='.'
USER='root'
SEND='send'
RECV='recv -e -F'
FLOCK='/usr/bin/flock -w 60 -n 9'
PORT='1337'
MBUFFER='-s 128k -m 1G -4'
USE_MBUFFER='no'
LOCK_DIR='/var/run/'$(basename $0)
test -d $LOCK_DIR || mkdir $LOCK_DIR
TEST=0
VERB=0

test -f $HOME/.keychain/$HOSTNAME-sh && . $HOME/.keychain/$HOSTNAME-sh

function zfssendrecv {
    local OPTIND OPTARG opt
    local zfs
    local snap
    local dest
    local inc=''
    local hold=0
    while getopts "z:s:d:i:I:h" opt; do
        case $opt in
            z)
                zfs=$OPTARG
                ;;
            s)
                snap=$OPTARG
                ;;
            d)
                dest=$OPTARG
                ;;
            i)
                inc="-i ${OPTARG}"
                ;;
            I)
                inc="-I ${OPTARG}"
                ;;
            h)
                hold=1
                ;;
        esac
    done
    shift $((OPTIND-1))
    local zfs_normal=$( echo $zfs|sed 's/[\:\|\/\\ ]/_/g' )
    local lock="${LOCK_DIR}/${zfs_normal}.lock"
    local zfs_recv_status
    local zfs_send_status
    local pipe_status
    (
        if ! $FLOCK; then
            $ulog "unable to lock ${lock}"
            return -2
        fi
        if [[ $TEST == 0 && $hold == 1 ]]; then
            zfs hold -r zfsrecvman $snap 2>&1 | $ulog
            local hold_status="${PIPESTATUS[0]}"
            if [[ $hold_status != 0 ]]; then
                $ulog "unable to place a hold on our snapshots: ${snap}"
                return -3
            fi
        fi
        $ulog "estimating size of sending ${snap}"
        local size=$( zfs $SEND -v -n $inc $snap 2>&1 | tail -1 | cut -d" " -f 5 )
        # could be 0 or 400 or 4K or 9.3g, etc.
        local suf=$( echo $size | sed -E 's/[0-9]+\.?[0-9]*//' | tr '[:lower:]' '[:upper:]' )
        size=$( echo $size | sed -E 's/[pPtTgGmMkKB]$//' ) # remove known suffixes
	if [[ $suf != 'B' ]]; then
        	size=$( echo "${size} * 1024" | bc | sed -E 's/\.[0-9]+//' ) # use bc to multiply decimals, sed to make ceil()
	fi
        case $suf in
            B)
		suf=''
		;;
            K)
                suf=''
                ;;
            M)
                suf='K'
                ;;
            G)
                suf='M'
                ;;
            T)
                suf='G'
                ;;
            P)
                suf='T'
                ;;
        esac
        $ulog "estimated size of sending ${snap} is ${size}${suf}"
        local pv_more="-s ${size}${suf}"
        if [[ $USE_MBUFFER == 'yes' ]]; then
            ssh "${USER}@${REMOTE}" "mbuffer ${MBUFFER} -q -I ${PORT} | zfs ${RECV} ${dest}" 2>&1 | $ulog &
            sleep 5
            zfs $SEND $inc $snap 2> >($ulog)|
            mbuffer $MBUFFER $MBUFFER_SEND_OPTS -O ${REMOTE}:${PORT}
            zfs_send_status="${PIPESTATUS[0]}"
            $ulog "zfs send exited with status: ${zfs_send_status}"
            $ulog "about to wait on zfs send (this may take a while and appear to have hung)"
            wait
            zfs_recv_status="${PIPESTATUS[0]}"
            $ulog "zfs recv exited with status: ${zfs_recv_status}"
        else
            zfs $SEND $inc $snap 2> >($ulog) | pv $PV_OPTS $pv_more | ssh "${USER}@${REMOTE}" "zfs ${RECV} ${dest}" 2>&1 | $ulog
            pipe_status=("${PIPESTATUS[@]}")
            zfs_send_status="${pipe_status[0]}"
            zfs_recv_status="${pipe_status[2]}"
            $ulog "zfs send exited with status: ${zfs_send_status}"
            $ulog "zfs recv exited with status: ${zfs_recv_status}"
        fi
        if [[ $zfs_send_status != 0 ]]; then
            return $zfs_send_status
        elif [[ $zfs_recv_status != 0 ]]; then
            return $zfs_recv_status
        else
            # both must be zero
            return 0
        fi
    ) 9>$lock
}

function terminal_options {
if [ -t 1 ]; then
    ISTERM=1
    LOGGER_EXTRA='-s' # enable logger output to stderr
    PV_OPTS='-perb' # enable all the magic output from pv
    MBUFFER_SEND_OPTS='' # don't do quiet mode, we have a term
    ulog="logger ${LOGGER_EXTRA} -p user.notice -t "$(basename $0 2>/dev/null)"[${$}]"
else
    ISTERM=0
    LOGGER_EXTRA='' # don't enable stderr output
    PV_OPTS='-q' # make pv quiet
    MBUFFER_SEND_OPTS='-q' # enable send side -q, no terminal
    ulog="logger ${LOGGER_EXTRA} -p user.notice -t "$(basename $0 2>/dev/null)"[${$}]"
fi
}

terminal_options

while getopts "p:f:L:mnvr:u:d:" opt; do
    case $opt in
        p)
            PORT=$OPTARG
            ;;
        f)
            FILTER=$OPTARG
            ;;
        L)
            PV_OPTS="${PV_OPTS} -L ${OPTARG}"
            ;;
        m)
            USE_MBUFFER='yes'
            ;;
        n)
            RECV="${RECV} -n"
            TEST=1
            VERB=1
            PV_OPTS='-q' # make pv quiet
            MBUFFER_SEND_OPTS='-q' # enable send side -q, no terminal
            ;;
        v)
            VERB=1
            ;;
        r)
            REMOTE=$OPTARG
            ;;
        u)
            USER=$OPTARG
            ;;
        d)
            DEST=$OPTARG
            ;;
            
    esac
done
if [[ "${REMOTE}Z" == 'Z' ]]; then
    echo 'must set remote with -r option'
    exit 1;
fi
if [[ "${DEST}Z" == 'Z' ]]; then
    echo 'must set dest with -d option'
    exit 1;
fi

if [[ $VERB == 1 ]]; then
    echo $RECV | grep -q -- -v || RECV="${RECV} -v"
fi

function set_options {
    local zfs=$1
    local target=$2
    local compress=$( zfs get -H -o value compression $zfs )
    local destroy=$( zfs get -H -o value zfssnapman:destroy $zfs )
    ssh "${USER}@${REMOTE}" "zfs set readonly=on ${target} ; zfs set compression=${compress} ${target} ; zfs set zfssnapman:destroy=${destroy} ${target}"
}

# zfs is the filesystem to be sent, including source pool
for zfs in $( zfs list -H -t filesystem,volume -o name,zfssendman:send | grep true | egrep "${FILTER}" | awk -F\\t '{if($2 == "true") {print $1;}}' | sort ); do
    # base is the zfs filesystem to send, minus the pool
    target="${DEST}/${zfs}"
    target_dir=$( dirname $target )
    # recv_last is the last snapshot on the recv side of this zfs
    if ! ssh "${USER}@${REMOTE}" zfs get written $target >/dev/null 2>/dev/null; then
        $ulog sending initial snapshot of $zfs to $target_dir
        zfssendrecv -z $zfs \
            -s $( zfs list -t snapshot -o name -s creation -d 1 -H $zfs | head -1 ) \
            -d $target_dir
        if [[ $TEST == 0 ]]; then
            set_options $zfs $target
        fi
        sleep 5
    fi
    
    recv_last=$( ssh "${USER}@${REMOTE}" zfs list -t snapshot -o name -s creation -d 1 -H $target | tail -1 ) || break;
    if ! echo $recv_last | grep -q "${zfs}@"; then
        $ulog "no snapshot on distant end, you must destroy the filesystem: $zfs"
        break;
    fi
    # send is the snapshot on the recv side after stripping off the DEST
    send=$( echo $recv_last|sed "s|$DEST/||" )
    
    zfs list -t snapshot -o name -d 1 -H $zfs | grep -q $send
    if [[ $? == 0 ]]; then
        # most recent snapshot on the send side
        current=$( zfs list -t snapshot -o name -s creation -d 1 -H $zfs | tail -1 )
        if [[ $send == $current ]]; then
            $ulog "${zfs} is in sync on source and destination (${target})"
        else
            $ulog sending $send through $current to $target
            zfssendrecv -z $zfs \
                -I $send \
                -s $current \
                -d $target_dir
            if [[ $? == 0 ]]; then
                $ulog "$zfs is in sync on source and destination"
                if [[ $TEST == 0 ]]; then
                    set_options $zfs $target
                fi
            else
                $ulog zfs exited with $? while sending $send through $current to $target
            fi
        fi
    else
        $ulog "the most recent snapshot ($recv_last) on the recv side does not exist on the send side ($send)"
    fi
done
