source env_variables

# we do clean staff at beginning of each run, so we can keep logs of the latset run
function cleanup_data() {
    rm -rf $WORK_DIR
    mkdir $WORK_DIR
    for target_db in "$@"; do
        run_sql "drop database if exists \`${target_db}\`" $TIDB_PORT $TIDB_PASSWORD
    done
    run_sql "drop database if exists dm_meta" $TIDB_PORT $TIDB_PASSWORD
}

function cleanup_data_upstream() {
    for target_db in "$@"; do
        run_sql "drop database if exists \`${target_db}\`" $MYSQL_PORT1 $MYSQL_PASSWORD1
        run_sql "drop database if exists \`${target_db}\`" $MYSQL_PORT2 $MYSQL_PASSWORD2
    done
}

function cleanup_process() {
    dm_master_num=$(ps aux >temp && grep "dm-master.test" temp | wc -l && rm temp)
    echo "$dm_master_num dm-master alive"
    pkill -hup dm-master.test 2>/dev/null || true

    dm_worker_num=$(ps aux >temp && grep "dm-worker.test" temp | wc -l && rm temp)
    echo "$dm_worker_num dm-worker alive"
    pkill -hup dm-worker.test 2>/dev/null || true

    dm_syncer_num=$(ps aux >temp && grep "dm-syncer.test" temp | wc -l && rm temp)
    echo "$dm_syncer_num dm-syncer alive"
    pkill -hup dm-syncer.test 2>/dev/null || true

    wait_process_exit dm-master.test
    wait_process_exit dm-worker.test
    wait_process_exit dm-syncer.test
}

function cleanup_tidb_server(){
    tidb_server_num=$(ps aux >temp && grep "tidb-server" temp | wc -l && rm temp)
    echo "tidb_server_num tidb-server alive"
    pkill -hup tidb-server 2>/dev/null || true

    wait_process_exit tidb-server
}

function kill_process() {
    keyword=$1
    ps aux | grep $keyword | grep -v 'grep' | awk '{print $2}' | xargs kill || true
    wait_process_exit $keyword
}

function kill_dm_master() {
    kill_process dm-master.test
}

function kill_dm_worker() {
    kill_process dm-worker.test
}

function wait_pattern_exit() {
    pattern=$1
    while true
    do
        if ! pgrep -f $pattern >/dev/null 2>&1; then
            echo "pattern $pattern already exit"
            return 0
        fi
        sleep 0.2
        echo "wait pattern $pattern exit..."
    done
}

if [ "$RESET_MASTER" = true ]; then
    run_sql "RESET MASTER" $MYSQL_PORT1 $MYSQL_PASSWORD1
    run_sql "RESET MASTER" $MYSQL_PORT2 $MYSQL_PASSWORD2
fi

function join_string() {
    local IFS="$1"; shift; echo "$*";
}

