#!/bin/bash

PROGNAME=${0##*/}
ES_SCROLL_PARAMS='size=1000&scroll=1m'
TMPFILE=

# Keep track of scroll ids to close them on exit
declare -A SCROLL_IDS=()

function exit_usage() {
    local status=${1:-0}
    [[ $status != 0 ]] && exec >&2
    echo "\
Usage: $PROGNAME [OPTION...] INDEX [QUERY] [PARAMS]
Search elasticsearch with a scroll until there is no more hit

Available options:
    -v, --verbose       Be verbose, prints on stderr
    -h, --help          Display this help

This script performs a scroll _search query on the given index.
QUERY is the data to POST for the query in the initial request.
If PARAM is given, it gets appended to the default parameters
of the initial search query, which are $ES_SCROLL_PARAMS"
    exit "$status"
}

function perr() {
    echo "$PROGNAME: $*" >&2
}

function _curl() {
    if [[ -n $VERBOSE ]]; then
        perr "Call es-curl $*"
    fi
    es-curl "$@"
}

function on_exit() {
    local i
    if [[ -n $TMPFILE && -e $TMPFILE ]]; then
        rm ${VERBOSE:+-v} "$TMPFILE" >&2
    fi
    for i in "${!SCROLL_IDS[@]}"; do
        _curl "/_search/scroll" -d "{\"scroll_id\":\"$i\"}" -s -o /dev/null -X DELETE
    done
}

# $1: index
# $2: query to POST
# $3: extra _search parameters
# $4: internal, scroll id
function es-scroll() {
    local index=$1 query=$2 params=$3 scroll_id=$4
    local meta

    if [[ -n $scroll_id ]]; then
        # continue scroll
        [[ $VERBOSE ]] && perr "Continue scroll $scroll_id"
        _curl "/_search/scroll" -d "{\"scroll\":\"1m\",\"scroll_id\":\"$scroll_id\"}" -s -o "$TMPFILE"
    else
        # initial query
        [[ $VERBOSE ]] && perr 'Send initial query'
        _curl "$index/_search?$ES_SCROLL_PARAMS${params:+&$params}" -d "$query" -s -o "$TMPFILE"
    fi
    [[ $? != 0 ]] && { perr "Error, cURL failed"; return 1; }
    [[ -s $TMPFILE ]] || { perr "Error, empty response"; return 1; }
    meta=($(jq -r '"\(._scroll_id)\n\(.hits.hits |length)"' "$TMPFILE"))
    [[ $? != 0 || ${#meta[@]} != 2 ]] && { perr "Error, jq failed to extract metadata"; return 1; }
    [[ -z ${meta[0]} ]] && { perr "Error, empty scroll id"; return 1; }
    SCROLL_IDS["${meta[0]}"]=1
    cat "$TMPFILE"
    [[ ${meta[1]} == 0 ]] && return 0
    "$FUNCNAME" "$index" "$query" "$params" "${meta[0]}"
}

ARGS=()
while (( $# > 0 )); do
    case "$1" in
        -v|--verbose) VERBOSE=1 ;;
        -h|--help) exit_usage 0 ;;
        --) shift; break ;;
        -*) exit_usage 1 ;;
        *) ARGS+=( "$1" ) ;;
    esac
    shift
done
ARGS+=( "$@" )

[[ -z ${ARGS[0]} ]] && exit_usage 1

trap on_exit EXIT
TMPFILE=$(TMPDIR=${TMPDIR:-/dev/shm} mktemp --suffix ".$PROGNAME")
es-scroll "${ARGS[@]:0:3}"
