#!/usr/bin/env bash

set -Eeuo pipefail

# Interactive cloud E2E usage:
#   bash e2e-test/run_snowflake_e2e.sh
#
# Non-interactive usage can provide every prompt through environment variables.
# The script starts a local TiDB/TiCDC playground, writes snapshot and
# incremental CSV files directly into a unique S3 run prefix, deploys the real
# Fivetran Connector SDK connection, triggers cloud syncs, then compares
# current TiDB rows with active Snowflake rows after e2e-test/test.sql has run.

fail() {
    printf 'ERROR: %s\n' "$*" >&2
    exit 1
}

log() {
    printf '%s\n' "$*"
}

run() {
    printf '+'
    printf ' %q' "$@"
    printf '\n'
    "$@" || fail "command failed ($?): $*"
}

run_redacted() {
    local label=$1
    shift

    log "+ $label"
    "$@" || fail "command failed ($?): $label"
}

require_command() {
    command -v "$1" >/dev/null 2>&1 || fail "missing required command: $1"
}

script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)" || exit 1
repo_root="$(cd "${script_dir}/.." && pwd -P)" || exit 1
cd "$repo_root" || exit 1

: "${PYTHON_VENV:=.venv}"
: "${SKIP_VENV:=0}"
: "${TIDB_VERSION:=v8.5.6}"
: "${TIUP_TAG:=tidb-fivetran-snowflake}"
: "${DB_NAME:=test}"
: "${TABLE_NAME:=fruit1}"
: "${E2E_RUNTIME_DIR:=e2e-test/.tmp/snowflake-e2e}"
: "${TEST_SQL_PATH:=e2e-test/test.sql}"
: "${WORKSPACE:=examples/snowflake_workspace}"
: "${TIDB_HOST:=127.0.0.1}"
: "${TIDB_PORT:=4000}"
: "${TIDB_USER:=root}"
: "${TIDB_PASSWORD:=}"
: "${TICDC_SERVER:=127.0.0.1:8300}"
: "${CDC_CLI:=./bin/cdc}"
: "${DUMPLING_THREADS:=8}"
: "${DUMPLING_ROWS:=200000}"
: "${DUMPLING_FILE_SIZE:=64MiB}"
: "${READY_TIMEOUT_SECONDS:=180}"
: "${FIVETRAN_DEPLOY_TIMEOUT_SECONDS:=600}"
: "${FIVETRAN_SYNC_TIMEOUT_SECONDS:=1200}"
: "${SNOWFLAKE_READY_TIMEOUT_SECONDS:=1200}"
: "${PLAYGROUND_LOG_TAIL_LINES:=200}"
: "${KEEP_PLAYGROUND:=0}"
: "${CLEANUP_S3_ON_EXIT:=0}"
: "${SHOW_RUNTIME_CONFIG:=1}"
: "${TIMEOUT_BIN:=timeout}"
: "${FIVETRAN_PYTHON_VERSION:=3.13}"
: "${RUN_ID:=tidb_fivetran_snowflake_e2e_$(date -u +%Y%m%d%H%M%S)}"
: "${CHANGEFEED_ID:=$RUN_ID}"
: "${SYNC_TASK_ID:=$RUN_ID}"
: "${FIVETRAN_API_KEY_BASE64:=}"
: "${FIVETRAN_DESTINATION_NAME:=}"
: "${FIVETRAN_CONNECTION_NAME:=$RUN_ID}"
: "${FIVETRAN_BASE_URL:=https://api.fivetran.com/v1}"
: "${RESET_FIVETRAN_STATE:=0}"
: "${S3_BUCKET:=}"
: "${S3_BASE_PREFIX:=tidb-fivetran-e2e}"
: "${S3_RUN_PREFIX:=}"
: "${S3_REGION:=}"
: "${S3_AK:=}"
: "${S3_SK:=}"
: "${S3_SESSION_TOKEN:=}"
: "${SNOWFLAKE_ACCOUNT:=}"
: "${SNOWFLAKE_USER:=}"
: "${SNOWFLAKE_PASSWORD:=}"
: "${SNOWFLAKE_WAREHOUSE:=}"
: "${SNOWFLAKE_DATABASE:=}"
: "${SNOWFLAKE_ROLE:=}"
: "${SNOWFLAKE_SCHEMA:=}"

tag="$TIUP_TAG"
db_name="$DB_NAME"
table_name="$TABLE_NAME"
changefeed_id="$CHANGEFEED_ID"
workspace="$WORKSPACE"
snapshot_dir="$workspace/snapshot"
increment_dir="$workspace/increment"
runtime_dir="$E2E_RUNTIME_DIR"
comparison_dir="$runtime_dir/comparison"
test_sql_path="$TEST_SQL_PATH"
playground_log="$runtime_dir/playground.log"
ticdc_config="$runtime_dir/ticdc.toml"
changefeed_config="$runtime_dir/changefeed.toml"
connector_config_path="$runtime_dir/fivetran_configuration.json"
snowflake_helper="$runtime_dir/snowflake_compare.py"
tiup_data_dir="$HOME/.tiup/data/$tag"
playground_pid=
cdc_cli=("$CDC_CLI" cli)
snowflake_python_cmd=()
connection_id=
ticdc_host="${TICDC_SERVER%:*}"
ticdc_port="${TICDC_SERVER##*:}"

prompt_required() {
    local var_name=$1
    local label=${2:-$1}
    local value=${!var_name:-}

    if [ -n "$value" ]; then
        return 0
    fi
    if [ ! -t 0 ]; then
        fail "$var_name is required; set it in the environment for non-interactive runs"
    fi

    printf '%s: ' "$label" >&2
    IFS= read -r value
    [ -n "$value" ] || fail "$var_name is required"
    printf -v "$var_name" '%s' "$value"
    export "$var_name"
}

prompt_optional() {
    local var_name=$1
    local label=${2:-$1}
    local value=${!var_name:-}

    if [ -n "$value" ] || [ ! -t 0 ]; then
        return 0
    fi

    printf '%s (optional): ' "$label" >&2
    IFS= read -r value
    printf -v "$var_name" '%s' "$value"
    export "$var_name"
}

