#!/usr/bin/env bash
set -euo pipefail

readonly SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd -P)"

source "$SCRIPT_DIR/tidb-data-lib.sh"

usage() {
  cat <<'EOF' >&2
Usage:
  scripts/tidb-import.sh <input-dir>

Restores a Dumpling schema-and-data dump into the database from
AGENTTEAM_DATABASE_DSN by using tiup tidb-lightning with the logical `tidb`
backend. The dump must include Dumpling schema files; the script generates a
temporary Lightning config that remaps the dumped database name onto the target
database in AGENTTEAM_DATABASE_DSN.

Environment variables:
  AGENTTEAM_DATABASE_DSN                 (required)
  AGENTTEAM_TIUP_BIN                     (default: tiup)
  AGENTTEAM_TIDB_LIGHTNING_COMPONENT     (default: tidb-lightning)
  AGENTTEAM_TIDB_IMPORT_LOG_LEVEL        (default: info)
  AGENTTEAM_TIDB_IMPORT_BACKEND          (default: tidb)
  AGENTTEAM_TIDB_TLS_CA_PATH             (optional)
  AGENTTEAM_TIDB_TLS_CERT_PATH           (optional)
  AGENTTEAM_TIDB_TLS_KEY_PATH            (optional)
EOF
}

agentteam_tidb_toml_escape_basic() {
  local value="$1"
  value="${value//\\/\\\\}"
  value="${value//\"/\\\"}"
  value="${value//$'\n'/\\n}"
  value="${value//$'\r'/\\r}"
  value="${value//$'\t'/\\t}"
  printf '%s' "$value"
}

agentteam_tidb_require_schema_files() {
  local dir_path="$1"

  if ! find "$dir_path" -type f \( -name '*-schema-create.sql' -o -name '*-schema-create.sql.gz' \) -print -quit | grep -q .; then
    agentteam_tidb_fail "input directory does not contain dumpling schema files: missing *-schema-create.sql in $dir_path"
  fi
  if ! find "$dir_path" -type f \( -name '*-schema.sql' -o -name '*-schema.sql.gz' \) -print -quit | grep -q .; then
    agentteam_tidb_fail "input directory does not contain dumpling schema files: missing *.table-schema.sql in $dir_path"
  fi
}

agentteam_tidb_write_lightning_config() {
  local config_path="$1"
  local input_dir="$2"
  local target_database="$3"
  local escaped_input_dir=""

  escaped_input_dir="$(agentteam_tidb_toml_escape_basic "$input_dir")"

  cat >"$config_path" <<EOF
[mydumper]
data-source-dir = "$escaped_input_dir"

[[mydumper.files]]
pattern = '^(.*?)-schema-create\.sql$'
schema = '$target_database'
type = "schema-schema"

[[mydumper.files]]
pattern = '^(.*?)\.(.*?)-schema\.sql$'
schema = '$target_database'
table = '\$2'
type = "table-schema"

[[mydumper.files]]
pattern = '^(.*?)\.(.*?)\.([0-9]+)\.(csv|parquet|sql)$'
schema = '$target_database'
table = '\$2'
type = '\$4'
key = '\$3'

[[mydumper.files]]
pattern = '^(.*?)-schema-create\.(sql)\.(gz)$'
schema = '$target_database'
type = "schema-schema"
compression = '$3'

[[mydumper.files]]
pattern = '^(.*?)\.(.*?)-schema\.(sql)\.(gz)$'
schema = '$target_database'
table = '\$2'
type = "table-schema"
compression = '\$4'

[[mydumper.files]]
pattern = '^(.*?)\.(.*?)\.([0-9]+)\.(sql|csv|parquet)\.(gz)$'
schema = '$target_database'
table = '\$2'
type = '\$4'
key = '\$3'
compression = '\$5'
EOF
}

case "${1:-}" in
  -h|--help|help)
    usage
    exit 0
    ;;
esac

if [[ $# -ne 1 ]]; then
  usage
  exit 1
fi

input_dir="$1"
agentteam_tidb_require_input_dir "$input_dir"
agentteam_tidb_require_schema_files "$input_dir"
agentteam_tidb_parse_core_storage_dsn "${AGENTTEAM_DATABASE_DSN:-${AGENTTEAM_CORE_STORAGE_DSN:-}}"
agentteam_tidb_validate_import_tls_support

config_file="$(mktemp "${TMPDIR:-/tmp}/agentteam-tidb-lightning.XXXXXX.toml")"
trap 'rm -f "$config_file"' EXIT
agentteam_tidb_write_lightning_config "$config_file" "$input_dir" "$agentteam_tidb_database"

tiup_bin="$(agentteam_tidb_resolve_tiup_bin)"
lightning_component="${AGENTTEAM_TIDB_LIGHTNING_COMPONENT:-tidb-lightning}"
log_level="${AGENTTEAM_TIDB_IMPORT_LOG_LEVEL:-info}"
backend="${AGENTTEAM_TIDB_IMPORT_BACKEND:-tidb}"

import_args=(
  "$lightning_component"
  "--backend" "$backend"
  "--config" "$config_file"
  "--tidb-host" "$agentteam_tidb_host"
  "--tidb-port" "$agentteam_tidb_port"
  "--tidb-user" "$agentteam_tidb_user"
  "-d" "$input_dir"
  "-L" "$log_level"
  "--log-file" "-"
)

if [[ -n "$agentteam_tidb_password" ]]; then
  import_args+=("--tidb-password" "$agentteam_tidb_password")
fi
agentteam_tidb_append_tls_file_flags import_args

"$tiup_bin" "${import_args[@]}"

echo "agentteam tidb import finished: database=$agentteam_tidb_database input=$input_dir"