#!/usr/bin/env bash

# NOTE(kwannoel):
# Do not run this script directly, it is meant to be sourced.
# Backwards compatibility tests consist of the following parts:
#
# 1. Setup old cluster binaries.
# 2. Seed old cluster.
# 3. Setup new cluster binaries.
# 4. Run validation on new cluster.
#
# Steps 1,3 are specific to the execution environment, CI / Local.
# This script only provides utilities for 2, 4.

################################### ENVIRONMENT CONFIG

# Duration to wait for recovery (seconds)
RECOVERY_DURATION=20

# Setup test directory
TEST_DIR=.risingwave/e2e_test/backwards-compat-tests/
# Keep this case on recently verified releases. It relies on Hummock system tables and
# risectl commands to build a mixed-table SST deterministically.
HUMMOCK_STALE_TABLE_IDS_MIN_VERSION=2.8.0
mkdir -p $TEST_DIR
cp -r e2e_test/backwards-compat-tests/slt/* $TEST_DIR

wait_for_process() {
  process_name="$1"

  while pgrep -x "$process_name" >/dev/null; do
    echo "Process $process_name is still running... Wait for 1 sec"
    sleep 1
  done
}

wait_all_process_exit() {
  wait_for_process meta-node
  wait_for_process compute-node
  wait_for_process frontend
  wait_for_process compactor
  echo "All processes has exited."
}

# Older versions of RW may not gracefully kill kafka.
# So we duplicate the definition here.
kill_cluster() {
  if tmux -L risedev ls &>/dev/null; then
    TMUX="tmux -L risedev"
  else
    TMUX="tmux"
  fi

  # Kill other components
  $TMUX list-windows -t risedev -F "#{window_name} #{pane_id}" |
    awk '{ print $2 }' |
    xargs -I {} $TMUX send-keys -t {} C-c C-d

  $TMUX kill-server
  test $? -eq 0 || {
    echo "Failed to stop all RiseDev components."
    exit 1
  }
  wait_all_process_exit
}

run_sql() {
  psql -h localhost -p 4566 -d dev -U root -c "$@"
}

run_sql_scalar() {
  psql -h localhost -p 4566 -d dev -U root -At -c "$@" | tr -d '[:space:]'
}

run_risectl() (
  set -euo pipefail
  set -a
  source .risingwave/config/risedev-env
  set +a
  .risingwave/bin/risingwave/risectl "$@"
)

check_version() {
  local VERSION=$1
  local raw_version=$(run_sql "SELECT version();")
  echo "--- Version"
  echo "raw_version: $raw_version"
  local version=$(echo $raw_version | grep -i risingwave | sed 's/^.*risingwave-\([0-9]*\.[0-9]*\.[0-9]\).*$/\1/i')
  if [[ "$version" != "$VERSION" ]]; then
    echo "Version mismatch, expected $VERSION, got $version"
    exit 1
  fi
}

create_kafka_topic() {
  RPK_BROKERS=message_queue:29092 \
  rpk topic create backwards_compat_test_kafka_source
}

insert_json_kafka() {
  local JSON=$1

  echo "$JSON" | \
  RPK_BROKERS=message_queue:29092 \
  rpk topic produce backwards_compat_test_kafka_source -f "%k,%v"
}

seed_json_kafka() {
  insert_json_kafka '{"user_id": 1},{"timestamp": "2023-07-28 07:11:00", "user_id": 1, "page_id": 1, "action": "gtrgretrg"}'
  insert_json_kafka '{"user_id": 2},{"timestamp": "2023-07-28 07:11:00", "user_id": 2, "page_id": 1, "action": "fsdfgerrg"}'
  insert_json_kafka '{"user_id": 3},{"timestamp": "2023-07-28 07:11:00", "user_id": 3, "page_id": 1, "action": "sdfergtth"}'
  insert_json_kafka '{"user_id": 4},{"timestamp": "2023-07-28 06:54:00", "user_id": 4, "page_id": 2, "action": "erwerhghj"}'
  insert_json_kafka '{"user_id": 5},{"timestamp": "2023-07-28 06:54:00", "user_id": 5, "page_id": 2, "action": "kiku7ikkk"}'
  insert_json_kafka '{"user_id": 6},{"timestamp": "2023-07-28 06:54:00", "user_id": 6, "page_id": 3, "action": "6786745ge"}'
  insert_json_kafka '{"user_id": 7},{"timestamp": "2023-07-28 06:54:00", "user_id": 7, "page_id": 3, "action": "fgbgfnyyy"}'
  insert_json_kafka '{"user_id": 8},{"timestamp": "2023-07-28 06:54:00", "user_id": 8, "page_id": 4, "action": "werwerwwe"}'
  insert_json_kafka '{"user_id": 9},{"timestamp": "2023-07-28 06:54:00", "user_id": 9, "page_id": 4, "action": "yjtyjtyyy"}'
}

seed_hummock_stale_table_ids() {
  rm -f "$TEST_DIR/hummock-stale-table-ids/enabled" \
    "$TEST_DIR/hummock-stale-table-ids/dropped_table_id"

  if version_lt "$OLD_VERSION" "$HUMMOCK_STALE_TABLE_IDS_MIN_VERSION"; then
    echo "--- HUMMOCK STALE TABLE IDS TEST: Skipped for old version ${OLD_VERSION}"
    return
  fi

  echo "--- HUMMOCK STALE TABLE IDS TEST: Seeding old cluster with mixed-table SST"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/hummock-stale-table-ids/seed.slt"

  local live_table_id
  live_table_id=$(run_sql_scalar "SELECT id FROM rw_catalog.rw_tables WHERE name = 'hummock_stale_live';")
  local dropped_table_id
  dropped_table_id=$(run_sql_scalar "SELECT id FROM rw_catalog.rw_tables WHERE name = 'hummock_stale_dropped';")
  if [[ -z "$live_table_id" || -z "$dropped_table_id" ]]; then
    echo "Failed to capture hummock stale table ids: live=${live_table_id}, dropped=${dropped_table_id}"
    exit 1
  fi

  local live_compaction_group_id
  live_compaction_group_id=$(run_sql_scalar "SELECT compaction_group_id FROM rw_catalog.rw_hummock_sstables WHERE table_ids @> '[${live_table_id}]'::jsonb LIMIT 1;")
  local dropped_compaction_group_id
  dropped_compaction_group_id=$(run_sql_scalar "SELECT compaction_group_id FROM rw_catalog.rw_hummock_sstables WHERE table_ids @> '[${dropped_table_id}]'::jsonb LIMIT 1;")
  if [[ -z "$live_compaction_group_id" || -z "$dropped_compaction_group_id" ]]; then
    echo "Failed to capture hummock compaction groups: live=${live_compaction_group_id}, dropped=${dropped_compaction_group_id}"
    exit 1
  fi

  # Put both tables into one compaction group so a manual L0 compaction can
  # rewrite their SST metadata into one mixed-table SST.
  if [[ "$live_compaction_group_id" != "$dropped_compaction_group_id" ]]; then
    run_risectl hummock merge-compaction-group \
      --left-group-id "$live_compaction_group_id" \
      --right-group-id "$dropped_compaction_group_id"
  fi

  # Target the exact L0 SSTs that contain either table id. This avoids depending
  # on the default manual selector deciding that there is enough data to compact.
  local target_sst_ids
  target_sst_ids=$(run_sql_scalar "
    SELECT string_agg(sstable_id::varchar, ',' ORDER BY sub_level_id, sstable_id)
    FROM rw_catalog.rw_hummock_sstables
    WHERE compaction_group_id = ${live_compaction_group_id}
      AND level_id = 0
      AND (
        table_ids @> '[${live_table_id}]'::jsonb
        OR table_ids @> '[${dropped_table_id}]'::jsonb
      );
  ")
  if [[ -z "$target_sst_ids" ]]; then
    echo "Failed to find L0 SSTs for hummock stale table ids test: live_table_id=${live_table_id}, dropped_table_id=${dropped_table_id}, compaction_group_id=${live_compaction_group_id}"
    exit 1
  fi

  # `risectl` does not expose a reliable shell exit status for "no task picked",
  # so the test verifies success by observing the resulting SST metadata below.
  run_risectl hummock trigger-manual-compaction \
    --compaction-group-id "$live_compaction_group_id" \
    --level 0 \
    --sst-ids "$target_sst_ids"

  # Fail closed unless the current Hummock version really contains an SST whose
  # metadata references both the live table and the table that will be dropped.
  local mixed_sst_count
  for _ in $(seq 1 60); do
    mixed_sst_count=$(run_sql_scalar "SELECT count(*) FROM rw_catalog.rw_hummock_sstables WHERE table_ids @> '[${live_table_id}]'::jsonb AND table_ids @> '[${dropped_table_id}]'::jsonb;")
    if [[ "$mixed_sst_count" != "0" ]]; then
      break
    fi
    sleep 1
  done
  if [[ "$mixed_sst_count" == "0" ]]; then
    echo "Failed to build mixed-table SST for hummock stale table ids test"
    exit 1
  fi

  echo "$dropped_table_id" > "$TEST_DIR/hummock-stale-table-ids/dropped_table_id"

  # Simulate the old-version upgrade state: the dropped table is unregistered
  # from compaction group configs, while old SST metadata can still contain it.
  run_sql "DROP TABLE hummock_stale_dropped;"
  run_sql "FLUSH;"

  local registered_stale_table_count
  registered_stale_table_count=$(run_sql_scalar "SELECT count(*) FROM rw_catalog.rw_hummock_compaction_group_configs WHERE member_tables @> '[${dropped_table_id}]'::jsonb;")
  if [[ "$registered_stale_table_count" != "0" ]]; then
    echo "Old cluster did not unregister dropped table id ${dropped_table_id}"
    exit 1
  fi

  local stale_sst_count
  stale_sst_count=$(run_sql_scalar "SELECT count(*) FROM rw_catalog.rw_hummock_sstables WHERE table_ids @> '[${dropped_table_id}]'::jsonb;")
  if [[ "$stale_sst_count" == "0" ]]; then
    echo "Old cluster did not retain stale SST table id ${dropped_table_id}; the backwards compatibility test would be ineffective"
    exit 1
  fi

  echo "--- HUMMOCK STALE TABLE IDS TEST: Validating old cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/hummock-stale-table-ids/validate_original.slt"
  touch "$TEST_DIR/hummock-stale-table-ids/enabled"
}

validate_hummock_stale_table_ids() {
  if [[ ! -f "$TEST_DIR/hummock-stale-table-ids/enabled" ]]; then
    echo "--- HUMMOCK STALE TABLE IDS TEST: Skipped"
    return
  fi

  echo "--- HUMMOCK STALE TABLE IDS TEST: Validating new cluster before compaction"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/hummock-stale-table-ids/validate_restart.slt"

  local dropped_table_id
  dropped_table_id=$(cat "$TEST_DIR/hummock-stale-table-ids/dropped_table_id")
  local stale_sst_count
  stale_sst_count=$(run_sql_scalar "SELECT count(*) FROM rw_catalog.rw_hummock_sstables WHERE table_ids @> '[${dropped_table_id}]'::jsonb;")
  if [[ "$stale_sst_count" != "0" ]]; then
    echo "Recovered new cluster still has SST metadata containing dropped table id ${dropped_table_id}"
    exit 1
  fi
}

# https://stackoverflow.com/a/4024263
version_le() {
  printf '%s\n' "$1" "$2" | sort -C -V
}

version_lt() {
  ! version_le "$2" "$1"
}

################################### Entry Points

get_old_version() {
  # For backwards compat test we assume we are testing the latest version of RW (i.e. latest main commit)
  # against the Nth latest release candidate, where N > 1. N can be larger,
  # in case some old cluster did not upgrade.
  if [[ -z ${VERSION_OFFSET:-} ]]; then
    local VERSION_OFFSET=1
  fi

  # First get the new version to offset against
  get_new_version
  echo "--- NEW VERSION: $NEW_VERSION"

  # Extract X.Y.Z tags
  echo "--- VERSION BRANCHES"
  local tags=$(git tag | grep -E "^v[0-9]+\.[0-9]+\.[0-9]+$" | tr -d 'v' | tr -d ' ')
  echo "$tags"

  # Then we sort them in descending order.
  echo "--- VERSIONS"
  local sorted_versions=$(echo -e "$tags" | sort -V)
  echo "$sorted_versions"

  # Find the index of NEW_VERSION in the sorted list
  local new_version_index=$(echo -e "$sorted_versions" | grep -n "$NEW_VERSION" | cut -d: -f1)
  if [[ -z $new_version_index ]]; then
    echo "Could not find NEW_VERSION ($NEW_VERSION) in git tags, looking for latest version in same minor series"
    # Extract major.minor from NEW_VERSION
    local major_minor=$(echo "$NEW_VERSION" | cut -d. -f1,2)
    # Find the latest version in the same minor series
    local latest_in_series=$(echo -e "$sorted_versions" | grep "^$major_minor\." | tail -n1)
    if [[ -n "$latest_in_series" ]]; then
      echo "Found latest version in series $major_minor: $latest_in_series"
      new_version_index=$(echo -e "$sorted_versions" | grep -n "$latest_in_series" | cut -d: -f1)
    else
      echo "No version found in series $major_minor, using latest tag as reference"
      # Get the total number of versions
      local total_versions=$(echo -e "$sorted_versions" | wc -l)
      # Use the latest version's index + 1 as reference
      new_version_index=$((total_versions + 1))
    fi
    echo "Using reference index: $new_version_index"
  fi

  # Calculate the target index by subtracting the offset
  local target_index=$((new_version_index - VERSION_OFFSET))

  # Get the version at the target index
  OLD_VERSION=$(echo -e "$sorted_versions" | sed -n "${target_index}p")
  if [[ -z $OLD_VERSION ]]; then
    echo "Error: Could not find version at offset $VERSION_OFFSET from reference version"
    exit 1
  fi
  echo "--- OLD VERSION: $OLD_VERSION"
}

get_new_version() {
  # Next, for $NEW_VERSION we just scrape it from `workspace.package.version`.
  NEW_VERSION=$(cat Cargo.toml | grep "\[workspace\.package\]" -A 5 | sed -n 's/version = \"\([0-9]*\.[0-9]*\.[0-9]*\).*/\1/p' | tr -d ' ')
}