prompt_required_fresh() {
    local var_name=$1
    local label=${2:-$1}
    local value=${!var_name:-}

    if [ -t 0 ]; then
        printf '%s: ' "$label" >&2
        IFS= read -r value
        [ -n "$value" ] || fail "$var_name is required"
    elif [ -z "$value" ]; then
        fail "$var_name is required; set it in the environment for non-interactive runs"
    fi

    printf -v "$var_name" '%s' "$value"
    export "$var_name"
}

prompt_optional_fresh() {
    local var_name=$1
    local label=${2:-$1}
    local value=${!var_name:-}

    if [ -t 0 ]; then
        printf '%s (optional): ' "$label" >&2
        IFS= read -r value
    fi

    printf -v "$var_name" '%s' "$value"
    export "$var_name"
}

prompt_secret_fresh() {
    local var_name=$1
    local label=${2:-$1}
    local value=${!var_name:-}

    if [ -t 0 ]; then
        printf '%s: ' "$label" >&2
        IFS= read -r -s value
        printf '\n' >&2
        [ -n "$value" ] || fail "$var_name is required"
    elif [ -z "$value" ]; then
        fail "$var_name is required; set it in the environment for non-interactive runs"
    fi

    printf -v "$var_name" '%s' "$value"
    export "$var_name"
}

prompt_secret() {
    local var_name=$1
    local label=${2:-$1}
    local value=${!var_name:-}

    if [ -n "$value" ]; then
        return 0
    fi
    if [ ! -t 0 ]; then
        fail "$var_name is required; set it in the environment for non-interactive runs"
    fi

    printf '%s: ' "$label" >&2
    IFS= read -r -s value
    printf '\n' >&2
    [ -n "$value" ] || fail "$var_name is required"
    printf -v "$var_name" '%s' "$value"
    export "$var_name"
}

normalize_s3_run_prefix() {
    if [ -z "$S3_RUN_PREFIX" ]; then
        local base_prefix
        base_prefix="${S3_BASE_PREFIX#/}"
        base_prefix="${base_prefix%/}"
        if [ -n "$base_prefix" ]; then
            S3_RUN_PREFIX="$base_prefix/$RUN_ID"
        else
            S3_RUN_PREFIX="$RUN_ID"
        fi
        export S3_RUN_PREFIX
    fi

    S3_RUN_PREFIX="${S3_RUN_PREFIX#/}"
    S3_RUN_PREFIX="${S3_RUN_PREFIX%/}"
    [ -n "$S3_RUN_PREFIX" ] || fail "S3_RUN_PREFIX must not be empty"
    export S3_RUN_PREFIX
}

collect_interactive_inputs() {
    prompt_secret "FIVETRAN_API_KEY_BASE64" "Fivetran API key, base64 encoded"
    prompt_required "FIVETRAN_DESTINATION_NAME" "Fivetran destination name"
    prompt_required "S3_BUCKET" "S3 bucket"
    prompt_required "S3_REGION" "S3 region"
    prompt_required_fresh "S3_AK" "S3 access key"
    prompt_secret_fresh "S3_SK" "S3 secret key"
    prompt_optional_fresh "S3_SESSION_TOKEN" "S3 session token"
    prompt_required "SNOWFLAKE_ACCOUNT" "Snowflake account"
    prompt_required "SNOWFLAKE_USER" "Snowflake user"
    prompt_secret "SNOWFLAKE_PASSWORD" "Snowflake password"
    prompt_required "SNOWFLAKE_WAREHOUSE" "Snowflake warehouse"
    prompt_required "SNOWFLAKE_DATABASE" "Snowflake database"
    prompt_optional "SNOWFLAKE_ROLE" "Snowflake role"

    if [ -z "$SNOWFLAKE_SCHEMA" ]; then
        SNOWFLAKE_SCHEMA="$FIVETRAN_CONNECTION_NAME"
        export SNOWFLAKE_SCHEMA
    fi

    normalize_s3_run_prefix
}