# shortcut for start task on one DM-worker
function dmctl_start_task_standalone() {
    if [ $# -ge 2 ]; then
        remove_meta=$2
    else
        remove_meta=""
    fi
    if [ $# -ge 1 ]; then
        task_conf=$1
    else
        task_conf="$cur/conf/dm-task.yaml"
    fi
    run_dm_ctl $WORK_DIR "127.0.0.1:$MASTER_PORT" \
        "start-task $task_conf $remove_meta" \
        "\"result\": true" 2 \
        "\"source\": \"$SOURCE_ID1\"" 1
}

# shortcut for start task on two DM-workers
function dmctl_start_task() {
    if [ $# -ge 2 ]; then
        remove_meta=$2
    else
        remove_meta=""
    fi
    if [ $# -ge 1 ]; then
        task_conf=$1
    else
        task_conf="$cur/conf/dm-task.yaml"
    fi
    run_dm_ctl $WORK_DIR "127.0.0.1:$MASTER_PORT" \
        "start-task $task_conf $remove_meta" \
        "\"result\": true" 3 \
        "\"source\": \"$SOURCE_ID1\"" 1 \
        "\"source\": \"$SOURCE_ID2\"" 1
}

# shortcut for stop task on two DM-workers
function dmctl_stop_task() {
    task_name=$1
    dmctl_operate_task $task_name stop-task
}

# shortcut for retryable stop task
function dmctl_stop_task_with_retry() {
    task_name=$1
    master_port=$2
    for ((k=0; k<10; k++)); do
        run_dm_ctl $WORK_DIR "127.0.0.1:$master_port" \
            "stop-task $task_name"
        run_dm_ctl $WORK_DIR "127.0.0.1:$master_port" \
            "query-status $task_name" \
            "task test has no source or not exist" 1 && return 0
        sleep 1
    done
    return 1
}

# shortcut for pause task on two DM-workers
function dmctl_pause_task() {
    task_name=$1
    dmctl_operate_task $task_name pause-task
}

# shortcut for stop task on two DM-workers
function dmctl_resume_task() {
    task_name=$1
    dmctl_operate_task $task_name resume-task
}

function dmctl_operate_task() {
    task_name=$1
    operate=$2
    run_dm_ctl $WORK_DIR "127.0.0.1:$MASTER_PORT" \
        "$operate $task_name" \
        "\"result\": true" 3 \
        "\"source\": \"$SOURCE_ID1\"" 1 \
        "\"source\": \"$SOURCE_ID2\"" 1
}

function dmctl_operate_source() {
    operate=$1
    source_toml=$2
    source_id=$3
    run_dm_ctl $WORK_DIR "127.0.0.1:$MASTER_PORT" \
        "operate-source $operate $source_toml" \
        "\"result\": true" 2 \
        "\"source\": \"$source_id\"" 1
}

# use this func to run case
# param 1: case name
# param 2: task config name
# param 3: init_table command
# param 4: clean_table command
# param 5: shard mode
function run_case() {
    case=$1
    task_conf=$2
    init_table_cmd=$3
    clean_table_cmd=$4
    shard_mode=$5

    echo "[$(date)] <<<<<< start DM-${case} ${shard_mode} >>>>>>"

    eval ${init_table_cmd}

    truncate -s 0 $WORK_DIR/master/log/dm-master.log
    truncate -s 0 $WORK_DIR/worker1/log/dm-worker.log
    truncate -s 0 $WORK_DIR/worker2/log/dm-worker.log

    run_dm_ctl $WORK_DIR "127.0.0.1:$MASTER_PORT" \
        "start-task $cur/conf/${task_conf}.yaml --remove-meta"

    # make sure task switch to sync unit
    # otherwise we may fail because dump schema inconsistent in shardddl
    # or fail because upstream schema change but down schema unchange in redump when restart worker
    if [[ "$task_conf" == *"single"* ]]; then
        run_dm_ctl_with_retry $WORK_DIR "127.0.0.1:$MASTER_PORT" \
            "query-status test" \
            "\"unit\": \"Sync\"" 1
    elif [[ "$task_conf" == *"double"* ]]; then
        run_dm_ctl_with_retry $WORK_DIR "127.0.0.1:$MASTER_PORT" \
            "query-status test" \
            "\"unit\": \"Sync\"" 2
    fi

    args=""
    for((i=5;i<=$#;i++)); do
        j=${!i}
        args="${args} $j "
    done
    DM_${case}_CASE $args

    if [[ "$task_conf" == *"single"* ]]; then
        run_dm_ctl $WORK_DIR "127.0.0.1:$MASTER_PORT" \
            "stop-task test" \
            "\"result\": true" 2
    elif [[ "$task_conf" == *"double"* ]]; then
        run_dm_ctl $WORK_DIR "127.0.0.1:$MASTER_PORT" \
            "stop-task test" \
            "\"result\": true" 3
    fi

    eval ${clean_table_cmd}

    echo "[$(date)] <<<<<< finish DM-${case} ${shard_mode} >>>>>>"
}

# shortcut for run sql in mysql1
function run_sql_source1() {
    run_sql "$1" $MYSQL_PORT1 $MYSQL_PASSWORD1
}

# shortcut for run sql in mysql2
function run_sql_source2() {
    run_sql "$1" $MYSQL_PORT2 $MYSQL_PASSWORD2
}

# shortcut for run sql in both mysql1 and mysql2
function run_sql_both_source() {
    run_sql_source1 "$1"
    run_sql_source2 "$1"
}

# shortcut for run tidb sql
function run_sql_tidb() {
    run_sql "$1" $TIDB_PORT $TIDB_PASSWORD
}

# shortcut for run tidb sql and check result with retry
function run_sql_tidb_with_retry() {
    rc=0
    for ((k=1; k<11; k++)); do
        # in retry scenario sometimes run_sql_tidb will fail because "table not exist", we should keep retrying so turn
        # off the error option temporarily.
        set +e
        run_sql_tidb "$1"
        set -e
        if grep -Fq "$2" "$TEST_DIR/sql_res.$TEST_NAME.txt"; then
            rc=1
            break
        fi
        echo "run tidb sql failed $k-th time, retry later"
        sleep 2
    done
    if [[ $rc = 0 ]]; then
        echo "TEST FAILED: OUTPUT DOES NOT CONTAIN '$2'"
        echo "____________________________________"
        cat "$TEST_DIR/sql_res.$TEST_NAME.txt"
        echo "^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^"
        exit 1
    fi
}

# shortcut for run tidb sql and check result with retry
function run_sql_tidb_with_retry_times() {
    rc=0
    for ((k=1; k<$3; k++)); do
        run_sql_tidb "$1"
        if grep -Fq "$2" "$TEST_DIR/sql_res.$TEST_NAME.txt"; then
            rc=1
            break
        fi
        echo "run tidb sql failed $k-th time, retry later"
        sleep 2
    done
    if [[ $rc = 0 ]]; then
        echo "TEST FAILED: OUTPUT DOES NOT CONTAIN '$2'"
        echo "____________________________________"
        cat "$TEST_DIR/sql_res.$TEST_NAME.txt"
        echo "^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^"
        exit 1
    fi
}

# shortcut for run tidb sql and check result with retry
# $1: table
# $2: fields numbers
# $3: value starts range(closed interval)
# $4: value ends range(closed interval)
# $5: database host
# $6: database port
# $7: database password
function run_sql_with_txn() {
    template="(?"
    for _ in {2..$2..1}
    do
        template+=",?"
    done
    template+=")"
    echo "begin;" | tee $WORK_DIR/tmp.sql
    seq $3 $4 | xargs -I\? echo "insert into $1 values $template;" | tee -a $WORK_DIR/tmp.sql
    echo "commit;" | tee -a $WORK_DIR/tmp.sql
    run_sql_file $WORK_DIR/tmp.sql $5 $6 $7
}

# shortcut for check log contain with retry
function check_log_contain_with_retry() {
    text=$1
    log1=$2
    log2=""
    if [[ "$#" -ge 3 ]]; then
        log2=$3
    fi
    rc=0
    for ((k=1;k<61;k++)); do
        if [[ ! -f $log1 ]]; then
            sleep 2
            echo "check log contain failed $k-th time (file not exist), retry later"
            continue
        fi
        got=`grep -a "$text" $log1 | wc -l`
        if [[ $got -ne 0 ]]; then
            rc=1
            break
        fi
        if [[ ! "$log2" = "" ]]; then
            if [[ ! -f $log2 ]]; then
                sleep 2
                echo "check log contain failed $k-th time (file not exist), retry later"
                continue
            fi
            got=`grep -a "$text" $log2 | wc -l`
            if [[ $got -ne 0 ]]; then
                rc=1
                break
            fi
        fi
        echo "check log contain failed $k-th time, retry later"
        sleep 2
    done
    if [[ $rc -eq 0 ]]; then
        echo "log doesn't contain $text"
        exit 1
    fi
}

# shortcut for init cluster with one master and two workers
function init_cluster(){
    run_dm_master $WORK_DIR/master $MASTER_PORT $cur/conf/dm-master.toml
    check_rpc_alive $cur/../bin/check_master_online 127.0.0.1:$MASTER_PORT

    cp $cur/conf/source1.yaml $WORK_DIR/source1.yaml
    cp $cur/conf/source2.yaml $WORK_DIR/source2.yaml
    sed -i "/relay-binlog-name/i\relay-dir: $WORK_DIR/worker1/relay_log" $WORK_DIR/source1.yaml
    sed -i "/relay-binlog-name/i\relay-dir: $WORK_DIR/worker2/relay_log" $WORK_DIR/source2.yaml

    run_dm_worker $WORK_DIR/worker1 $WORKER1_PORT $cur/conf/dm-worker1.toml
    check_rpc_alive $cur/../bin/check_worker_online 127.0.0.1:$WORKER1_PORT
    dmctl_operate_source create $WORK_DIR/source1.yaml $SOURCE_ID1

    run_dm_worker $WORK_DIR/worker2 $WORKER2_PORT $cur/conf/dm-worker2.toml
    check_rpc_alive $cur/../bin/check_worker_online 127.0.0.1:$WORKER2_PORT
    dmctl_operate_source create $WORK_DIR/source2.yaml $SOURCE_ID2
}

function get_master_status() {
	arr=$(echo "show master status;" | MYSQL_PWD=123456 mysql -uroot -h$1 -P$2 | awk 'NR==2')
	echo $arr
}

function check_master_port_offline() {
    master_ports_1_indexed=(0 $MASTER_PORT1 $MASTER_PORT2 $MASTER_PORT3 $MASTER_PORT4 $MASTER_PORT5)
    master_peer_ports_1_indexed=(0 $MASTER_PEER_PORT1 $MASTER_PEER_PORT2 $MASTER_PEER_PORT3 $MASTER_PEER_PORT4 $MASTER_PEER_PORT5)
    idx=$1
    check_port_offline ${master_ports_1_indexed[$idx]} 20
    check_port_offline ${master_peer_ports_1_indexed[$idx]} 20
}

function check_rows_equal() {
    num=$1
    check_contains "$num. row"
    ((num++))
    check_not_contains "$num. row"
}