#!/usr/bin/bash

PROGNAME=${0##*/}
set -o pipefail

[[ -f /etc/profile.d/logcenter-path.sh ]] &&
    source /etc/profile.d/logcenter-path.sh

function exit_usage() {
    local status=${1:-0}
    [[ $status != 0 ]] && exec >&2
    echo "\
Usage: $PROGNAME [OPTION...] action
Helper for maintenance operations of an Elasticsearch cluster

Available options:
    -E, --es-host       Set ES_HOST environement variable
    -h, --help          Display this help

Supported actions:
$(declare -f -F |sed -nre '/^declare -f action_/{ s,^[^ ]+ -f action_,    - ,; s,_,-,g; p; }')
"
    exit "$status"
}

function jfmt() {
    jq -c
}

function info() {
    echo "$(date +%Y-%m-%dT%H:%M:%S.%3N%:z) $*" >&2
}

# wait-until [-m max-wait-sec] command...
function wait-until() {
    local explain="$FUNCNAME ${*@Q}" try=0 max=0
    [[ $1 == -m ]] && { max=$2; shift 2; }
    while :; do
        (( try++ ))
        echo "[$try/${max:-0}] $explain"
        "$@" && return 0
        (( max > 0 && try >= max )) && return 1
        sleep 1
    done
}

# $1: name to lookup, default $HOSTNAME
# domain part is ignored
function is-node-present() {
    declare -a nodes
    local lookup=${1:-$HOSTNAME}
    lookup=${lookup,,}
    lookup=${lookup%%.*}
    nodes=( $(es-curl _cat/nodes?h=n) ) || return 1
    nodes=( "${nodes[@]%%.*}" )
    nodes=( "${nodes[@],,}" )
    local n
    for n in "${nodes[@]}"; do
        [[ $n == $lookup ]] && return 0
    done
    return 1
}

function is-cluster-green() {
    local status
    status=$(es-curl _cat/health?h=st) || return 1
    [[ $status == green ]] && return 0
    return 1
}

function action_status() {
    info '[status] Cluster settings'
    es-curl _cluster/settings |jfmt
}

function action_begin_rolling() {
    # do not touch indices.recovery.max_bytes_per_sec if a persistent
    # value is already defined
    local persistent_max=$(es-curl _cluster/settings |
        jq -r '.persistent.indices.recovery.max_bytes_per_sec//""')

    info '[begin-rolling] Apply transient cluster settings'
    jq -n --arg max "$persistent_max" '{"transient":
        ({"cluster.routing.rebalance.enable":"none"} +
         (if $max=="" then {"indices.recovery.max_bytes_per_sec":"125mb"} else {} end))
    }' |es-curl _cluster/settings -X PUT -d @- |jfmt
}

function action_end_rolling() {
    info '[end-rolling] Reset transient cluster settings'
    es-curl _cluster/settings -X PUT -d '{ "transient": {
            "cluster.routing.rebalance.enable": null,
            "indices.recovery.max_bytes_per_sec": null
        }}' |jfmt
}

function action_begin_node() {
    info '[begin-node] Apply transient cluster shard allocation constraint'
    es-curl _cluster/settings -X PUT -d '{ "transient": {
            "cluster.routing.allocation.enable": "primaries"
        }}' |jfmt
    info '[begin-node] Issue flush on indices'
    es-curl _flush -XPOST |jfmt
}

function action_end_node() {
    info '[end-node] Reset transient cluster shard allocation'
    es-curl _cluster/settings -X PUT -d '{ "transient": {
            "cluster.routing.allocation.enable": null
        }}' |jfmt
}

# es-maintenance wait-node [-n $HOSTNAME] [-m 300]
function action_wait_node() {
    local name= max=
    while (( $# > 0 )); do
        case "$1" in
            -n|--name) name=$2; shift ;;
            -m|--max) max=$2; shift ;;
            *) return 1 ;;
        esac
        shift
    done
    [[ -z $name ]] && name=$HOSTNAME
    info "[wait-node] Wait for $name to join cluster"
    wait-until -m "$max" is-node-present "$name"
}

# es-maintenance wait-green [-m 300]
function action_wait_green() {
    local max=
    while (( $# > 0 )); do
        case "$1" in
            -m|--max) max=$2; shift ;;
            *) return 1 ;;
        esac
        shift
    done
    info '[wait-green] Wait for cluster green state'
    wait-until -m "$max" is-cluster-green
}

args=()
while (( $# > 0 )); do
    case "$1" in
        -h|--help) exit_usage 0 ;;
        -E|--es-host) export ES_HOST=$2; shift ;;
        *) args+=( "$1" ) ;;
    esac
    shift
done

[[ -z $args ]] && exit_usage 1
action_fn=action_${args//[^[:alnum:]]/_}
declare -f -F "$action_fn" > /dev/null || exit_usage 1
"$action_fn" "${args[@]:1}"