mask_secret_value() {
    local value=$1
    local length=${#value}
    local tail

    if [ "$length" -eq 0 ]; then
        printf '<empty>'
        return 0
    fi

    if [ "$length" -le 4 ]; then
        printf '<redacted len=%s>' "$length"
        return 0
    fi

    tail=${value: -4}
    printf '<redacted len=%s tail=%s>' "$length" "$tail"
}

log_config_plain() {
    local var_name=$1
    local value=${!var_name:-}

    if [ -z "$value" ]; then
        log "  $var_name=<empty>"
    else
        log "  $var_name=$value"
    fi
}

log_config_secret() {
    local var_name=$1
    local value=${!var_name:-}

    log "  $var_name=$(mask_secret_value "$value")"
}

log_runtime_configuration() {
    if [ "$SHOW_RUNTIME_CONFIG" != "1" ]; then
        return 0
    fi

    log "Runtime configuration:"
    log_config_plain "RUN_ID"
    log_config_plain "CHANGEFEED_ID"
    log_config_plain "SYNC_TASK_ID"
    log_config_plain "FIVETRAN_DESTINATION_NAME"
    log_config_plain "FIVETRAN_CONNECTION_NAME"
    log_config_plain "FIVETRAN_BASE_URL"
    log_config_secret "FIVETRAN_API_KEY_BASE64"
    log_config_plain "S3_BUCKET"
    log_config_plain "S3_REGION"
    log_config_plain "S3_BASE_PREFIX"
    log_config_plain "S3_RUN_PREFIX"
    log_config_secret "S3_AK"
    log_config_secret "S3_SK"
    log_config_secret "S3_SESSION_TOKEN"
    log_config_plain "SNOWFLAKE_ACCOUNT"
    log_config_plain "SNOWFLAKE_USER"
    log_config_secret "SNOWFLAKE_PASSWORD"
    log_config_plain "SNOWFLAKE_WAREHOUSE"
    log_config_plain "SNOWFLAKE_DATABASE"
    log_config_plain "SNOWFLAKE_ROLE"
    log_config_plain "SNOWFLAKE_SCHEMA"
    log_config_plain "TIDB_HOST"
    log_config_plain "TIDB_PORT"
    log_config_plain "TIDB_USER"
    log_config_secret "TIDB_PASSWORD"
    log_config_plain "TICDC_SERVER"
    log_config_plain "WORKSPACE"
    log_config_plain "E2E_RUNTIME_DIR"
    log_config_plain "TEST_SQL_PATH"
}

kill_tagged_tiup_processes() {
    command -v pgrep >/dev/null 2>&1 || return 0

    local pids
    pids=$(pgrep -f "$tiup_data_dir/" 2>/dev/null || true)
    if [ -z "$pids" ]; then
        return 0
    fi

    # shellcheck disable=SC2086
    kill -TERM $pids >/dev/null 2>&1 || true
    sleep 2
    pids=$(pgrep -f "$tiup_data_dir/" 2>/dev/null || true)
    if [ -n "$pids" ]; then
        # shellcheck disable=SC2086
        kill -KILL $pids >/dev/null 2>&1 || true
    fi
}

cleanup() {
    if [ -n "${playground_pid:-}" ]; then
        if [ "$KEEP_PLAYGROUND" = "1" ]; then
            log "Leaving playground running with pid $playground_pid because KEEP_PLAYGROUND=1"
        else
            kill -TERM -- -"$playground_pid" >/dev/null 2>&1 || true
            kill -TERM -- "$playground_pid" >/dev/null 2>&1 || true
            sleep 2
            kill -KILL -- -"$playground_pid" >/dev/null 2>&1 || true
            kill -KILL -- "$playground_pid" >/dev/null 2>&1 || true
            wait "$playground_pid" 2>/dev/null || true
        fi
    fi

    rm -f "$ticdc_config" "$changefeed_config" "$connector_config_path" "$snowflake_helper"

    if [ "$CLEANUP_S3_ON_EXIT" = "1" ] && [ -n "${S3_BUCKET:-}" ] && [ -n "${S3_RUN_PREFIX:-}" ]; then
        aws_s3_rm_prefix "$S3_RUN_PREFIX" || true
    fi

    if [ "$KEEP_PLAYGROUND" != "1" ]; then
        tiup clean "$tag" >/dev/null 2>&1 || true
        kill_tagged_tiup_processes
    fi
}

on_error() {
    local status=$?
    local line_no=${BASH_LINENO[0]:-$LINENO}
    local command=${BASH_COMMAND:-unknown}

    trap - ERR
    printf 'ERROR: failed unexpectedly at line %s: %s\n' "$line_no" "$command" >&2
    return "$status"
}

on_exit() {
    status=$?
    if [ "$status" -ne 0 ] && [ -n "${playground_pid:-}" ] && [ -f "$playground_log" ]; then
        {
            printf '\n===== playground log tail (%s lines) =====\n' "$PLAYGROUND_LOG_TAIL_LINES"
            tail -n "$PLAYGROUND_LOG_TAIL_LINES" "$playground_log" || true
            printf '===== end playground log tail =====\n'
        } >&2
    fi
    cleanup
    exit "$status"
}

trap on_exit EXIT
trap on_error ERR

activate_python() {
    if [ "$SKIP_VENV" = "1" ]; then
        return 0
    fi

    if [ -f "${PYTHON_VENV}/bin/activate" ]; then
        source "${PYTHON_VENV}/bin/activate"
        return 0
    fi

    if command -v fivetran >/dev/null 2>&1; then
        log "No ${PYTHON_VENV}/bin/activate found; using current Python environment."
        return 0
    fi

    fail "missing ${PYTHON_VENV}/bin/activate and fivetran is not on PATH"
}

ensure_fivetran_cli() {
    if command -v fivetran >/dev/null 2>&1; then
        return 0
    fi

    if [ "$SKIP_VENV" != "1" ] && command -v uv >/dev/null 2>&1; then
        log "fivetran CLI is missing from the active environment; running uv sync to install dev dependencies."
        run uv sync
        if [ -f "${PYTHON_VENV}/bin/activate" ]; then
            source "${PYTHON_VENV}/bin/activate"
        fi
    fi

    if command -v fivetran >/dev/null 2>&1; then
        return 0
    fi

    fail "missing required command: fivetran; run 'uv sync' or install fivetran-connector-sdk in ${PYTHON_VENV}"
}

validate_identifier() {
    local value=$1
    local label=$2

    if [[ ! "$value" =~ ^[A-Za-z_][A-Za-z0-9_]*$ ]]; then
        fail "invalid $label: $value"
    fi
}

validate_port() {
    local value=$1
    local label=$2

    case "$value" in
        '' | *[!0-9]*)
            fail "invalid $label port: $value"
            ;;
    esac

    if [ "$value" -lt 1 ] || [ "$value" -gt 65535 ]; then
        fail "invalid $label port: $value"
    fi
}

validate_ticdc_server() {
    if [ -z "$ticdc_host" ] || [ "$ticdc_host" = "$ticdc_port" ]; then
        fail "invalid TICDC_SERVER: $TICDC_SERVER; expected host:port"
    fi

    validate_port "$ticdc_port" TICDC_SERVER
}

wait_until() {
    local description=$1
    local timeout_seconds=$2
    shift 2

    local deadline=$((SECONDS + timeout_seconds))
    while [ "$SECONDS" -lt "$deadline" ]; do
        if "$@"; then
            return 0
        fi
        sleep 2
    done

    fail "$description did not become ready within ${timeout_seconds}s"
}

tidb_ready() {
    mysql --protocol=tcp -h "$TIDB_HOST" -P "$TIDB_PORT" -u "$TIDB_USER" --password="$TIDB_PASSWORD" -e "select 1" >/dev/null 2>&1
}

ticdc_ready() {
    "${cdc_cli[@]}" changefeed list --server="$TICDC_SERVER" >/dev/null 2>&1
}

mysql_exec() {
    local sql=$1

    mysql --protocol=tcp -h "$TIDB_HOST" -P "$TIDB_PORT" -u "$TIDB_USER" --password="$TIDB_PASSWORD" -e "$sql" ||
        fail "mysql failed: $sql"
}

