#!/bin/bash

# Usage:
#  CONCURRENCY=8 ./docker-run-random.sh

set -euo pipefail

CONCURRENCY="${CONCURRENCY:-8}"
# Tend to use less CPU to generate more process context switches.
CPU="${CPU:-4}"
MEMORY="${MEMORY:-5g}"
DOCKER_STOP_TIMEOUT="${DOCKER_STOP_TIMEOUT:-10}"

TMP_PATH=""
LOG_PATH=""
PATH_SUFFIX=""
TESTNAME="all"
TIDB_VERSION="v8.5.3"
PD_VERSION="v8.5.0"
TIKV_VERSION=""
TICDC_VERSION="v8.5"
REBUILD_IMAGE=1
declare -a RUN_ARGS
RUN_ARGS=()
KEEP_TMP_ON_ERROR=0
MEMORY_PROFILE=0

print_help_option() {
	printf '  %-*s %s\n' 50 "$1" "$2"
}

show_help() {
	echo "Usage: $0 [OPTIONS]"
	echo "OPTIONS:"
	print_help_option "--help" "Display this message"
	print_help_option "--tmp-path <temporary path>" "Set the path for temporary data generated during testing"
	print_help_option "--log-path <log path>" "Set the path for logs"
	print_help_option "--path-with-suffix" "Add git commit and timestamp to the path as suffix"
	print_help_option "--test <all/with_tidb/upgrade/replication/cdc>" "Set the name of test case to run"
	print_help_option "--tidb-version <v8.5.3/...>" "Set the TiDB version for with_tidb/replication/upgrade image build"
	print_help_option "--pd-version <v8.5.0/...>" "Set the PD version for with_tidb/replication/upgrade image build"
	print_help_option "--tikv-version <git_hash>" "Set the version of old TiKV for \"upgrade\" test"
	print_help_option "--ticdc-version <v8.5/...>" "Set the TiCDC version for cdc image build"
	print_help_option "--no-rebuild-image" "Do NOT rebuild the testing Docker image"
	print_help_option "--keep-tmp-on-error" "Keep temporary data on error for debugging"
	print_help_option "--memory-profile" "Enable memory profiling"
	print_help_option "--serverless <0/1>" "Set FOR_SERVERLESS for random tests (default: 1)"
	print_help_option "--test-duration 600" "Set the duration of test in seconds"
	print_help_option "--cdc-keyspace-count <N>" "Set CDC keyspace count (default: 2)"
	print_help_option "--enable-columnar-multi-column-fulltext-workload" "Force multi-column FTS workload"
	print_help_option "--disable-columnar-multi-column-fulltext-workload" "Force legacy single-column FTS workload"
	print_help_option "--disable-cmek" "Disable CMEK in random tests"
}

PWD=$(pwd)

