#!/bin/bash

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

set -e

CURRENT_PATH="$(
	cd "$(dirname "${BASH_SOURCE[0]}")"
	pwd
)"
OUTPUT_PATH="$(readlink -f "${CURRENT_PATH}/../../output")"

readonly CURRENT_PATH
readonly OUTPUT_PATH

function log() {
	local level="${1}"
	local message="${2}"
	local date
	date="$(date +'%Y-%m-%d %H:%M:%S')"
	if [[ "${level}" == 'INFO' ]]; then
		level="[\033[32;1m ${level}  \033[0m]"
	elif [[ "${level}" == 'WARNING' ]]; then
		level="[\033[33;1m${level}\033[0m]"
	elif [[ "${level}" == 'ERROR' ]]; then
		level="[\033[31;1m ${level} \033[0m]"
	fi
	echo -e "${level} ${date} - ${message}"
}

function log_info() {
	local message="${1}"
	log 'INFO' "${message}"
}

function log_warning() {
	local message="${1}"
	log 'WARNING' "${message}"
}

function log_error() {
	local message="${1}"
	log 'ERROR' "${message}"
	exit 1
}

function help() {
	cat >&2 <<EOF
${BASH_SOURCE[0]} start|stop|clean [OPTIONS ...]

    start -n <NUM> -l <LIBRARY_PATH> -p <BASE_PORT>

             Start the FE cluster.
      -n     The number of FEs.
      -l     The FE library path (default: doris/output/fe/lib)
      -p     The base port to generate all needed ports (default: 9030).

    stop     Stop the FE cluster.

    clean    Clean the data (rm -rf "\$(pwd)"/fe*).
EOF
	exit 1
}

function parse_options() {
	NUMBER_INSTANCES=1
	PORT=9010
	LIBS_PATH="${OUTPUT_PATH}/fe/lib"

	ACTION="${1}"
	if ! shift 1; then
		help
	fi

	local option
	while getopts "+n:p:l:" option; do
		case "${option}" in
		n)
			NUMBER_INSTANCES="${OPTARG}"
			;;
		p)
			PORT="${OPTARG}"
			;;
		l)
			LIBS_PATH="$(readlink -f "${OPTARG}")"
			;;
		*)
			help
			;;
		esac
	done

	readonly NUMBER_INSTANCES
	readonly PORT
	readonly LIBS_PATH
}

function prepare() {
	local port="${1}"
	local doris_home="${2}"
	local log_dir="${3}"
	local conf_dir="${doris_home}/conf"
	local meta_dir="${doris_home}/doris-meta"
	local http_port=8030

	mkdir -p "${doris_home}"
	mkdir -p "${log_dir}"
	mkdir -p "${conf_dir}"
	mkdir -p "${meta_dir}"

	cat >"${conf_dir}/fe.conf" <<EOF
LOG_DIR = ${log_dir}
meta_dir = ${meta_dir}

edit_log_port = $((port + id))
rpc_port = $((port + 10 + id))
query_port = $((port + 20 + id))
http_port = $((http_port + id))

sys_log_level = INFO
sys_log_verbose_modules = org.apache.doris

enable_ssl = false
priority_networks = 127.0.0.1/32
EOF
}

function mysql_command() {
	local command="${1}"
	local query_port
	query_port="$(grep 'query_port' "${CURRENT_PATH}/fe1/conf/fe.conf" | sed 's/query_port[[:space:]]*=[[:space:]]*//')"
	mysql -h127.0.0.1 -P"${query_port}" -uroot -e "${command}" -s -N
}

function create_role_and_version() {
	local id="${1}"
	local http_port
	local header
	local meta_dir

	http_port="$(grep 'http_port' "${CURRENT_PATH}/fe1/conf/fe.conf" | sed 's/http_port[[:space:]]*=[[:space:]]*//')"
	while true; do
		mysql_command "ALTER SYSTEM ADD FOLLOWER \"127.0.0.1:$((PORT + id))\"" >/dev/null 2>&1 || true
		if header="$(curl --noproxy '127.0.0.1' --silent --head --fail -H "CLIENT_NODE_HOST:127.0.0.1" -H "CLIENT_NODE_PORT:$((PORT + id))" "http://127.0.0.1:${http_port}/check")"; then
			break
		fi
		log_info "Waiting for the Master FE to be ready ..."
		sleep 1
	done
	meta_dir="$(grep 'meta_dir' "${CURRENT_PATH}/fe${id}/conf/fe.conf" | sed 's/meta_dir[[:space:]]*=[[:space:]]*//')"

	mkdir -p "${meta_dir}/image"
	echo "${header}" | sed -n '/cluster_id/,/token/p' | sed '{s/cluster_id: /clusterId=/;s/token: /token=/;}' \
		>"${meta_dir}/image/VERSION"

	cat >"${meta_dir}/image/ROLE" <<EOF
role=FOLLOWER
name=$(
		set -e
		mysql_command "SHOW FRONTENDS" | grep "127.0.0.1_$((PORT + id))" | awk '{print $1}'
	)
EOF
}

function start_fe() {
	local id="${1}"
	local port="${2}"
	local libs_path="${3}"
	local doris_home
	local helper

	doris_home="$(pwd)/fe${id}"
	local pid_dir="${doris_home}"
	local log_dir="${doris_home}/log"

	if [[ -f "${pid_dir}/fe.pid" ]]; then
		if xargs kill -0 <"${pid_dir}/fe.pid"; then
			log_warning "FE(${id}) is running..."
			return
		fi
	fi

	prepare "${port}" "${doris_home}" "${log_dir}"

	declare -a classpath=("${libs_path}")
	for lib in "${libs_path}"/*.jar; do
		classpath+=("${lib}")
	done

	if [[ "${id}" -gt 1 ]]; then
		helper="-helper 127.0.0.1:$((PORT + 1))"
		create_role_and_version "${id}"
	fi

	pushd "${doris_home}" >/dev/null
	DORIS_HOME="${doris_home}" \
		PID_DIR="${doris_home}" \
		nohup java -cp "$(
		IFS=:
		echo "${classpath[*]}"
	)" org.apache.doris.DorisFE ${helper:+${helper}} >>"${log_dir}/fe.out" 2>&1 &
	if kill -0 $!; then
		log_info "Start FE(${id}) successfully"
	else
		log_warning "Failed to start FE(${id})"
	fi
	popd >/dev/null
}

function start() {
	if [[ -z "${LIBS_PATH}" ]]; then
		help
	fi

	log_info "NUMBER:       ${NUMBER_INSTANCES}"
	log_info "BASE PORT:    ${PORT}"
	log_info "LIBRARY_PATH: ${LIBS_PATH}"

	for i in $(seq "${NUMBER_INSTANCES}"); do
		start_fe "${i}" "${PORT}" "${LIBS_PATH}"
	done
}

function stop() {
	local pid
	while read -r file; do
		pid="$(<"${file}")"
		if kill "${pid}"; then
			log_info "Stopped FE (${pid}) successfully"
		else
			log_warning "Failed to stop FE (${pid})"
		fi
	done < <(find "$(pwd)"/fe* -name 'fe.pid')
}

function clean() {
	rm -rf "$(pwd)"/fe*
	log_info "Clean all data successfully"
}

function main() {
	parse_options "${@}"

	case "${ACTION}" in
	start) start ;;
	stop) stop ;;
	clean) clean ;;
	*)
		help
		;;
	esac
}

main "${@}"