mysql_try_exec() {
    local sql=$1

    if ! mysql --protocol=tcp -h "$TIDB_HOST" -P "$TIDB_PORT" -u "$TIDB_USER" --password="$TIDB_PASSWORD" -e "$sql"; then
        log "Optional TiDB setting failed and was ignored: $sql"
    fi
}

mysql_exec_file() {
    local file=$1

    [ -f "$file" ] || fail "missing SQL file: $file"
    mysql --protocol=tcp -h "$TIDB_HOST" -P "$TIDB_PORT" -u "$TIDB_USER" --password="$TIDB_PASSWORD" <"$file" ||
        fail "mysql failed while executing $file"
}

mysql_export_tsv() {
    local sql=$1
    local output_path=$2

    mysql --protocol=tcp -h "$TIDB_HOST" -P "$TIDB_PORT" -u "$TIDB_USER" --password="$TIDB_PASSWORD" \
        --batch --raw --skip-column-names -e "$sql" >"$output_path" ||
        fail "mysql failed while exporting comparison query"
}

wait_for_increment_outputs() {
    local deadline=$((SECONDS + READY_TIMEOUT_SECONDS))
    while [ "$SECONDS" -lt "$deadline" ]; do
        if aws_s3_object_exists "$increment_s3_uri/metadata" &&
            aws_s3_prefix_has_regex "$increment_s3_uri" '/CDC[^/]*\.csv$' &&
            aws_s3_prefix_has_regex "$increment_s3_uri" '/CDC\.index$' &&
            aws_s3_prefix_has_regex "$increment_s3_uri" '/schema_[^/]*\.json$'; then
            return 0
        fi
        sleep 2
    done

    fail "TiCDC did not write expected increment files under $increment_s3_uri"
}

metadata_checkpoint_ts() {
    aws_s3_cp_stdout "$increment_s3_uri/metadata" |
        sed -n 's/.*"checkpoint-ts"[[:space:]]*:[[:space:]]*"\{0,1\}\([0-9][0-9]*\)"\{0,1\}.*/\1/p' |
        tail -n 1
}

increment_row_commit_ts() {
    local needle=$1
    local key
    local row_commit_ts

    while IFS= read -r key; do
        row_commit_ts=$(
            aws_s3_cp_stdout "s3://$S3_BUCKET/$key" 2>/dev/null |
                grep -F "$needle" |
                awk -F',' 'NR == 1 {print $4; exit}'
        )
        if [ -n "$row_commit_ts" ]; then
            printf '%s\n' "$row_commit_ts"
            return 0
        fi
    done < <(aws_s3_prefix_keys "$increment_s3_uri" | grep -E '/CDC[^/]*\.csv$')

    return 1
}

final_ddl_row_checkpoint_ready() {
    local row_commit_ts
    local checkpoint_ts

    row_commit_ts=$(increment_row_commit_ts "fruit-after-truncate" || true)
    checkpoint_ts=$(metadata_checkpoint_ts || true)
    if [[ ! "$row_commit_ts" =~ ^[0-9]+$ ]] || [[ ! "$checkpoint_ts" =~ ^[0-9]+$ ]]; then
        return 1
    fi

    [ "$checkpoint_ts" -ge "$row_commit_ts" ]
}

test_sql_final_row_checkpoint_ready() {
    local row_commit_ts
    local checkpoint_ts

    row_commit_ts=$(increment_row_commit_ts "cross-family value after ddl" || true)
    checkpoint_ts=$(metadata_checkpoint_ts || true)
    if [[ ! "$row_commit_ts" =~ ^[0-9]+$ ]] || [[ ! "$checkpoint_ts" =~ ^[0-9]+$ ]]; then
        return 1
    fi

    [ "$checkpoint_ts" -ge "$row_commit_ts" ]
}

