#!/bin/bash

set -eu

DOCKER_STOP_TIMEOUT=${DOCKER_STOP_TIMEOUT:-10}
SHUTDOWN_EXIT_MARGIN=2
if [ "$DOCKER_STOP_TIMEOUT" -le "$SHUTDOWN_EXIT_MARGIN" ]; then
    echo "DOCKER_STOP_TIMEOUT must be greater than $SHUTDOWN_EXIT_MARGIN seconds" >&2
    exit 1
fi
SHUTDOWN_GRACE_PERIOD=$((DOCKER_STOP_TIMEOUT - SHUTDOWN_EXIT_MARGIN))
TEST_PID=""
PENDING_SIGNAL=""
PENDING_EXIT_CODE=""

defer_terminate() {
    if [ -z "$PENDING_SIGNAL" ]; then
        PENDING_SIGNAL=$1
        PENDING_EXIT_CODE=$2
    fi
}

terminate() {
    local signal=$1
    local exit_code=$2
    local kill_timer_pid

    # Drop the traps so a repeated signal during cleanup cannot re-enter the handler.
    trap - TERM INT
    # Gate on the owned process group, not the leader alone: the leader may
    # already be reaped while same-group cluster processes are still running.
    if [ -n "$TEST_PID" ] && kill -0 -- "-$TEST_PID" 2>/dev/null; then
        # Forward the signal to the test's whole process group (leading '-') so the
        # test binary and the cluster processes it spawned shut down gracefully.
        kill -s "$signal" -- "-$TEST_PID" 2>/dev/null || true
        # Arm a dead-man's switch: SIGKILL the process group if it is still alive
        # after the grace period.
        (
            sleep "$SHUTDOWN_GRACE_PERIOD"
            kill -KILL -- "-$TEST_PID" 2>/dev/null || true
        ) &
        kill_timer_pid=$!

        # Wait for the main test binary to exit.
        wait "$TEST_PID" 2>/dev/null || true
        # Cluster children may outlive the test binary; poll until the whole
        # process group is gone or the kill timer above has fired.
        while kill -0 -- "-$TEST_PID" 2>/dev/null; do
            if ! kill -0 "$kill_timer_pid" 2>/dev/null; then
                break
            fi
            sleep 0.1
        done
        # Disarm the kill timer so it cannot fire a stray SIGKILL later.
        kill "$kill_timer_pid" 2>/dev/null || true
        wait "$kill_timer_pid" 2>/dev/null || true
        # Last resort if anything in the group survived the timer's SIGKILL.
        if kill -0 -- "-$TEST_PID" 2>/dev/null; then
            kill -KILL -- "-$TEST_PID" 2>/dev/null || true
        fi
    fi
    # Exit with the conventional 128+signal code (143=TERM, 130=INT).
    exit "$exit_code"
}

trap 'terminate TERM 143' TERM
trap 'terminate INT 130' INT

DOCKER_ID=$1
TESTNAME=$2
shift 2

export RUST_BACKTRACE=1
export LOG_LEVEL=info

# Components pattern for env_logger. E.g. export RUST_LOG="info,raft=debug"
# Also see "--env-logger" of make-bin.sh
RUST_LOG="info"

# Optional keep temporary data on error for debugging.
# Note: When the container stops, the files from the last loop will not be automatically removed. However, they will be cleared during the next iterations.
KEEP_TMP_ON_ERROR=0
LOG_PATH="/random"
MEMORY_PROFILE=0

TEST_DURATION=120 # seconds

FOR_SERVERLESS=1
USE_TIFLASH=1
GLOBAL_TXN_FILE=1
TXN_FILE_PESSIMISTIC=0
USE_REMOTE_COP=1

TPC_WORKLOAD=1
TPCC_TXNS_THRESHOLD=100

JEPSEN_WORKLOAD=1
JEPSEN_TXN_FILE=1
JEPSEN_TXNS_THRESHOLD=100

LOAD_DATA_TASK_TIMEOUT_SEC=90

UNIQUE_WORKLOAD=0
COLUMNAR_WORKLOAD=0