# Get $OLD_VERSION and $NEW_VERSION for Risingwave
get_rw_versions() {
  get_old_version
  get_new_version

  if version_le "$OLD_VERSION" "$NEW_VERSION"
  then
    echo "OLD_VERSION: $OLD_VERSION"
    echo "NEW_VERSION: $NEW_VERSION"
  else
    echo "ERROR: $OLD_VERSION >= $NEW_VERSION"
    exit 1
  fi
}

# Setup table and materialized view.
# Run updates and deletes on the table.
# Get the results.
seed_old_cluster() {
  echo "--- Start cluster on old_version: $OLD_VERSION"
  # Caller should make sure the test env has these.
  # They are called here because the current tests
  # may not be backwards compatible, so we need to call
  # them in old cluster environment.
  cp -r e2e_test/streaming/nexmark $TEST_DIR
  cp -r e2e_test/nexmark/* $TEST_DIR/nexmark

  cp -r e2e_test/batch/tpch $TEST_DIR
  cp -r e2e_test/tpch/* $TEST_DIR/tpch

  ./risedev clean-data
  # `ENABLE_PYTHON_UDF` and `ENABLE_JS_UDF` are set for backwards-compartibility
  ENABLE_PYTHON_UDF=1 ENABLE_JS_UDF=1 ENABLE_UDF=1 ./risedev d full-without-monitoring && rm .risingwave/log/*

  check_version "$OLD_VERSION"

  echo "--- BASIC TEST: Seeding old cluster with data"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/basic/seed.slt"

  echo "--- BASIC TEST: Validating old cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/basic/validate_original.slt"

  echo "--- HASH JOIN WATERMARK TEST: Seeding old cluster with data"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/hash-join-watermark/seed.slt"

  echo "--- HASH JOIN WATERMARK TEST: Validating old cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/hash-join-watermark/validate_original.slt"

  echo "--- EOWC OVER WINDOW NUMBERING TEST: Seeding old cluster with data"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/eowc-over-window-numbering/seed.slt"

  echo "--- EOWC OVER WINDOW NUMBERING TEST: Validating old cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/eowc-over-window-numbering/validate_original.slt"

  echo "--- NEXMARK TEST: Seeding old cluster with data"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/nexmark-backwards-compat/seed.slt"

  echo "--- NEXMARK TEST: Validating old cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/nexmark-backwards-compat/validate_original.slt"

  echo "--- TPCH TEST: Seeding old cluster with data"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/tpch-backwards-compat/seed.slt"

  echo "--- TPCH TEST: Validating old cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/tpch-backwards-compat/validate_original.slt"

  echo "--- KAFKA TEST: Seeding old cluster with data"
  create_kafka_topic
  seed_json_kafka
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/kafka/seed.slt"
  # use the old syntax for version at most 1.5.4
  if version_le "$OLD_VERSION" "1.5.4"; then
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/kafka/upsert/deprecate_upsert.slt"
  else
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/kafka/upsert/include_key_as.slt"
  fi

  echo "--- KAFKA TEST: wait 5s for kafka to process data"
  sleep 5

  echo "--- KAFKA TEST: Validating old cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/kafka/validate_original.slt"

  # Test version columns backwards compatibility, if OLD_VERSION <= 2.6.0
  if version_le "$OLD_VERSION" "2.6.0"; then
    echo "--- VERSION COLUMNS TEST: Seeding old cluster with data"
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/version-columns/seed.slt"

    echo "--- VERSION COLUMNS TEST: Validating old cluster"
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/version-columns/validate_original.slt"
  fi

  # Test legacy streaming parallelism session parameter migration for versions in [2.8.0, 2.9.0).
  if version_le "2.8.0" "$OLD_VERSION" && version_lt "$OLD_VERSION" "2.9.0"; then
    echo "--- STREAMING PARALLELISM TEST: Seeding old cluster"
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/streaming-parallelism/seed.slt"

    echo "--- STREAMING PARALLELISM TEST: Validating old cluster"
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/streaming-parallelism/validate_original.slt"
  fi

  # Test invalid WITH options, if OLD_VERSION <= 1.5.0
  if version_le "$OLD_VERSION" "1.5.0"; then
    echo "--- KAFKA TEST (invalid options): Seeding old cluster with data"
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/kafka/invalid_options/seed.slt"

    echo "--- KAFKA TEST (invalid options): wait 5s for kafka to process data"
    sleep 5

    echo "--- KAFKA TEST (invalid options): Validating old cluster"
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/kafka/invalid_options/validate_original.slt"
  fi

  echo "--- SINK INTO TABLE TEST: Seeding old cluster with data"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/sink_into_table/seed.slt"

  echo "--- SINK INTO TABLE TEST: Validating old cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/sink_into_table/validate_original.slt"

  echo "--- ASOF JOIN TEST: Seeding old cluster with data"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/asof-join/seed.slt"

  echo "--- ASOF JOIN TEST: Validating old cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/asof-join/validate_original.slt"

  echo "--- CDC TEST: Seeding old cluster with data"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/cdc/seed.slt"

  echo "--- CDC TEST: Validating old cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/cdc/validate_original.slt"

  # work around https://github.com/risingwavelabs/risingwave/issues/18650
  echo "--- wait for a version checkpoint"
  sleep 60

  seed_hummock_stale_table_ids

  echo "--- Killing cluster"
  kill_cluster
  echo "--- Killed cluster"
}

validate_new_cluster() {
  echo "--- Start cluster on latest"
  ENABLE_UDF=1 ./risedev d full-without-monitoring

  echo "--- Wait ${RECOVERY_DURATION}s for Recovery on Old Cluster Data"
  sleep $RECOVERY_DURATION

  check_version "$NEW_VERSION"

  echo "--- BASIC TEST: Validating new cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/basic/validate_restart.slt"

  echo "--- HASH JOIN WATERMARK TEST: Validating new cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/hash-join-watermark/validate_restart.slt"

  echo "--- EOWC OVER WINDOW NUMBERING TEST: Validating new cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/eowc-over-window-numbering/validate_restart.slt"

  echo "--- NEXMARK TEST: Validating new cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/nexmark-backwards-compat/validate_restart.slt"

  echo "--- TPCH TEST: Validating new cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/tpch-backwards-compat/validate_restart.slt"

  echo "--- KAFKA TEST: Seeding new cluster with data"
  seed_json_kafka

  echo "--- KAFKA TEST: Validating new cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/kafka/validate_restart.slt"

  # Test version columns backwards compatibility, if OLD_VERSION <= 2.6.0
  if version_le "$OLD_VERSION" "2.6.0"; then
    echo "--- VERSION COLUMNS TEST: Validating new cluster"
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/version-columns/validate_restart.slt"
  fi

  # Test legacy streaming parallelism session parameter migration for versions in [2.8.0, 2.9.0).
  if version_le "2.8.0" "$OLD_VERSION" && version_lt "$OLD_VERSION" "2.9.0"; then
    echo "--- STREAMING PARALLELISM TEST: Validating new cluster"
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/streaming-parallelism/validate_restart.slt"
  fi

  # Test invalid WITH options, if OLD_VERSION <= 1.5.0
  if version_le "$OLD_VERSION" "1.5.0"; then
    echo "--- KAFKA TEST (invalid options): Validating new cluster"
    sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/kafka/invalid_options/validate_restart.slt"
  fi

  echo "--- SINK INTO TABLE TEST: Validating new cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/sink_into_table/validate_restart.slt"

  echo "--- ASOF JOIN TEST: Validating new cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/asof-join/validate_restart.slt"

  echo "--- CDC TEST: Validating new cluster"
  sqllogictest -d dev -h localhost -p 4566 "$TEST_DIR/cdc/validate_restart.slt"

  validate_hummock_stale_table_ids

  kill_cluster
}
