#!/bin/bash

# Restores indices saved by backup-elastic-daily.
#
# Usage: restore-elastic [options] index [index...]
#        restore-elastic [options] -a
#
#   -a          restore every index found in the backup
#   -c          also restore cluster metadata (ingest pipelines, templates);
#               existing ones are kept unless -f is given
#   -d DIR      backup directory (default: $BACKUP_ROOT/elastic/current)
#   -f          delete and recreate indices (and templates) that already exist
#   -r NAME     restore a single index under a different name (aliases are
#               not restored in this case)
#   -u URL      Elasticsearch URL (overrides ES_URL from the config)
#
# Connection settings (ES_URL, ES_USER, ES_PASS, ES_API_KEY, ES_CA_CERT,
# ES_INSECURE) are read from /etc/backup-daily.conf, like the backup script.
# ES_BULK_DOCS (default 5000) sets how many documents are sent per _bulk,
# ES_RETRY (default 3) how many times a failed request is retried.

#######################################

usage() {
	sed -n '/^# Usage/,/^$/s/^# \{0,1\}//p' "$0" >&2
	exit 2
}

if ! [ -f /etc/backup-daily.conf ]; then
	echo "/etc/backup-daily.conf not found"
	exit 4
fi

BACKUP_ROOT=
ES_URL=http://localhost:9200
ES_USER=
ES_PASS=
ES_API_KEY=
ES_CA_CERT=
ES_INSECURE=0
ES_BULK_DOCS=5000
ES_RETRY=3

. /etc/backup-daily.conf

ALL=0
CLUSTER=0
FORCE=0
RENAME=
SRC_DIR=

while getopts "acd:fr:u:h" opt; do
	case $opt in
		a) ALL=1 ;;
		c) CLUSTER=1 ;;
		d) SRC_DIR=$OPTARG ;;
		f) FORCE=1 ;;
		r) RENAME=$OPTARG ;;
		u) ES_URL=$OPTARG ;;
		*) usage ;;
	esac
done
shift $((OPTIND - 1))

if [ -z "$SRC_DIR" ]; then
	[ -z $BACKUP_ROOT ] && exit 11
	SRC_DIR=$BACKUP_ROOT/elastic/current
fi
[ -d "$SRC_DIR" ] || { echo "$SRC_DIR not found" >&2; exit 21; }

for cmd in curl jq split; do
	command -v $cmd >/dev/null || { echo "$cmd not found"; exit 31; }
done