RESTART_TSO_SVC=1
IA_TABLE_RATIO=0.5
ASYNC_COMMIT_RATIO=0.1
TXN_CHECK_BACKUP_TS=1
OSS_CHAOS_RATIO=0.2

UPGRADE_TEST_DURATION="60s"

ENABLE_VALUE_CACHE=1
ENABLE_TIFLASH_WRITE_NODE=0
ENABLE_COLUMNAR_NORMAL_WORKLOAD=1
ENABLE_COLUMNAR_DYNAMIC_WORKLOAD=1
ENABLE_COLUMNAR_PARTITION_WORKLOAD=1
ENABLE_COLUMNAR_VECTOR_WORKLOAD=1
ENABLE_COLUMNAR_FULLTEXT_WORKLOAD=1
ENABLE_COLUMNAR_MULTI_COLUMN_FULLTEXT_WORKLOAD_SETTING="${ENABLE_COLUMNAR_MULTI_COLUMN_FULLTEXT_WORKLOAD:--1}"
ENABLE_BLOB=0
ENABLE_CMEK=1

TIDB_NEXT_GEN=0
REP_LOG_LEVEL="info"
REP_LAZY_LOAD_SHARDS_RATIO=0.5
REP_CDC_NEWARCH=0
CDC_KEYSPACE_COUNT=2

while [ $# -gt 0 ]; do
    case "$1" in
    --keep-tmp-on-error)
        KEEP_TMP_ON_ERROR=1
        ;;
    --log-path)
        LOG_PATH="$2"
        shift
        ;;
    --memory-profile)
        MEMORY_PROFILE=1
        ;;
    --test-duration)
        TEST_DURATION="$2"
        shift
        ;;
    --no-tiflash)
        USE_TIFLASH=0
        ;;
    --no-txn-file)
        GLOBAL_TXN_FILE=0
        ;;
    --txn-file-pessimistic)
        TXN_FILE_PESSIMISTIC="$2"
        shift
        ;;
    --no-remote-cop)
        USE_REMOTE_COP=0
        ;;
    --no-tpc)
        TPC_WORKLOAD=0
        ;;
    --tpcc-txns-threshold)
        TPCC_TXNS_THRESHOLD="$2"
        shift
        ;;
    --no-jepsen)
        JEPSEN_WORKLOAD=0
        ;;
    --jepsen-no-txn-file)
        JEPSEN_TXN_FILE=0
        ;;
    --jepsen-txns-threshold)
        JEPSEN_TXNS_THRESHOLD="$2"
        shift
        ;;
    --load-data-task-timeout-sec)
        LOAD_DATA_TASK_TIMEOUT_SEC="$2"
        shift
        ;;
    --unique-workload)
        UNIQUE_WORKLOAD=1
        ;;
    --columnar-workload)
        COLUMNAR_WORKLOAD=1
        ;;
    --no-restart-tso-svc)
        RESTART_TSO_SVC=0
        ;;
    --ia-table-ratio)
        IA_TABLE_RATIO="$2"
        shift
        ;;
    --async-commit-ratio)
        ASYNC_COMMIT_RATIO="$2"
        shift
        ;;
    --no-check-backup-ts)
        TXN_CHECK_BACKUP_TS=0
        ;;
    --oss-chaos-ratio)
        OSS_CHAOS_RATIO="$2"
        shift
        ;;
    --upgrade-test-duration)
        UPGRADE_TEST_DURATION="$2"
        shift
        ;;
    --disable-value-cache)
        ENABLE_VALUE_CACHE=0
        ;;
    --enable-tiflash-write-node)
        ENABLE_TIFLASH_WRITE_NODE=1
        ;;
    --disable-columnar-normal-workload)
        ENABLE_COLUMNAR_NORMAL_WORKLOAD=0
        ;;
    --disable-columnar-dynamic-workload)
        ENABLE_COLUMNAR_DYNAMIC_WORKLOAD=0
        ;;
    --disable-columnar-partition-workload)
        ENABLE_COLUMNAR_PARTITION_WORKLOAD=0
        ;;
    --disable-columnar-vector-workload)
        ENABLE_COLUMNAR_VECTOR_WORKLOAD=0
        ;;
    --disable-columnar-fulltext-workload)
        ENABLE_COLUMNAR_FULLTEXT_WORKLOAD=0
        ;;
    --enable-columnar-multi-column-fulltext-workload)
        ENABLE_COLUMNAR_MULTI_COLUMN_FULLTEXT_WORKLOAD_SETTING=1
        ;;
    --disable-columnar-multi-column-fulltext-workload)
        ENABLE_COLUMNAR_MULTI_COLUMN_FULLTEXT_WORKLOAD_SETTING=0
        ;;
    --enable-blob)
        ENABLE_BLOB=1
        ;;
    --disable-cmek)
        ENABLE_CMEK=0
        ;;
    --tidb-next-gen)
        TIDB_NEXT_GEN=1
        ;;
    --rep-debug-log)
        RUST_LOG="info,merged_engine=debug,replication_worker=debug"
        REP_LOG_LEVEL="debug"
        ;;
    --rep-lazy-load-shards-ratio)
        REP_LAZY_LOAD_SHARDS_RATIO="$2"
        shift
        ;;
    --rep-cdc-newarch)
        REP_CDC_NEWARCH=1
        ;;
    --cdc-keyspace-count)
        CDC_KEYSPACE_COUNT="$2"
        shift
        ;;
    --serverless)
        FOR_SERVERLESS="$2"
        shift
        ;;
    *)
        echo "Usage: $0 DOCKER_ID TESTNAME [--keep-tmp-on-error] [--log-path LOG_PATH] [--memory-profile] [--serverless <0/1>]"
        exit 1
        ;;
    esac
    shift
