/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.sdk.extensions.sql.impl;

import static java.util.stream.Collectors.toList;
import static java.util.stream.Collectors.toMap;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;

import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.beam.sdk.extensions.sql.TableNameExtractionUtils;
import org.apache.beam.sdk.extensions.sql.meta.CustomTableResolver;
import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider;
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.jdbc.CalciteSchema;
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlNode;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;

/** Utils to wire up the custom table resolution into Calcite's planner. */
@SuppressWarnings({
  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
class TableResolutionUtils {

  /**
   * Extract table names from the FROM clauses, register them with root TableProviders that support
   * custom table schema resolution, e.g. DataCatalog.
   *
   * <p>Go over top-level schemas in the JdbcConnection, and for all top-level table providers that
   * support custom table resolution, register all the parsed table names with them.
   *
   * <p>This way when a table provider has custom name-resolution strategy it can analyze whether it
   * supports the name without using Calcite's logic. E.g. for DataCatalog we need to assemble the
   * table name back into a single string and then query the back-end, whereas Calcite would require
   * us to call the back-end for each part of the table name.
   *
   * <p>The logic is:
   *
   * <pre>
   *   - if it's a compound identifier (table name contains multiple parts):
   *       - get the first part of the identifier, assume it represents a top-level schema;
   *       - find a top-level table provider with the same name;
   *       - register the table identifier with it, if supported;
   *       - if not supported, then ignore the table identifier, everything will be resolved using
   *         existing Calcite's logic;
   *
   *   - if it's a simple identifier (contains only a table name without a schema part),
   *     or if there was no matching top-level schema:
   *       - register with the default schema, if it supports custom table resolution;
   *       - if it does not, existing Calcite logic will still work as is;
   * </pre>
   */
  static void setupCustomTableResolution(JdbcConnection connection, SqlNode parsed) {
    List<TableName> tableNames = TableNameExtractionUtils.extractTableNamesFromNode(parsed);
    String currentSchemaName = getCurrentSchemaName(connection);

    SchemaWithName defaultSchema = SchemaWithName.create(connection, currentSchemaName);

    if (defaultSchema.supportsCustomResolution()) {
      registerWithDefaultSchema(connection, tableNames, defaultSchema);
    }

    registerWithTopLevelSchemas(connection, tableNames);
  }

  /** Current (default) schema name in the JdbcConnection. */
  private static String getCurrentSchemaName(JdbcConnection connection) {
    try {
      return connection.getSchema();
    } catch (SQLException e) {
      throw new IllegalStateException(
          "Unable to get current schema name from JdbcConnection. "
              + "Assuming table names in the query are fully-qualified from the root.",
          e);
    }
  }

  /**
   * Simple identifiers have to be resolved by the default schema, as well as compoung identifiers
   * that don't have a matching top-level schema (meaning that a user didn't specify a top-level
   * schema and expected it to be inferred).
   */
  private static void registerWithDefaultSchema(
      JdbcConnection connection, List<TableName> tableNames, SchemaWithName defaultSchema) {
    Set<String> topLevelSchemas = connection.getRootSchema().getSubSchemaNames();

    List<TableName> simpleIdentifiers =
        tableNames.stream().filter(TableName::isSimple).collect(toList());

    List<TableName> withoutMatchingSchemas =
        tableNames.stream()
            .filter(name -> name.isCompound() && !topLevelSchemas.contains(name.getPrefix()))
            .collect(toList());

    List<TableName> explicitlyInDefaulSchema =
        tableNames.stream()
            .filter(name -> name.isCompound() && name.getPrefix().equals(defaultSchema.name))
            .map(TableName::removePrefix)
            .collect(toList());

    List<TableName> shouldGoIntoDefaultSchema =
        ImmutableList.<TableName>builder()
            .addAll(simpleIdentifiers)
            .addAll(withoutMatchingSchemas)
            .addAll(explicitlyInDefaulSchema)
            .build();

    defaultSchema.getCustomTableResolver().registerKnownTableNames(shouldGoIntoDefaultSchema);
  }

  /**
   * Register compound table identifiers with the matching custom resolvers that correspond to the
   * top-level schemas.
   */
  private static void registerWithTopLevelSchemas(
      JdbcConnection connection, List<TableName> tableNames) {

    Map<String, CustomTableResolver> topLevelResolvers = getCustomTopLevelResolvers(connection);

    topLevelResolvers.forEach(
        (topLevelSchemaName, resolver) ->
            resolver.registerKnownTableNames(tablesForSchema(tableNames, topLevelSchemaName)));
  }

  /** Get the custom schema resolvers for all top-level schemas that support custom resolution. */
  private static Map<String, CustomTableResolver> getCustomTopLevelResolvers(
      JdbcConnection connection) {
    return connection.getRootSchema().getSubSchemaNames().stream()
        .map(topLevelSchemaName -> SchemaWithName.create(connection, topLevelSchemaName))
        .filter(schema -> !schema.getName().equals(getCurrentSchemaName(connection)))
        .filter(SchemaWithName::supportsCustomResolution)
        .collect(toMap(SchemaWithName::getName, SchemaWithName::getCustomTableResolver));
  }

  /**
   * Get the compound identifiers that have the first component matching the given top-level schema
   * name and remove the first component.
   */
  private static List<TableName> tablesForSchema(
      List<TableName> tableNames, String topLevelSchema) {
    return tableNames.stream()
        .filter(TableName::isCompound)
        .filter(t -> t.getPrefix().equals(topLevelSchema))
        .map(TableName::removePrefix)
        .collect(toList());
  }

  /**
   * A utility class that keeps track of schema name and other properties.
   *
   * <p>Sole purpose is to reduce inline boilerplate and encapsulate stuff.
   */
  private static class SchemaWithName {
    String name;
    org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Schema schema;

    static SchemaWithName create(JdbcConnection connection, String name) {
      SchemaWithName schemaWithName = new SchemaWithName();
      schemaWithName.name = name;
      schemaWithName.schema =
          CalciteSchema.from(connection.getRootSchema().getSubSchema(name)).schema;
      return schemaWithName;
    }

    /** Whether this schema/table provider supports custom table resolution. */
    boolean supportsCustomResolution() {
      return isBeamSchema() && tableProviderSupportsCustomResolution();
    }

    /** Whether this Calcite schema is actually an instance of BeamCalciteSchema. */
    boolean isBeamSchema() {
      return schema instanceof BeamCalciteSchema;
    }

    /** Whether the table provider is an instance of CustomTableResolver. */
    boolean tableProviderSupportsCustomResolution() {
      return getTableProvider() instanceof CustomTableResolver;
    }

    /** Gets the table provider that backs the BeamCalciteSchema. */
    TableProvider getTableProvider() {
      checkState(isBeamSchema());
      return ((BeamCalciteSchema) schema).getTableProvider();
    }

    /** Schema name. */
    String getName() {
      return name;
    }

    /** Custom table resolver in the provider. */
    CustomTableResolver getCustomTableResolver() {
      checkState(supportsCustomResolution());
      return (CustomTableResolver) getTableProvider();
    }
  }
}