if [ $ALL = 1 ]; then
	[ $# -gt 0 ] && usage
	set -- $(cd "$SRC_DIR" && ls *.meta.json 2>/dev/null | sed 's/\.meta\.json$//')
fi
[ $# -eq 0 ] && [ $CLUSTER = 0 ] && usage
if [ -n "$RENAME" ] && [ $# -ne 1 ]; then
	echo "-r needs exactly one index" >&2
	exit 2
fi

#######################################

WORK=$(mktemp -d)
trap 'rm -rf "$WORK"' EXIT
chmod 700 $WORK

# curl config keeps credentials out of the process list
CURLCFG=$WORK/curl.cfg
curl_quote() { printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g'; }
{
	echo 'silent'
	echo 'show-error'
	echo 'fail'
	echo 'header = "Content-Type: application/json"'
	# transient failures only (timeouts, 408/429/5xx, connection refused);
	# bodies read from stdin are buffered by curl, so retries resend them
	echo "retry = $ES_RETRY"
	echo 'retry-connrefused'
	[ -n "$ES_USER" ] && echo "user = \"$(curl_quote "$ES_USER:$ES_PASS")\""
	[ -n "$ES_API_KEY" ] && echo "header = \"Authorization: ApiKey $(curl_quote "$ES_API_KEY")\""
	[ -n "$ES_CA_CERT" ] && echo "cacert = \"$(curl_quote "$ES_CA_CERT")\""
	[ "$ES_INSECURE" = 1 ] && echo 'insecure'
} > $CURLCFG

# es METHOD PATH [curl args...]
es() {
	local method=$1 path=$2
	shift 2
	curl -K "$CURLCFG" -X "$method" "${ES_URL%/}/$path" "$@"
}

exists() {
	curl -K "$CURLCFG" -o /dev/null "${ES_URL%/}/$1" 2>/dev/null
}

# decompress FILE - to stdout, by extension
decompress() {
	case $1 in
		*.gz) gzip -dc "$1" ;;
		*.zst) zstd -dcq "$1" ;;
	esac
}

uri() {
	jq -rn --arg s "$1" '$s | @uri'
}

# sends one chunk of _bulk lines from stdin, called by split --filter
# (some split implementations run two filters at once, hence mktemp)
cat > $WORK/bulk <<'EOF'
#!/bin/bash
RESP=$(mktemp -p "$WORK") || exit 1
trap 'rm -f "$RESP"' EXIT
curl -K "$CURLCFG" -X POST "$BULK_URL" --data-binary @- > $RESP || exit 1
jq -e '.errors == false' $RESP >/dev/null && exit 0
jq -r '[.items[][] | select(.error)] | .[0:5][] | "  \(._id): \(.error.reason)"' $RESP >&2
exit 1
EOF
chmod 700 $WORK/bulk

FAILED=

#######################################

# cluster metadata: pipelines first, then component and index templates

# put_all KIND FILE JQ_FILTER - JQ_FILTER yields {name, body} objects
put_all() {
	local kind=$1 file=$SRC_DIR/_cluster/$2 filter=$3 name
	[ -f "$file" ] || return 0
	jq -c "$filter" "$file" | while read -r item; do
		name=$(jq -r '.name' <<< "$item")
		if [ $FORCE = 0 ] && exists "$kind/$(uri "$name")"; then
			echo "$kind/$name exists, skipped"
			continue
		fi
		if jq -c '.body' <<< "$item" | es PUT "$kind/$(uri "$name")" -d @- >/dev/null; then
			echo "$kind/$name restored"
		else
			echo "$kind/$name failed" >&2
			echo "$kind/$name" >> $WORK/failed
		fi
	done
}

if [ $CLUSTER = 1 ]; then
	put_all _ingest/pipeline ingest_pipeline.json \
		'to_entries[] | {name: .key, body: .value}'
	put_all _component_template component_template.json \
		'(.component_templates // [])[] | {name, body: .component_template}'
	put_all _index_template index_template.json \
		'(.index_templates // [])[] | {name, body: .index_template}'
	put_all _template template.json \
		'to_entries[] | {name: .key, body: .value}'
	[ -f $WORK/failed ] && FAILED="$FAILED $(tr '\n' ' ' < $WORK/failed)"
fi

#######################################

restore_index() {
	local idx=$1 target=${RENAME:-$1}
	local META=$SRC_DIR/$idx.meta.json DATA= ext
	local utarget=$(uri "$target")
	local expected count

	for ext in zst gz; do
		[ -f "$SRC_DIR/$idx.ndjson.$ext" ] && DATA=$SRC_DIR/$idx.ndjson.$ext
	done
	if ! [ -f "$META" ] || [ -z "$DATA" ]; then
		echo "$idx: not found in $SRC_DIR" >&2
		return 1
	fi
	case $DATA in
		*.gz) command -v gzip >/dev/null || { echo "gzip not found" >&2; return 1; } ;;
		*.zst) command -v zstd >/dev/null || { echo "zstd not found" >&2; return 1; } ;;
	esac

	if exists "$utarget"; then
		if [ $FORCE = 0 ]; then
			echo "$target: index exists, use -f to replace it" >&2
			return 1
		fi
		es DELETE "$utarget" >/dev/null || return 1
	fi

	# create without replicas and refresh for a faster load
	jq --arg r "$RENAME" '
		.settings.index.number_of_replicas = "0" |
		.settings.index.refresh_interval = "-1" |
		if $r != "" then .aliases = {} else . end' "$META" \
		| es PUT "$utarget" -d @- >/dev/null || return 1

	# 2 lines per document
	decompress "$DATA" \
		| BULK_URL="${ES_URL%/}/$utarget/_bulk" WORK=$WORK CURLCFG=$CURLCFG \
		  split -l $((ES_BULK_DOCS * 2)) --filter="$WORK/bulk"
	local status=("${PIPESTATUS[@]}")
	[ "${status[0]}" = 0 ] && [ "${status[1]}" = 0 ] || return 1

	# put back the original (or default, if not set) values
	jq '{index: {
		number_of_replicas: .settings.index.number_of_replicas,
		refresh_interval: .settings.index.refresh_interval }}' "$META" \
		| es PUT "$utarget/_settings" -d @- >/dev/null || return 1

	es POST "$utarget/_refresh" >/dev/null || return 1

	expected=$(( $(decompress "$DATA" | wc -l) / 2 ))
	count=$(es GET "$utarget/_count" | jq '.count')
	if [ "$count" != "$expected" ]; then
		echo "$target: expected $expected documents, got $count" >&2
		return 1
	fi

	echo "$target: $count documents restored"
}

for idx in "$@"
do
	if ! restore_index "$idx"; then
		echo "${RENAME:-$idx}: restore failed" >&2
		FAILED="$FAILED ${RENAME:-$idx}"
	fi
done

if [ -n "$FAILED" ]; then
	echo "failed:$FAILED" >&2
	exit 1
fi

exit 0