done

export RUST_LOG
export MEMORY_PROFILE

export USE_TIFLASH
export GLOBAL_TXN_FILE
export TXN_FILE_PESSIMISTIC
export USE_REMOTE_COP

export TPC_WORKLOAD
export TPCC_TXNS_THRESHOLD

export JEPSEN_WORKLOAD
export JEPSEN_TXN_FILE
export JEPSEN_TXNS_THRESHOLD

export LOAD_DATA_TASK_TIMEOUT_SEC

export UNIQUE_WORKLOAD
export COLUMNAR_WORKLOAD

export RESTART_TSO_SVC
export IA_TABLE_RATIO
export ASYNC_COMMIT_RATIO
export TXN_CHECK_BACKUP_TS
export OSS_CHAOS_RATIO

export TEST_DURATION
export TEST_DUR_BEFORE_UPGRADE="$UPGRADE_TEST_DURATION"
export TEST_DUR_AFTER_UPGRADE="$UPGRADE_TEST_DURATION"
export TEST_DUR_AFTER_DOWNGRADE="$UPGRADE_TEST_DURATION"

export FOR_SERVERLESS
export ENABLE_VALUE_CACHE
export ENABLE_WAL_DOUBLE_WRITE
export ENABLE_TIFLASH_WRITE_NODE
export ENABLE_COLUMNAR_NORMAL_WORKLOAD
export ENABLE_COLUMNAR_DYNAMIC_WORKLOAD
export ENABLE_COLUMNAR_PARTITION_WORKLOAD
export ENABLE_COLUMNAR_VECTOR_WORKLOAD
export ENABLE_COLUMNAR_FULLTEXT_WORKLOAD
export ENABLE_BLOB
export ENABLE_CMEK
export TIDB_NEXT_GEN
export REP_LOG_LEVEL
export REP_LAZY_LOAD_SHARDS_RATIO
export REP_CDC_NEWARCH
export CDC_KEYSPACE_COUNT

