#!/usr/bin/env sh

nodes_file="nodes.txt"
### nodes.txt
# <origin host>   <user> <ssh port> <timezone> <compressed flag>
# <mirror 1 host> <user> <ssh port> <zfs root>
# <mirror 2 host> <user> <ssh port> <zfs root>
# ...

datasets_file="datasets.txt"
### datasets.txt
# <dataset 1>
# <dataset 2>
# <dataset 3>
# ...

if [ ! -f $nodes_file ]; then echo "No nodes specified."; exit 2; fi
if [ ! -f $datasets_file ]; then echo "No datasets specified."; exit 3; fi

ZFS_SEND_OPTS="-R"
ZFS_SNAPSHOT_OPTS="-r"
ZFS_RECV_OPTS="-svu"

read LOCAL_HOST LOCAL_USER LOCAL_PORT LOCAL_TZ SEND_COMPRESSED < $nodes_file
case $SEND_COMPRESSED in
    [tTyY]*) ZFS_SEND_OPTS="$ZFS_SEND_OPTS -c" ;;
esac

ORIGIN_PERMS="send,snapshot,hold"
MIRROR_PERMS="compression,receive,create"

SEND_DATASETS=""
SNAPSHOT_DATASETS=""
while read dataset dataset_options; do
    case $dataset in
        \#*) continue;;
    esac

    snapshot_dataset="yes"
    send_dataset="yes"
    case $dataset_options in 
        *"nosnap"* ) snapshot_dataset="";;
    esac
    case $dataset_options in 
        *"nosend"* ) send_dataset="";;
    esac

    if [ $send_dataset ]; then
        SEND_DATASETS="$SEND_DATASETS $dataset"
    fi
    if [ $snapshot_dataset ]; then
        SNAPSHOT_DATASETS="$SNAPSHOT_DATASETS $dataset"
    fi
done < datasets.txt

list_datasets()
{
    for dataset in $SEND_DATASETS; do
        zfs list -H $dataset
    done
}

list_remote_datasets() {
    while read addr user ssh_port zroot; do
        echo "=== $user@$addr $zroot ==="
        if [ $BG_TEST_MODE = "TRUE" ]; then
            echo ssh -i ~/.ssh/id_rsa -o port=$ssh_port $user@$addr\
                "sh -c 'for dataset in $SEND_DATASETS;
                    do zfs list -H $zroot/\$dataset;
                done'"
        else
            ssh -i ~/.ssh/id_rsa -o port=$ssh_port $user@$addr \
                "sh -c 'for dataset in $SEND_DATASETS;
                    do zfs list -H $zroot/\$dataset;
                done';"
        fi
    done < sendnodes.tmp
}

get_latest_snapshot()
{
    zfs list -t snapshot $dataset | tail -n1 | cut -d'@' -f2 | cut -d' ' -f 1
}

get_latest_remote_snapshot()
{
    ssh -i ~/.ssh/id_rsa -o port=$ssh_port $user@$addr \
        "zfs list -t snapshot $zroot/$dataset    |
            tail -n1|cut -d'@' -f2|cut -d' ' -f 1"
}

zfs_allow_origin()
{
    for dataset in $SEND_DATASETS; do
        if [ $BG_TEST_MODE = "TRUE" ]; then
            echo zfs allow -u $LOCAL_USER $ORIGIN_PERMS $dataset
        else
            zfs allow -u $LOCAL_USER $ORIGIN_PERMS $dataset
        fi
    done
}

zfs_allow_mirrors()
{
    while read addr user ssh_port zroot; do
        echo "=== $user@$addr $zroot ==="
        if [ $BG_TEST_MODE = "TRUE" ]; then
            echo ssh -i ~/.ssh/id_rsa -o port=$ssh_port $user@$addr \
                "for dataset in $SEND_DATASETS; do
                    zfs allow -u $user $MIRROR_PERMS $zroot/\$dataset
                done"
        else
            ssh -i ~/.ssh/id_rsa -o port=$ssh_port $user@$addr \
                "for dataset in $SEND_DATASETS; do
                    zfs allow -u $user $MIRROR_PERMS $zroot/\$dataset
                done"
        fi
    done < sendnodes.tmp
}

create_snapshots()
{
    snapshot_id="`date -I`-`date +%s`-$LOCAL_TZ"
    for dataset in $SNAPSHOT_DATASETS; do
        if [ $BG_TEST_MODE = "TRUE" ]; then
            echo zfs snapshot $ZFS_SNAPSHOT_OPTS $dataset@$snapshot_id
        else    
	        zfs snapshot $ZFS_SNAPSHOT_OPTS $dataset@$snapshot_id
        fi
    done
}

send_latest()
{
    while read addr user ssh_port zroot; do
        echo "=== $user@$addr $zroot ==="
        for dataset in $SEND_DATASETS; do
            origin_snapshot=`get_latest_snapshot $dataset`
            if [ $BG_TEST_MODE = "TRUE" ]; then
                echo zfs send $ZFS_SEND_OPTS $dataset@$origin_snapshot  \| \
                    ssh -i ~/.ssh/id_rsa -o port=$ssh_port $user@$addr \
                    zfs recv $ZFS_RECV_OPTS $zroot/$dataset
            else
                zfs send $ZFS_SEND_OPTS $dataset@$origin_snapshot |
                    ssh -i ~/.ssh/id_rsa -o port=$ssh_port \
                    $user@$addr                            \
                    zfs recv $ZFS_RECV_OPTS $zroot/$dataset
            fi
        done
    done < sendnodes.tmp
}

send_increment()
{
    while read addr user ssh_port zroot; do
        echo "=== $user@$addr $zroot ==="
        for dataset in $SEND_DATASETS; do
            origin_snapshot=`get_latest_snapshot`
            remote_snapshot=`get_latest_remote_snapshot`
            if [ $BG_TEST_MODE = "TRUE" ]; then
                echo zfs send $ZFS_SEND_OPTS     \
                    -i $dataset@$remote_snapshot \
                    $dataset@$origin_snapshot    \| \
                        ssh -i ~/.ssh/id_rsa -o port=$ssh_port $user@$addr  \
                            zfs recv $ZFS_RECV_OPTS $zroot/$dataset
            else
                zfs send $ZFS_SEND_OPTS          \
                    -i $dataset@$remote_snapshot \
                    $dataset@$origin_snapshot    |
                        ssh -i ~/.ssh/id_rsa -o port=$ssh_port $user@$addr  \
                            zfs recv $ZFS_RECV_OPTS $zroot/$dataset
            fi
        done
    done < sendnodes.tmp
}


################################### MAIN ####################################

### Options
set -- `getopt "tl:" "$@"` || {
    echo "Usage: `basename $0` [-tl] [command]" 1>&2
    exit 1
}
BG_TEST_MODE=FALSE
while :; do
    case "$1" in
        -t)
            BG_TEST_MODE=TRUE
            echo Test mode.
        ;;
        --) shift; break ;;
    esac
    shift
done

tail -n +2 $nodes_file|grep -v '^#' > sendnodes.tmp

### Command
case "$1" in
    list) list_datasets ;;
    list-remote) list_remote_datasets ;;
    snapshot) create_snapshots ;;
    configure-mirrors) zfs_allow_mirrors ;;
    configure-origin) zfs_allow_origin ;;
    mirror-reset) send_latest ;;
    mirror-increment) send_increment ;;
    #mirror-prune) send_increment_with_F ;;
esac

rm sendnodes.tmp