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

# Uso:
#   ALLOW_DROP_DATABASE=1 ./mysqlsh_restore_benchmark.sh <DUMP_DIR> <DATABASE> [SKIP_VIEWS_TRIGGERS]

THREADS="${THREADS:-16}"
DUMP_DIR="${1:-}"
DATABASE="${2:-}"
SKIP_VIEWS_TRIGGERS="${3:-1}"

if [[ -z "$DUMP_DIR" || -z "$DATABASE" ]]; then
  echo "Uso: $0 <DUMP_DIR> <DATABASE> [SKIP_VIEWS_TRIGGERS]"
  exit 1
fi

if [[ ! -d "$DUMP_DIR" ]]; then
  echo "Error: no existe el directorio de dump: $DUMP_DIR"
  exit 1
fi

if [[ ! "$THREADS" =~ ^[0-9]+$ ]] || [[ "$THREADS" -lt 1 ]] || [[ "$THREADS" -gt 64 ]]; then
  echo "Error: THREADS debe ser un entero entre 1 y 64"
  exit 1
fi

DB_USER="${DB_USER:-}"
DB_PASSWORD="${DB_PASSWORD:-}"
DB_HOST="${DB_HOST:-}"
DB_PORT="${DB_PORT:-}"
MYSQL_CNF="${MYSQL_CNF:-}"
MYSQLSH_BIN="${MYSQLSH_BIN:-mysqlsh}"

if ! command -v "$MYSQLSH_BIN" >/dev/null 2>&1; then
  echo "Error: no se encontró MYSQLSH_BIN=$MYSQLSH_BIN"
  exit 1
fi

if [[ -n "$MYSQL_CNF" ]]; then
  if [[ ! -f "$MYSQL_CNF" ]]; then
    echo "Error: MYSQL_CNF no existe: $MYSQL_CNF"
    exit 1
  fi
  if [[ ! -r "$MYSQL_CNF" ]]; then
    echo "Error: MYSQL_CNF sin permiso de lectura: $MYSQL_CNF"
    echo "Sugerencia: ajustar permisos/grupo para el usuario del proceso web"
    exit 1
  fi
  DB_USER="$(awk -F= '/^[[:space:]]*user[[:space:]]*=/{gsub(/^[[:space:]]+|[[:space:]]+$/, "", $2); print $2; exit}' "$MYSQL_CNF" || true)"
  DB_PASSWORD="$(awk -F= '/^[[:space:]]*password[[:space:]]*=/{gsub(/^[[:space:]]+|[[:space:]]+$/, "", $2); print $2; exit}' "$MYSQL_CNF" || true)"
  DB_HOST="$(awk -F= '/^[[:space:]]*host[[:space:]]*=/{gsub(/^[[:space:]]+|[[:space:]]+$/, "", $2); print $2; exit}' "$MYSQL_CNF" || true)"
  DB_PORT="$(awk -F= '/^[[:space:]]*port[[:space:]]*=/{gsub(/^[[:space:]]+|[[:space:]]+$/, "", $2); print $2; exit}' "$MYSQL_CNF" || true)"
fi

if [[ -z "$DB_USER" || -z "$DB_PASSWORD" || -z "$DB_HOST" || -z "$DB_PORT" ]]; then
  echo "Error: faltan credenciales/host/puerto para conexión MySQL"
  exit 1
fi

ALLOW_DROP_DATABASE="${ALLOW_DROP_DATABASE:-0}"
if [[ "$ALLOW_DROP_DATABASE" != "1" ]]; then
  echo "Error: restore cancelado para evitar DROP accidental."
  echo "Para continuar, ejecuta con ALLOW_DROP_DATABASE=1"
  exit 1
fi

BENCHMARK_FALLBACK_TO_HARDENED="${BENCHMARK_FALLBACK_TO_HARDENED:-0}"
if [[ "$BENCHMARK_FALLBACK_TO_HARDENED" != "0" && "$BENCHMARK_FALLBACK_TO_HARDENED" != "1" ]]; then
  echo "Error: BENCHMARK_FALLBACK_TO_HARDENED debe ser 0 o 1"
  exit 1