write_runtime_configs() {
    mkdir -p "$runtime_dir" "$comparison_dir" || fail "failed to prepare $runtime_dir"

    cat >"$ticdc_config" <<'EOF'
newarch = true
EOF

    cat >"$changefeed_config" <<'EOF'
[sink]
[sink.csv]
include-commit-ts = true
EOF

    cat >"$snowflake_helper" <<'PY'
#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import os
import sys
import time

try:
    import snowflake.connector
except ModuleNotFoundError as exc:
    raise SystemExit(
        "missing snowflake-connector-python; install it or make uv available so the shell script can run "
        "`uv run --no-project --with snowflake-connector-python python`"
    ) from exc

NULL_SENTINEL = "__NULL__"

TABLE_CANDIDATES = {
    "fruit": ("SNOWFLAKE_FRUIT_TABLE", ["test_fruit_1", "test_fruit1"]),
    "matrix": ("SNOWFLAKE_MATRIX_TABLE", ["ft_ddl_type_test_ft_ddl_type_matrix"]),
    "create": ("SNOWFLAKE_CREATE_AFTER_SNAPSHOT_TABLE", ["ft_ddl_type_test_ft_ddl_create_after_snapshot"]),
    "truncate": ("SNOWFLAKE_TRUNCATE_TABLE", ["ft_ddl_type_test_ft_ddl_truncate_target"]),
}

PROFILE_COLUMNS = {
    "fruit": ["id", "name", "color", "sweetness"],
    "matrix": [
        "id",
        "c_varchar",
        "c_added_nullable",
        "c_default_toggle",
        "c_changed_rename_type",
        "c_type_cross_family_number",
    ],
    "create": ["id", "payload"],
    "truncate": ["id", "payload"],
}


def connect():
    kwargs = {
        "account": require_env("SNOWFLAKE_ACCOUNT"),
        "user": require_env("SNOWFLAKE_USER"),
        "password": require_env("SNOWFLAKE_PASSWORD"),
        "warehouse": require_env("SNOWFLAKE_WAREHOUSE"),
        "database": require_env("SNOWFLAKE_DATABASE"),
    }
    role = os.environ.get("SNOWFLAKE_ROLE")
    if role:
        kwargs["role"] = role
    return snowflake.connector.connect(**kwargs)


def require_env(name: str) -> str:
    value = os.environ.get(name)
    if not value:
        raise SystemExit(f"{name} is required")
    return value


def quote_ident(value: str) -> str:
    return '"' + value.replace('"', '""') + '"'


def profile_candidates(profile: str) -> list[str]:
    if profile not in TABLE_CANDIDATES:
        raise SystemExit(f"unknown profile: {profile}")
    env_name, defaults = TABLE_CANDIDATES[profile]
    result: list[str] = []
    override = os.environ.get(env_name)
    if override:
        result.append(override)
    for item in defaults:
        if item not in result:
            result.append(item)
    return result


def resolve_relation(cur, schema_name: str, candidates: list[str]) -> tuple[str, str]:
    cur.execute(
        """
        select table_schema, table_name
        from information_schema.tables
        where upper(table_schema) = upper(%s)
        """,
        (schema_name,),
    )
    rows = [(str(row[0]), str(row[1])) for row in cur.fetchall()]
    for candidate in candidates:
        for actual_schema, actual_table in rows:
            if actual_table.upper() == candidate.upper():
                return actual_schema, actual_table
    raise RuntimeError(
        "none of the expected Snowflake tables exists in schema "
        f"{schema_name!r}: {', '.join(candidates)}"
    )


def relation_exists(cur, schema_name: str, candidates: list[str]) -> bool:
    try:
        resolve_relation(cur, schema_name, candidates)
    except RuntimeError:
        return False
    return True


def load_columns(cur, schema_name: str, table_name: str) -> dict[str, str]:
    cur.execute(
        """
        select column_name
        from information_schema.columns
        where table_schema = %s and table_name = %s
        """,
        (schema_name, table_name),
    )
    return {str(row[0]).upper(): str(row[0]) for row in cur.fetchall()}


def column(columns: dict[str, str], name: str) -> str:
    actual = columns.get(name.upper())
    if actual is None:
        raise RuntimeError(f"Snowflake column {name!r} was not found")
    return quote_ident(actual)


def active_condition(columns: dict[str, str]) -> str:
    deleted = columns.get("_FIVETRAN_DELETED")
    if deleted is None:
        return "TRUE"
    return f"COALESCE({quote_ident(deleted)}, FALSE) = FALSE"


def to_string_expr(columns: dict[str, str], name: str) -> str:
    return f"COALESCE(TO_VARCHAR({column(columns, name)}), '{NULL_SENTINEL}')"


def export_profile(profile: str, schema_name: str, output_path: str) -> None:
    candidates = profile_candidates(profile)
    wanted_columns = PROFILE_COLUMNS[profile]
    with connect() as conn:
        cur = conn.cursor()
        actual_schema, actual_table = resolve_relation(cur, schema_name, candidates)
        columns = load_columns(cur, actual_schema, actual_table)
        select_list = ", ".join(to_string_expr(columns, name) for name in wanted_columns)
        id_column = column(columns, wanted_columns[0])
        sql = (
            f"select {select_list} "
            f"from {quote_ident(actual_schema)}.{quote_ident(actual_table)} "
            f"where {active_condition(columns)} "
            f"order by TRY_TO_NUMBER({id_column}), {id_column}"
        )
        cur.execute(sql)
        rows = cur.fetchall()

    with open(output_path, "w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle, delimiter="\t", lineterminator="\n")
        for row in rows:
            writer.writerow(["" if value is None else str(value) for value in row])


def count_active(profile: str, schema_name: str) -> int:
    candidates = profile_candidates(profile)
    with connect() as conn:
        cur = conn.cursor()
        actual_schema, actual_table = resolve_relation(cur, schema_name, candidates)
        columns = load_columns(cur, actual_schema, actual_table)
        cur.execute(
            f"select count(*) from {quote_ident(actual_schema)}.{quote_ident(actual_table)} "
            f"where {active_condition(columns)}"
        )
        return int(cur.fetchone()[0])


def check_absent(schema_name: str, candidates: list[str]) -> None:
    with connect() as conn:
        cur = conn.cursor()
        if relation_exists(cur, schema_name, candidates):
            raise SystemExit(
                "unexpected Snowflake table exists in schema "
                f"{schema_name!r}: {', '.join(candidates)}"
            )


def wait_count(profile: str, schema_name: str, expected: int, timeout_seconds: int) -> None:
    deadline = time.monotonic() + timeout_seconds
    last_error = ""
    while time.monotonic() < deadline:
        try:
            actual = count_active(profile, schema_name)
            if actual == expected:
                return
            last_error = f"active count was {actual}, expected {expected}"
        except Exception as exc:
            last_error = str(exc)
        time.sleep(10)
    raise SystemExit(f"Snowflake profile {profile} did not reach active count {expected}: {last_error}")


def main() -> int:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)

    export_parser = subparsers.add_parser("export-profile")
    export_parser.add_argument("profile")
    export_parser.add_argument("schema")
    export_parser.add_argument("output")

    count_parser = subparsers.add_parser("count-active")
    count_parser.add_argument("profile")
    count_parser.add_argument("schema")

    absent_parser = subparsers.add_parser("check-absent")
    absent_parser.add_argument("schema")
    absent_parser.add_argument("table", nargs="+")

    wait_parser = subparsers.add_parser("wait-count")
    wait_parser.add_argument("profile")
    wait_parser.add_argument("schema")
    wait_parser.add_argument("expected", type=int)
    wait_parser.add_argument("timeout_seconds", type=int)

    args = parser.parse_args()
    if args.command == "export-profile":
        export_profile(args.profile, args.schema, args.output)
    elif args.command == "count-active":
        print(count_active(args.profile, args.schema))
    elif args.command == "check-absent":
        check_absent(args.schema, args.table)
    elif args.command == "wait-count":
        wait_count(args.profile, args.schema, args.expected, args.timeout_seconds)
    else:
        raise AssertionError(args.command)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
PY
    chmod 700 "$snowflake_helper"
}

resolve_cdc_cli() {
    if [ -x "$CDC_CLI" ]; then
        cdc_cli=("$CDC_CLI" cli)
        return 0
    fi

    cdc_cli=(tiup "ctl:${TIDB_VERSION}" cdc)
}