mkdir -p "$LOG_PATH"/logs "$LOG_PATH"/error-logs
for i in $(seq -w 1 100000); do
    export TMPDIR="/random-tmp/$i"
    if [ "$ENABLE_COLUMNAR_MULTI_COLUMN_FULLTEXT_WORKLOAD_SETTING" -lt 0 ]; then
        export ENABLE_COLUMNAR_MULTI_COLUMN_FULLTEXT_WORKLOAD=$((10#$i % 2))
    else
        export ENABLE_COLUMNAR_MULTI_COLUMN_FULLTEXT_WORKLOAD="$ENABLE_COLUMNAR_MULTI_COLUMN_FULLTEXT_WORKLOAD_SETTING"
    fi
    mkdir -p "$TMPDIR"

    if [ "$MEMORY_PROFILE" -eq 1 ]; then
        export MALLOC_CONF="prof_leak:true,prof:true,lg_prof_interval:30,prof_final:true,prof_prefix:$TMPDIR/jeprof.out"
    fi

    LOG="$LOG_PATH"/logs/random_"$TESTNAME"_"$i"_"$DOCKER_ID".log
    CLUSTER_LOGS="$LOG_PATH"/logs/random_"$TESTNAME"_"$i"_"$DOCKER_ID"_tc
    mkdir -p "$CLUSTER_LOGS"
    # Random test get logs file path from env variable "LOG_FILE"
    export LOG_FILE="$CLUSTER_LOGS"/tikv.log
    export TEST_ID="$i"
    # Job control gives the workload an owned process group so PID 1 can signal
    # and reap it without broadcasting to unrelated container processes.
    PENDING_SIGNAL=""
    PENDING_EXIT_CODE=""
    trap 'defer_terminate TERM 143' TERM
    trap 'defer_terminate INT 130' INT
    set -m
    /random/random-bin test_random_"$TESTNAME" --nocapture >"$LOG" 2>&1 &
    TEST_PID=$!
    trap 'terminate TERM 143' TERM
    trap 'terminate INT 130' INT
    set +m
    if [ -n "$PENDING_SIGNAL" ]; then
        terminate "$PENDING_SIGNAL" "$PENDING_EXIT_CODE"
    fi
    wait "$TEST_PID" || true
    # The test leader can exit while cluster processes in its group are still
    # running. Kill and drain the group before starting the next iteration.
    kill -KILL -- "-$TEST_PID" 2>/dev/null || true
    for ((drain_attempt = 0; drain_attempt < SHUTDOWN_GRACE_PERIOD * 10; drain_attempt++)); do
        if ! kill -0 -- "-$TEST_PID" 2>/dev/null; then
            break
        fi
        sleep 0.1
    done
    if kill -0 -- "-$TEST_PID" 2>/dev/null; then
        echo "process group $TEST_PID survived SIGKILL after ${SHUTDOWN_GRACE_PERIOD}s" >&2
        exit 1
    fi
    TEST_PID=""

    sync -d "$LOG"
    if grep -q 'TEST SUCCEED' "$LOG"; then
        grep 'TEST SUCCEED' "$LOG"
        rm "$LOG"
        rm -rf "$CLUSTER_LOGS" || true
        rm -rf "$TMPDIR" || true
    else
        sync -d "$LOG_FILE"
        grep -E 'CRIT|FATAL|panicked' "$LOG_FILE" >>"$LOG" || true

        # For logs of TiDB cluster components
        if compgen -G "$TMPDIR/tc*" >/dev/null; then
            chmod +r "$TMPDIR"/tc*/*.log || true
            cp "$TMPDIR"/tc*/*.log "$TMPDIR"/tc*/*.toml "$CLUSTER_LOGS" || true
        fi
        for prefix in rep cdc; do
            if compgen -G "$TMPDIR/${prefix}*" >/dev/null; then
                chmod +r "$TMPDIR"/"$prefix"*/*.log || true
                cp "$TMPDIR"/"$prefix"*/*.log "$TMPDIR"/"$prefix"*/*.toml "$CLUSTER_LOGS" || true
                chmod +r "$TMPDIR"/"$prefix"*/sync_diff*/output/*.log || true
                cp -r "$TMPDIR"/"$prefix"*/sync_diff* "$CLUSTER_LOGS" || true
            fi
        done

        mv "$LOG" "$LOG_PATH"/error-logs/
        if [ "$KEEP_TMP_ON_ERROR" -ne 1 ]; then
            rm -rf "$TMPDIR" || true
        fi
    fi
done