fi

js_quote() {
  local value="$1"
  value="${value//\\/\\\\}"
  value="${value//\'/\\\'}"
  printf "'%s'" "$value"
}

dialect_for_file() {
  case "$1" in
    *.tsv|*.tsv.gz|*.tsv.zst) echo "tsv" ;;
    *.csv|*.csv.gz|*.csv.zst) echo "csv-unix" ;;
    *.txt|*.txt.gz|*.txt.zst) echo "default" ;;
    *) echo "default" ;;
  esac
}

table_for_file() {
  local file_name base_name logical_name table_name

  file_name="$(basename "$1")"
  base_name="$file_name"

  # Quita extensión de compresión cuando aplica.
  case "$base_name" in
    *.zst|*.gz) base_name="${base_name%.*}" ;;
  esac

  # Quita extensión de formato de datos.
  case "$base_name" in
    *.tsv|*.csv|*.txt) logical_name="${base_name%.*}" ;;
    *) logical_name="$base_name" ;;
  esac

  # Para archivos chunked tipo schema@tabla@@0, conservar solo schema@tabla.
  logical_name="${logical_name%%@@*}"

  # Extrae la tabla del patrón schema@table (preferir schema objetivo si coincide).
  if [[ "$logical_name" == "$TARGET_DB"@* ]]; then
    table_name="${logical_name#"$TARGET_DB"@}"
  elif [[ "$logical_name" == *"@"* ]]; then
    table_name="${logical_name#*@}"
  else
    table_name="$logical_name"
  fi

  # Algunos dumps generan chunks con sufijo @0 (sin doble @@).
  # Ejemplo: schema@tabla@0.tsv.zst -> tabla
  if [[ "$table_name" =~ @[0-9]+$ ]]; then
    table_name="${table_name%@*}"
  fi

  printf '%s' "$table_name"
}

export TZ="America/Guatemala"
TARGET_DB="$DATABASE"

if [[ ! "$TARGET_DB" =~ ^[A-Za-z0-9_]+$ ]]; then
  echo "Error: nombre de base de datos invalido: $TARGET_DB"
  exit 1
fi

case "$TARGET_DB" in
  mysql|sys|information_schema|performance_schema)
    echo "Error: base protegida, restore cancelado: $TARGET_DB"
    exit 1
    ;;
esac

WORK_DIR="/tmp/load_dump_${TARGET_DB}"
DDL_FILE="/tmp/mysqlsh_restore_benchmark_${TARGET_DB}.sql"
JS_FILE="/tmp/mysqlsh_restore_benchmark_${TARGET_DB}.js"
DDL_BOOTSTRAP_JS_FILE="/tmp/mysqlsh_restore_benchmark_ddl_bootstrap_${TARGET_DB}.js"

append_ddl_content() {
  local ddl_path="$1"

  case "$ddl_path" in
    *.sql)
      cat "$ddl_path"
      ;;
    *.sql.zst)
      if ! command -v zstd >/dev/null 2>&1; then
        echo "Error: zstd no está instalado para leer DDL comprimido: $ddl_path"
        exit 1
      fi
      zstd -dc "$ddl_path"
      ;;
    *.sql.gz)
      gzip -dc "$ddl_path"
      ;;
    *)
      echo "Error: formato DDL no soportado: $ddl_path"
      exit 1
      ;;
  esac
}

normalize_ddl_stream() {
  local source_schema="$1"
  local target_schema="$2"

  local -a sed_rules=()
  sed_rules+=(
    -e 's/\/\*!([0-9]{5}) DEFINER=`[^`]+`@`[^`]+` \*\//\/\* stripped DEFINER *\//g'
    -e 's/DEFINER=`[^`]+`@`[^`]+`[ ]*//g'
    -e 's/SQL SECURITY DEFINER/SQL SECURITY INVOKER/g'
  )

  if [[ -n "$source_schema" && "$source_schema" != "$target_schema" ]]; then
    sed_rules+=(
      -e "s/CREATE DATABASE IF NOT EXISTS \\\`$source_schema\\\`/CREATE DATABASE IF NOT EXISTS \\\`$target_schema\\\`/g"
      -e "s/USE \\\`$source_schema\\\`/USE \\\`$target_schema\\\`/g"
      -e "s/\\\`$source_schema\\\`\./\\\`$target_schema\\\`\./g"
    )
  fi

  if [[ "$SKIP_VIEWS_TRIGGERS" == "1" ]]; then
    sed -E "${sed_rules[@]}" \
      -e '/^-- begin trigger /,/^-- end trigger /d' \
      -e '/^-- begin view /,/^-- end view /d'
  else
    sed -E "${sed_rules[@]}"
  fi
}