resolve_snowflake_python() {
    if python3 -c 'import snowflake.connector' >/dev/null 2>&1; then
        snowflake_python_cmd=(python3)
        return 0
    fi

    if command -v uv >/dev/null 2>&1; then
        snowflake_python_cmd=(uv run --no-project --quiet --with snowflake-connector-python python)
        return 0
    fi

    fail "missing snowflake-connector-python; install it for python3 or install uv"
}

run_snowflake_helper() {
    "${snowflake_python_cmd[@]}" "$snowflake_helper" "$@"
}

prepare_test_sql_environment() {
    mysql_try_exec "set global tidb_enable_check_constraint = on;"
    mysql_try_exec "set global tidb_enable_foreign_key = on;"
}

apply_ddl_test_cases() {
    local ddl_sql

    ddl_sql="
use $db_name;
insert into $table_name(name) values ('fruit-11');
alter table $table_name add column color varchar(20) default null;
insert into $table_name(name, color) values ('fruit-12', 'red');
alter table $table_name add column sweetness int not null default 0;
insert into $table_name(name, color, sweetness) values ('fruit-13', 'yellow', 7);
alter table $table_name modify column name varchar(64) not null default '';
insert into $table_name(name, color, sweetness) values ('fruit-14-long-name-after-widen', 'purple', 8);
alter table $table_name comment = 'snowflake e2e ddl coverage';
create table ddl_no_pk(note varchar(20));
insert into ddl_no_pk(note) values ('ignored');
truncate table $table_name;
insert into $table_name(id, name, color, sweetness) values (100, 'fruit-after-truncate', 'green', 9);
"
    mysql_exec "$ddl_sql"
}

configure_aws_environment() {
    export AWS_ACCESS_KEY_ID="$S3_AK"
    export AWS_SECRET_ACCESS_KEY="$S3_SK"
    export AWS_DEFAULT_REGION="$S3_REGION"
    if [ -n "$S3_SESSION_TOKEN" ]; then
        export AWS_SESSION_TOKEN="$S3_SESSION_TOKEN"
    else
        unset AWS_SESSION_TOKEN
    fi
}

aws_s3_cmd() {
    local aws_cmd=(aws)

    "${aws_cmd[@]}" "$@"
}

aws_s3_prefix_keys() {
    local remote_uri=$1

    aws_s3_cmd s3 ls "$remote_uri/" --recursive |
        awk 'NF >= 4 {print $4}'
}

aws_s3_prefix_has_regex() {
    local remote_uri=$1
    local regex=$2

    aws_s3_prefix_keys "$remote_uri" | grep -E "$regex" >/dev/null
}

aws_s3_object_exists() {
    local remote_uri=$1

    aws_s3_cmd s3 ls "$remote_uri" >/dev/null 2>&1
}

aws_s3_cp_stdout() {
    local remote_uri=$1

    aws_s3_cmd s3 cp "$remote_uri" - --only-show-errors
}

aws_s3_rm_prefix() {
    local remote_prefix=$1
    local remote_uri="s3://$S3_BUCKET/$remote_prefix"

    log "+ aws s3 rm $remote_uri --recursive"
    aws_s3_cmd s3 rm "$remote_uri" --recursive --only-show-errors
}

build_ticdc_s3_sink_uri() {
    printf '%s?protocol=csv&output-old-value=true\n' "$increment_s3_uri"
}

write_connector_config() {
    export CONNECTOR_CONFIG_PATH="$connector_config_path"
    python3 - <<'PY'
from __future__ import annotations

import json
import os

bucket = os.environ["S3_BUCKET"]
prefix = os.environ["S3_RUN_PREFIX"].strip("/")
s3_url = f"s3://{bucket}/{prefix}"

config = {
    "sync_task_id": os.environ["SYNC_TASK_ID"],
    "s3_url": s3_url,
    "s3_region": os.environ["S3_REGION"],
    "s3_ak": os.environ["S3_AK"],
    "s3_sk": os.environ["S3_SK"],
    "snapshot_prefix": "snapshot",
    "incremental_prefix": "increment",
    "checkpoint_interval_rows": "10",
    "enable_diagnostic_columns": "true",
    "incremental": json.dumps({"csv_output_old_value": True}, separators=(",", ":")),
}

with open(os.environ["CONNECTOR_CONFIG_PATH"], "w", encoding="utf-8") as handle:
    json.dump(config, handle, indent=2)
    handle.write("\n")
PY
    chmod 600 "$connector_config_path"
}

fivetran_api_python() {
    local method=$1
    local path=$2
    local body=${3:-}

    python3 - "$method" "$path" "$body" <<'PY'
from __future__ import annotations

import json
import os
import sys
from urllib import error, request

method, path, body = sys.argv[1], sys.argv[2], sys.argv[3]
base_url = os.environ.get("FIVETRAN_BASE_URL", "https://api.fivetran.com/v1").rstrip("/")
headers = {
    "Authorization": f"Basic {os.environ['FIVETRAN_API_KEY_BASE64']}",
    "Accept": "application/json;version=2",
}
payload = None
if body:
    payload = body.encode("utf-8")
    headers["Content-Type"] = "application/json"
req = request.Request(f"{base_url}{path}", data=payload, headers=headers, method=method)
try:
    with request.urlopen(req, timeout=60) as resp:
        data = json.load(resp)
except error.HTTPError as exc:
    print(exc.read().decode("utf-8", errors="replace"), file=sys.stderr)
    raise
if data.get("code") != "Success":
    raise SystemExit(f"Fivetran API returned {data.get('code')}: {data.get('message')}")
print(json.dumps(data, separators=(",", ":")))
PY
}