while [[ $# -gt 0 ]]; do
	case "$1" in
	--tmp-path)
		TMP_PATH="$2"
		shift
		;;
	--log-path)
		LOG_PATH="$2"
		shift
		;;
	--path-with-suffix)
		PATH_SUFFIX=-$(date +%Y%m%d-%H%M%S)-$(git rev-parse --short HEAD)
		;;
	--test)
		TESTNAME="$2"
		shift
		;;
	--test-duration)
		RUN_ARGS+=("--test-duration" "$2")
		shift
		;;
	--serverless)
		RUN_ARGS+=("--serverless" "$2")
		shift
		;;
	--tidb-version)
		TIDB_VERSION="$2"
		shift
		;;
	--pd-version)
		PD_VERSION="$2"
		shift
		;;
	--tikv-version)
		TIKV_VERSION="$2"
		shift
		;;
	--ticdc-version)
		TICDC_VERSION="$2"
		shift
		;;
	--no-rebuild-image)
		REBUILD_IMAGE=0
		;;
	--keep-tmp-on-error)
		KEEP_TMP_ON_ERROR=1
		RUN_ARGS+=("--keep-tmp-on-error")
		;;
	--memory-profile)
		MEMORY_PROFILE=1
		RUN_ARGS+=("--memory-profile")
		;;
	--no-tiflash)
		RUN_ARGS+=("--no-tiflash")
		;;
	--no-txn-file)
		RUN_ARGS+=("--no-txn-file")
		;;
	--txn-file-pessimistic)
		RUN_ARGS+=("--txn-file-pessimistic" "$2")
		shift
		;;
	--no-remote-cop)
		RUN_ARGS+=("--no-remote-cop")
		;;
	--no-tpc)
		RUN_ARGS+=("--no-tpc")
		;;
	--tpcc-txns-threshold)
		RUN_ARGS+=("--tpcc-txns-threshold" "$2")
		shift
		;;
	--no-jepsen)
		RUN_ARGS+=("--no-jepsen")
		;;
	--jepsen-no-txn-file)
		RUN_ARGS+=("--jepsen-no-txn-file")
		;;
	--jepsen-txns-threshold)
		RUN_ARGS+=("--jepsen-txns-threshold" "$2")
		shift
		;;
	--load-data-task-timeout-sec)
		RUN_ARGS+=("--load-data-task-timeout-sec" "$2")
		shift
		;;
	--unique-workload)
		RUN_ARGS+=("--unique-workload")
		;;
	--columnar-workload)
		RUN_ARGS+=("--columnar-workload")
		;;
	--no-restart-tso-svc)
		RUN_ARGS+=("--no-restart-tso-svc")
		;;
	--ia-table-ratio)
		RUN_ARGS+=("--ia-table-ratio" "$2")
		shift
		;;
	--async-commit-ratio)
		RUN_ARGS+=("--async-commit-ratio" "$2")
		shift
		;;
	--no-check-backup-ts)
		RUN_ARGS+=("--no-check-backup-ts")
		;;
	--oss-chaos-ratio)
		RUN_ARGS+=("--oss-chaos-ratio" "$2")
		shift
		;;
	--upgrade-test-duration)
		RUN_ARGS+=("--upgrade-test-duration" "$2")
		shift
		;;
	--disable-value-cache)
		RUN_ARGS+=("--disable-value-cache")
		;;
	--enable-tiflash-write-node)
		RUN_ARGS+=("--enable-tiflash-write-node")
		;;
	--disable-columnar-normal-workload)
		RUN_ARGS+=("--disable-columnar-normal-workload")
		;;
	--disable-columnar-dynamic-workload)
		RUN_ARGS+=("--disable-columnar-dynamic-workload")
		;;
	--disable-columnar-partition-workload)
		RUN_ARGS+=("--disable-columnar-partition-workload")
		;;
	--disable-columnar-vector-workload)
		RUN_ARGS+=("--disable-columnar-vector-workload")
		;;
	--disable-columnar-fulltext-workload)
		RUN_ARGS+=("--disable-columnar-fulltext-workload")
		;;
	--enable-columnar-multi-column-fulltext-workload)
		RUN_ARGS+=("--enable-columnar-multi-column-fulltext-workload")
		;;
	--disable-columnar-multi-column-fulltext-workload)
		RUN_ARGS+=("--disable-columnar-multi-column-fulltext-workload")
		;;
	--enable-blob)
		RUN_ARGS+=("--enable-blob")
		;;
	--disable-cmek)
		RUN_ARGS+=("--disable-cmek")
		;;
	--tidb-next-gen)
		RUN_ARGS+=("--tidb-next-gen")
		;;
	--rep-debug-log)
		RUN_ARGS+=("--rep-debug-log")
		;;
	--rep-lazy-load-shards-ratio)
		RUN_ARGS+=("--rep-lazy-load-shards-ratio" "$2")
		shift
		;;
	--rep-cdc-newarch)
		RUN_ARGS+=("--rep-cdc-newarch")
		;;
	--cdc-keyspace-count)
		RUN_ARGS+=("--cdc-keyspace-count" "$2")
		shift
		;;
	--help)
		show_help
		exit 0
		;;
	*)
		show_help
		exit 1
		;;
	esac
	shift