rm -rf "$WORK_DIR"
mkdir -p "$WORK_DIR"

echo "Restore benchmark hacia: ${DB_USER}@${DB_HOST}:${DB_PORT}"
echo "Dump origen: $DUMP_DIR"
echo "Base destino: $TARGET_DB"
echo "Hilos: $THREADS"

cp -a "$DUMP_DIR"/. "$WORK_DIR"/

SOURCE_SCHEMA=""

# 1) Detectar por nombres de archivos de datos (schema@tabla@@chunk.ext)
first_data_file="$(find "$WORK_DIR" -type f \( -name '*@*@@*.tsv' -o -name '*@*@@*.tsv.zst' -o -name '*@*@@*.tsv.gz' -o -name '*@*@@*.csv' -o -name '*@*@@*.csv.zst' -o -name '*@*@@*.csv.gz' -o -name '*@*@@*.txt' -o -name '*@*@@*.txt.zst' -o -name '*@*@@*.txt.gz' \) | sort | head -n 1 || true)"
if [[ -n "$first_data_file" ]]; then
  first_data_name="$(basename "$first_data_file")"
  SOURCE_SCHEMA="${first_data_name%%@*}"
fi

# 2) Detectar desde @.json (portable, sin extensiones de awk no POSIX)
if [[ -z "$SOURCE_SCHEMA" && -f "$WORK_DIR/@.json" ]]; then
  SOURCE_SCHEMA="$(sed -nE 's/.*"schemas"[[:space:]]*:[[:space:]]*\[[[:space:]]*"([^"]+)".*/\1/p' "$WORK_DIR/@.json" | head -n 1 || true)"
fi

if [[ -z "$SOURCE_SCHEMA" ]]; then
  first_schema_file="$(find "$WORK_DIR" -type f -name '*@*.sql*' | sort | head -n 1 || true)"
  if [[ -n "$first_schema_file" ]]; then
    first_schema_name="$(basename "$first_schema_file")"
    SOURCE_SCHEMA="${first_schema_name%%@*}"
  fi
fi

echo "Esquema origen detectado: ${SOURCE_SCHEMA:-desconocido}"
echo "Esquema destino aplicado: $TARGET_DB"

MYSQL_PWD="$DB_PASSWORD" mysql -h "$DB_HOST" -u "$DB_USER" -P "$DB_PORT" -e \
  "DROP DATABASE IF EXISTS \`$TARGET_DB\`; CREATE DATABASE \`$TARGET_DB\`;"

mapfile -d '' ddl_files < <(
  find "$WORK_DIR" -type f \( -name '*.sql' -o -name '*.sql.zst' -o -name '*.sql.gz' \) -print0 | sort -z
)
if [[ "${#ddl_files[@]}" -gt 0 ]]; then
  : > "$DDL_FILE"
  for ddl in "${ddl_files[@]}"; do
    append_ddl_content "$ddl" | normalize_ddl_stream "$SOURCE_SCHEMA" "$TARGET_DB" >> "$DDL_FILE"
    printf '\n' >> "$DDL_FILE"
  done
  MYSQL_PWD="$DB_PASSWORD" mysql -h "$DB_HOST" -u "$DB_USER" -P "$DB_PORT" "$TARGET_DB" < "$DDL_FILE"

  tables_created="$(MYSQL_PWD="$DB_PASSWORD" mysql -N -s -h "$DB_HOST" -u "$DB_USER" -P "$DB_PORT" -e "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema='${TARGET_DB}'" 2>/dev/null || echo 0)"
  if [[ -z "$tables_created" || ! "$tables_created" =~ ^[0-9]+$ ]]; then
    tables_created=0
  fi
  echo "Tablas detectadas tras DDL: $tables_created"
  if [[ "$tables_created" -eq 0 ]]; then
    echo "Aviso: el DDL manual no creó tablas. Intentando bootstrap de esquema con util.loadDump (solo DDL)..."

    {
      printf 'var targetSchema = %s;\n' "$(js_quote "$TARGET_DB")"
      printf 'var sourceSchema = %s;\n' "$(js_quote "${SOURCE_SCHEMA:-}")"
      printf 'var sourceDump = %s;\n' "$(js_quote "$WORK_DIR")"
      printf 'var loadOptions = {\n'
      printf '  ignoreVersion: true,\n'
      printf '  resetProgress: true,\n'
      printf '  loadDdl: true,\n'
      printf '  loadData: false,\n'
      printf '  loadUsers: false,\n'
      printf '  showProgress: true\n'
      printf '};\n'
      printf 'if (sourceSchema && sourceSchema !== targetSchema) {\n'
      printf '  loadOptions.includeSchemas = [sourceSchema];\n'
      printf '}\n'
      printf 'loadOptions.schema = targetSchema;\n'
      printf 'try {\n'
      printf '  util.loadDump(sourceDump, loadOptions);\n'
      printf '} catch (e) {\n'
      printf '  var msg = (e && e.message) ? String(e.message) : "";\n'
      printf '  var canRetry = loadOptions.includeSchemas && (msg.indexOf("0 will be loaded") !== -1 || msg.indexOf("Invalid option: schema") !== -1);\n'
      printf '  if (!canRetry) {\n'
      printf '    throw e;\n'
      printf '  }\n'
      printf '  print("Aviso: loadDump DDL falló con includeSchemas; reintentando sin includeSchemas...");\n'
      printf '  var retryOptions = {\n'
      printf '    ignoreVersion: loadOptions.ignoreVersion,\n'
      printf '    resetProgress: loadOptions.resetProgress,\n'
      printf '    schema: loadOptions.schema,\n'
      printf '    loadDdl: loadOptions.loadDdl,\n'
      printf '    loadData: loadOptions.loadData,\n'
      printf '    loadUsers: loadOptions.loadUsers,\n'
      printf '    showProgress: loadOptions.showProgress\n'
      printf '  };\n'
      printf '  util.loadDump(sourceDump, retryOptions);\n'
      printf '}\n'
    } > "$DDL_BOOTSTRAP_JS_FILE"

    "$MYSQLSH_BIN" "${DB_USER}:${DB_PASSWORD}@${DB_HOST}:${DB_PORT}" --verbose=0 --js < "$DDL_BOOTSTRAP_JS_FILE"

    tables_created="$(MYSQL_PWD="$DB_PASSWORD" mysql -N -s -h "$DB_HOST" -u "$DB_USER" -P "$DB_PORT" -e "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema='${TARGET_DB}'" 2>/dev/null || echo 0)"
    if [[ -z "$tables_created" || ! "$tables_created" =~ ^[0-9]+$ ]]; then
      tables_created=0
    fi
    echo "Tablas detectadas tras bootstrap DDL: $tables_created"

    if [[ "$tables_created" -eq 0 ]]; then
      echo "Aviso: benchmark no logró crear tablas tras DDL/manual bootstrap."
      echo "Diagnóstico: util.loadDump no materializó estructuras; revisa metadatos del dump y compatibilidad MySQL Shell."

      if [[ "$BENCHMARK_FALLBACK_TO_HARDENED" != "1" ]]; then
        echo "Error: benchmark detenido para conservar diagnóstico de causa raíz."
        echo "Sugerencia: corrige el dump o ejecuta con BENCHMARK_FALLBACK_TO_HARDENED=1 para continuar con restore normal."
        rm -rf "$WORK_DIR" "$DDL_FILE" "$JS_FILE" "$DDL_BOOTSTRAP_JS_FILE"
        exit 1
      fi

      echo "Fallback habilitado: se ejecutará restore normal para completar la carga."

      fallback_hardened=""
      for candidate in \
        "$(cd "$(dirname "$0")" && pwd)/mysqlsh_restore_hardened.sh" \
        "/opt/scripts/mysqlsh_restore_hardened.sh" \
        "/usr/bin/mysqlsh_restore_hardened.sh"
      do
        if [[ -f "$candidate" ]]; then
          fallback_hardened="$candidate"
          break
        fi
      done

      if [[ -z "$fallback_hardened" ]]; then
        echo "Error: no se encontró script de fallback mysqlsh_restore_hardened.sh"
        rm -rf "$WORK_DIR" "$DDL_FILE" "$JS_FILE" "$DDL_BOOTSTRAP_JS_FILE"
        exit 1
      fi

      rm -rf "$WORK_DIR" "$DDL_FILE" "$JS_FILE" "$DDL_BOOTSTRAP_JS_FILE"

      echo "Fallback script: $fallback_hardened"
      ALLOW_DROP_DATABASE=1 \
      DB_USER="$DB_USER" \
      DB_PASSWORD="$DB_PASSWORD" \
      DB_HOST="$DB_HOST" \
      DB_PORT="$DB_PORT" \
      MYSQL_CNF="$MYSQL_CNF" \
      "$fallback_hardened" "$DUMP_DIR" "$TARGET_DB" "" "$SKIP_VIEWS_TRIGGERS"

      echo "Fallback restore normal finalizado: $TARGET_DB"
      exit 0
    fi
  fi
else
  echo "Error: no se encontraron archivos DDL (.sql/.sql.zst/.sql.gz) para crear tablas"
  rm -rf "$WORK_DIR" "$DDL_FILE" "$JS_FILE" "$DDL_BOOTSTRAP_JS_FILE"
  exit 1
fi

mapfile -d '' data_files < <(
  find "$WORK_DIR" -type f \
    ! -name '*.sql' \
      ! -name '*.sql.zst' \
      ! -name '*.sql.gz' \
    ! -name '*.idx' \
    ! -name '*.idx.zst' \
    ! -name '*.idx.gz' \
    ! -name '@.json' \
    ! -name '@.done.json' \
    ! -name '*.json' \
    -print0 | sort -z
)

if [[ "${#data_files[@]}" -eq 0 ]]; then
  echo "Error: no se encontraron archivos de datos para importar"
  rm -rf "$WORK_DIR" "$DDL_FILE" "$JS_FILE"
  exit 1
fi

{
  printf 'try {\n'
  printf '  session.runSql("SET GLOBAL local_infile = 1")\n'
  printf '} catch (e) {\n'
  printf '  print("Aviso: no se pudo habilitar local_infile manualmente: " + e.message)\n'
  printf '}\n'
  printf 'var targetSchema = %s;\n' "$(js_quote "$TARGET_DB")"
  printf 'var threads = %s;\n' "$(js_quote "$THREADS")"
  printf 'var dataFiles = [\n'
  for data_file in "${data_files[@]}"; do
    printf '  { path: %s, table: %s, dialect: %s },\n' \
      "$(js_quote "$data_file")" \
      "$(js_quote "$(table_for_file "$data_file")")" \
      "$(js_quote "$(dialect_for_file "$data_file")")"
  done
  printf '];\n'
  printf 'for (var i = 0; i < dataFiles.length; i++) {\n'
  printf '  var currentFile = dataFiles[i];\n'
  printf '  print("Importando " + currentFile.path);\n'
  printf '  util.importTable(currentFile.path, { schema: targetSchema, table: currentFile.table, dialect: currentFile.dialect, threads: parseInt(threads, 10), showProgress: true });\n'
  printf '}\n'
} > "$JS_FILE"

echo "Ejecutando importación benchmark con util.importTable..."
"$MYSQLSH_BIN" "${DB_USER}:${DB_PASSWORD}@${DB_HOST}:${DB_PORT}" --verbose=0 --js < "$JS_FILE"

rm -rf "$WORK_DIR" "$DDL_FILE" "$JS_FILE" "$DDL_BOOTSTRAP_JS_FILE"
echo "Restore benchmark finalizado: $TARGET_DB"