fivetran_connection_id_by_name() {
    local group_id=$1
    local connection_name=$2

    python3 - "$group_id" "$connection_name" <<'PY'
from __future__ import annotations

import json
import os
import sys
from urllib import error, parse, request

group_id, connection_name = sys.argv[1], sys.argv[2]
base_url = os.environ.get("FIVETRAN_BASE_URL", "https://api.fivetran.com/v1").rstrip("/")
query = parse.urlencode({"group_id": group_id})
req = request.Request(
    f"{base_url}/connections?{query}",
    headers={
        "Authorization": f"Basic {os.environ['FIVETRAN_API_KEY_BASE64']}",
        "Accept": "application/json;version=2",
    },
    method="GET",
)
try:
    with request.urlopen(req, timeout=60) as resp:
        payload = json.load(resp)
except error.HTTPError as exc:
    print(exc.read().decode("utf-8", errors="replace"), file=sys.stderr)
    raise
if payload.get("code") != "Success":
    raise SystemExit(f"Fivetran API returned {payload.get('code')}: {payload.get('message')}")
items = payload["data"]["items"]
matches = [item for item in items if item.get("schema") == connection_name]
if not matches:
    raise SystemExit(f"connection {connection_name} not found in destination group {group_id}")
if len(matches) > 1:
    raise SystemExit(f"connection {connection_name} is ambiguous in destination group {group_id}")
print(matches[0]["id"])
PY
}

resolve_fivetran_connection_id() {
    local group_id

    group_id=$(
        python3 - <<'PY'
from __future__ import annotations

import json
import os
import sys
from urllib import error, request

base_url = os.environ.get("FIVETRAN_BASE_URL", "https://api.fivetran.com/v1").rstrip("/")
req = request.Request(
    f"{base_url}/groups",
    headers={
        "Authorization": f"Basic {os.environ['FIVETRAN_API_KEY_BASE64']}",
        "Accept": "application/json;version=2",
    },
    method="GET",
)
try:
    with request.urlopen(req, timeout=60) as resp:
        payload = json.load(resp)
except error.HTTPError as exc:
    print(exc.read().decode("utf-8", errors="replace"), file=sys.stderr)
    raise
if payload.get("code") != "Success":
    raise SystemExit(f"Fivetran API returned {payload.get('code')}: {payload.get('message')}")
items = payload["data"]["items"]
matches = [item for item in items if item.get("name") == os.environ["FIVETRAN_DESTINATION_NAME"]]
if not matches:
    raise SystemExit(f"destination {os.environ['FIVETRAN_DESTINATION_NAME']} not found")
if len(matches) > 1:
    raise SystemExit(f"destination {os.environ['FIVETRAN_DESTINATION_NAME']} is ambiguous")
print(matches[0]["id"])
PY
    )
    connection_id=$(fivetran_connection_id_by_name "$group_id" "$FIVETRAN_CONNECTION_NAME")
}

deploy_fivetran_connector() {
    write_connector_config
    run_redacted "fivetran deploy $repo_root --force --api-key *** --destination $FIVETRAN_DESTINATION_NAME --connection $FIVETRAN_CONNECTION_NAME --configuration $connector_config_path --python-version $FIVETRAN_PYTHON_VERSION" \
        "$TIMEOUT_BIN" "${FIVETRAN_DEPLOY_TIMEOUT_SECONDS}s" fivetran deploy "$repo_root" \
        --force \
        --api-key "$FIVETRAN_API_KEY_BASE64" \
        --destination "$FIVETRAN_DESTINATION_NAME" \
        --connection "$FIVETRAN_CONNECTION_NAME" \
        --configuration "$connector_config_path" \
        --python-version "$FIVETRAN_PYTHON_VERSION"
    resolve_fivetran_connection_id
}

set_fivetran_paused() {
    local paused=$1

    fivetran_api_python PATCH "/connections/$connection_id" "{\"paused\": $paused}" >/dev/null
}

reset_fivetran_state_if_requested() {
    if [ "$RESET_FIVETRAN_STATE" != "1" ]; then
        return 0
    fi

    log "+ Fivetran API reset connection state"
    fivetran_api_python PATCH "/connections/$connection_id/state" '{"state": {}}' >/dev/null
}

trigger_fivetran_sync() {
    local phase=$1
    local sync_body='{"force": true}'

    # Fivetran REST shape: /connections/{connection_id}/sync
    log "+ Fivetran API sync force=true ($phase)"
    set_fivetran_paused false
    fivetran_api_python POST "/connections/$connection_id/sync" "$sync_body" >/dev/null
}

wait_for_initial_snowflake_snapshot() {
    log "+ wait for Snowflake initial snapshot rows"
    run_snowflake_helper wait-count fruit "$SNOWFLAKE_SCHEMA" 10 "$SNOWFLAKE_READY_TIMEOUT_SECONDS"
}

export_tidb_comparison_profiles() {
    mkdir -p "$comparison_dir" || fail "failed to prepare comparison directory"

    mysql_export_tsv "
select
    cast(id as char),
    coalesce(name, '__NULL__'),
    coalesce(color, '__NULL__'),
    coalesce(cast(sweetness as char), '__NULL__')
from $db_name.$table_name
order by id;
" "$comparison_dir/fruit.tidb.tsv"

    mysql_export_tsv "
select
    cast(id as char),
    coalesce(c_varchar, '__NULL__'),
    coalesce(c_added_nullable, '__NULL__'),
    coalesce(c_default_toggle, '__NULL__'),
    coalesce(cast(c_changed_rename_type as char), '__NULL__'),
    coalesce(cast(c_type_cross_family_number as char), '__NULL__')
from ft_ddl_type_test.ft_ddl_type_matrix
order by id;
" "$comparison_dir/matrix.tidb.tsv"

    mysql_export_tsv "
select
    cast(id as char),
    coalesce(payload, '__NULL__')
from ft_ddl_type_test.ft_ddl_create_after_snapshot
order by id;
" "$comparison_dir/create.tidb.tsv"

    mysql_export_tsv "
select
    cast(id as char),
    coalesce(payload, '__NULL__')
from ft_ddl_type_test.ft_ddl_truncate_target
order by id;
" "$comparison_dir/truncate.tidb.tsv"
}