done

if [ "$MEMORY_PROFILE" -eq 1 ] && [ "$KEEP_TMP_ON_ERROR" -ne 1 ]; then
	echo "WARNING: --keep-tmp-on-error is not enabled. The profile dumps will be removed after each test."
fi

declare -a BUILD_IMAGE_ARGS
BUILD_IMAGE_ARGS=()
if [ "$REBUILD_IMAGE" -eq 1 ]; then
	BUILD_IMAGE_ARGS+=(--pull --no-cache)
fi

IMAGE="amazonlinux:2023.5.20241001.1"
if [[ "$TESTNAME" =~ ^(with_tidb|replication|cdc)$ ]]; then
	declare -a TIDB_BUILD_ARGS
	TIDB_BUILD_ARGS=(
		--build-arg TIDB_VERSION="$TIDB_VERSION"
		--build-arg PD_VERSION="$PD_VERSION"
		--build-arg TICDC_VERSION="$TICDC_VERSION"
	)
	docker build "${BUILD_IMAGE_ARGS[@]}" -t random-tidb \
		"${TIDB_BUILD_ARGS[@]}" \
		- <Dockerfile.tidb
	IMAGE="random-tidb"
elif [ "$TESTNAME" = "upgrade" ]; then
	if [ -z "$TIKV_VERSION" ]; then
		echo "ERROR: --tikv-version is required for upgrade test"
		exit 1
	fi
	docker build "${BUILD_IMAGE_ARGS[@]}" -t random-upgrade \
		--build-arg TIDB_VERSION="$TIDB_VERSION" \
		--build-arg PD_VERSION="$PD_VERSION" \
		--build-arg TIKV_VERSION="$TIKV_VERSION" \
		- <Dockerfile.upgrade
	IMAGE="random-upgrade"
elif [ "$MEMORY_PROFILE" -eq 1 ]; then
	# Parse profile dumps in an environment different with container for random test would fail to translate the addresses to symbols.
	# So provide an image with tools.
	docker build "${BUILD_IMAGE_ARGS[@]}" -t random-profile - <Dockerfile.profile
	IMAGE="random-profile"
fi

declare -a LOG_VOLUME_ARGS
LOG_VOLUME_ARGS=()
if [ -n "$LOG_PATH" ]; then
	ABSOLUTE_LOG_PATH=$(readlink -f "${LOG_PATH}${PATH_SUFFIX}")
	LOG_VOLUME_ARGS=(-v "${ABSOLUTE_LOG_PATH}:/random-logs")
	RUN_ARGS+=("--log-path" "/random-logs")
fi

for ((i = 0; i < "$CONCURRENCY"; i++)); do
	declare -a TMP_VOLUME_ARGS
	TMP_VOLUME_ARGS=()
	if [ -n "$TMP_PATH" ]; then
		# Containers should use different $TMPDIR, otherwise they will conflict on TIKV_LOCK_FILES.
		mkdir -p "${TMP_PATH}${PATH_SUFFIX}/${i}"
		ABSOLUTE_TMP_PATH=$(readlink -f "${TMP_PATH}${PATH_SUFFIX}/$i")
		TMP_VOLUME_ARGS=(-v "$ABSOLUTE_TMP_PATH:/random-tmp")
	fi

	docker run --name random-"$TESTNAME"-"$i" --cpus="$CPU" --memory="$MEMORY" \
		--stop-timeout="$DOCKER_STOP_TIMEOUT" \
		--env "DOCKER_STOP_TIMEOUT=$DOCKER_STOP_TIMEOUT" \
		-itd \
		-v "$PWD":/random \
		"${TMP_VOLUME_ARGS[@]}" \
		"${LOG_VOLUME_ARGS[@]}" \
		"$IMAGE" \
		/bin/bash /random/run-random.sh "$i" "$TESTNAME" "${RUN_ARGS[@]}"
done