wait_for_snowflake_profile_match() {
    local profile=$1
    local expected=$2
    local actual="$comparison_dir/$profile.snowflake.tsv"
    local tmp_actual="$actual.tmp"
    local deadline=$((SECONDS + FIVETRAN_SYNC_TIMEOUT_SECONDS))
    local last_error="$comparison_dir/$profile.snowflake.err"

    while [ "$SECONDS" -lt "$deadline" ]; do
        if run_snowflake_helper export-profile "$profile" "$SNOWFLAKE_SCHEMA" "$tmp_actual" >"$last_error" 2>&1; then
            if cmp -s "$expected" "$tmp_actual"; then
                mv "$tmp_actual" "$actual"
                return 0
            fi
        fi
        sleep 10
    done

    log "Snowflake profile $profile did not match TiDB before timeout."
    if [ -s "$last_error" ]; then
        log "Last Snowflake helper error:"
        sed -n '1,120p' "$last_error" >&2 || true
    fi
    run_snowflake_helper export-profile "$profile" "$SNOWFLAKE_SCHEMA" "$actual" || true
    diff -u "$expected" "$actual" || true
    fail "Snowflake profile $profile does not match TiDB active rows"
}

compare_snowflake_to_tidb() {
    export_tidb_comparison_profiles
    wait_for_snowflake_profile_match fruit "$comparison_dir/fruit.tidb.tsv"
    wait_for_snowflake_profile_match matrix "$comparison_dir/matrix.tidb.tsv"
    wait_for_snowflake_profile_match create "$comparison_dir/create.tidb.tsv"
    wait_for_snowflake_profile_match truncate "$comparison_dir/truncate.tidb.tsv"

    run_snowflake_helper check-absent "$SNOWFLAKE_SCHEMA" "test_ddl_no_pk"
    run_snowflake_helper check-absent "$SNOWFLAKE_SCHEMA" "ft_ddl_type_test_ft_ddl_no_stable_pk"
}

activate_python
ensure_fivetran_cli
collect_interactive_inputs
log_runtime_configuration

for tool in tiup mysql aws "$TIMEOUT_BIN" grep tr sed awk diff python3; do
    require_command "$tool"
done

[ -f "$test_sql_path" ] || fail "missing $test_sql_path"
validate_identifier "$db_name" DB_NAME
validate_identifier "$table_name" TABLE_NAME
validate_port "$TIDB_PORT" TIDB_PORT
validate_ticdc_server
configure_aws_environment
snapshot_s3_uri="s3://$S3_BUCKET/$S3_RUN_PREFIX/snapshot"
increment_s3_uri="s3://$S3_BUCKET/$S3_RUN_PREFIX/increment"

rm -rf "$snapshot_dir" "$increment_dir" "$runtime_dir"
mkdir -p "$workspace" "$runtime_dir" || fail "failed to prepare workspace"
write_runtime_configs
resolve_cdc_cli
resolve_snowflake_python

tiup clean "$tag" >/dev/null 2>&1 || true

log "+ tiup playground $TIDB_VERSION --tag $tag --host $TIDB_HOST --db 1 --db.port $TIDB_PORT --pd 1 --kv 1 --tiflash 0 --ticdc 1 --ticdc.host $ticdc_host --ticdc.port $ticdc_port --ticdc.config $ticdc_config"
tiup playground "$TIDB_VERSION" --tag "$tag" --host "$TIDB_HOST" --db 1 --db.port "$TIDB_PORT" --pd 1 --kv 1 --tiflash 0 --ticdc 1 --ticdc.host "$ticdc_host" --ticdc.port "$ticdc_port" --ticdc.config "$ticdc_config" >"$playground_log" 2>&1 &
playground_pid=$!

wait_until "TiDB" "$READY_TIMEOUT_SECONDS" tidb_ready
wait_until "TiCDC" "$READY_TIMEOUT_SECONDS" ticdc_ready

seed_sql="drop database if exists $db_name; create database $db_name; use $db_name; create table fruit1(id bigint not null auto_increment primary key, name varchar(20) not null default ''); insert into fruit1(name) values ('apple'), ('banana'), ('cherry'), ('date'), ('elderberry'), ('fig'), ('grape'), ('honeydew'), ('kiwi'), ('lemon');"
mysql_exec "$seed_sql"

dumpling_cmd=(
    tiup dumpling
    -u "$TIDB_USER"
    -P "$TIDB_PORT"
    -h "$TIDB_HOST"
    --filetype csv
    -t "$DUMPLING_THREADS"
    -r "$DUMPLING_ROWS"
    -F "$DUMPLING_FILE_SIZE"
    -B "$db_name"
    --s3.region "$S3_REGION"
    -o "$snapshot_s3_uri"
)
if [ -n "$TIDB_PASSWORD" ]; then
    dumpling_cmd+=(-p "$TIDB_PASSWORD")
fi
run "${dumpling_cmd[@]}"

if ! aws_s3_object_exists "$snapshot_s3_uri/metadata" ||
    ! aws_s3_prefix_has_regex "$snapshot_s3_uri" "/$db_name\.$table_name\.[^/]*\.csv$"; then
    fail "Dumpling did not create expected CSV snapshot files under $snapshot_s3_uri"
fi

deploy_fivetran_connector
reset_fivetran_state_if_requested
trigger_fivetran_sync snapshot
wait_for_initial_snowflake_snapshot

sink_uri=$(build_ticdc_s3_sink_uri)
run_redacted "cdc changefeed create --server=$TICDC_SERVER --sink-uri=$increment_s3_uri?protocol=csv&output-old-value=true --changefeed-id=$changefeed_id --config=$changefeed_config" \
    "${cdc_cli[@]}" changefeed create --server="$TICDC_SERVER" --sink-uri="$sink_uri" --changefeed-id="$changefeed_id" --config="$changefeed_config" --no-confirm

apply_ddl_test_cases
prepare_test_sql_environment
mysql_exec_file "$test_sql_path"
wait_for_increment_outputs
wait_until "TiCDC final DDL row checkpoint" "$READY_TIMEOUT_SECONDS" final_ddl_row_checkpoint_ready
wait_until "TiCDC test.sql final row checkpoint" "$READY_TIMEOUT_SECONDS" test_sql_final_row_checkpoint_ready

trigger_fivetran_sync incremental
compare_snowflake_to_tidb

log "OK: Snowflake schema $SNOWFLAKE_SCHEMA matches TiDB active rows for the Snowflake E2E profiles"
