/*
 * 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.io.snowflake.services;

import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import javax.sql.DataSource;
import org.apache.beam.sdk.io.snowflake.data.SnowflakeTableSchema;
import org.apache.beam.sdk.io.snowflake.enums.CreateDisposition;
import org.apache.beam.sdk.io.snowflake.enums.WriteDisposition;
import org.apache.beam.sdk.transforms.SerializableFunction;

/** Implemenation of {@link SnowflakeServices.BatchService} used in production. */
@SuppressWarnings({
  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
public class SnowflakeBatchServiceImpl implements SnowflakeServices.BatchService {
  private static final String SNOWFLAKE_GCS_PREFIX = "gcs://";
  private static final String GCS_PREFIX = "gs://";

  /** Writing data to Snowflake in batch mode. */
  @Override
  public void write(SnowflakeBatchServiceConfig config) throws Exception {
    copyToTable(config);
  }

  /** Reading data from Snowflake tables in batch processing. */
  @Override
  public String read(SnowflakeBatchServiceConfig config) throws Exception {
    return copyIntoStage(config);
  }

  /**
   * Copies data from specified table to stage (bucket and directory). Uses Snowflake's <a
   * href="https://docs.snowflake.com/en/sql-reference/sql/copy-into-location.html">COPY method</a>.
   * All the details needed for COPY are inside passed configuration.
   *
   * @param config object with configuration to perform COPY query.
   * @return destination where files were copied into
   * @throws SQLException in case COPY query failed.
   */
  private String copyIntoStage(SnowflakeBatchServiceConfig config) throws SQLException {
    SerializableFunction<Void, DataSource> dataSourceProviderFn = config.getDataSourceProviderFn();
    String database = config.getDatabase();
    String schema = config.getSchema();
    String table = config.getTable();
    String query = config.getQuery();
    String storageIntegrationName = config.getStorageIntegrationName();
    String stagingBucketDir = config.getStagingBucketDir();

    String source;
    if (query != null) {
      // Query must be surrounded with brackets
      source = String.format("(%s)", query);
    } else {
      source = getTablePath(database, schema, table);
    }

    String copyQuery =
        String.format(
            "COPY INTO '%s' FROM %s STORAGE_INTEGRATION=%s FILE_FORMAT=(TYPE=CSV COMPRESSION=GZIP FIELD_OPTIONALLY_ENCLOSED_BY='%s' ESCAPE='\\\\');",
            getProperBucketDir(stagingBucketDir),
            source,
            storageIntegrationName,
            getASCIICharRepresentation(config.getQuotationMark()));

    runStatement(copyQuery, getConnection(dataSourceProviderFn), null);

    return stagingBucketDir.concat("*");
  }

  private String getASCIICharRepresentation(String input) {
    return String.format("0x%x", new BigInteger(1, input.getBytes(StandardCharsets.UTF_8)));
  }

  /**
   * Copies staged data from bucket directory to table. Uses Snowflake's <a
   * href="https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html">COPY method</a>.
   * All the details needed for COPY are inside passed configuration.
   *
   * @param config object with configuration to perform COPY query.
   * @throws SQLException
   */
  private void copyToTable(SnowflakeBatchServiceConfig config) throws SQLException {

    SerializableFunction<Void, DataSource> dataSourceProviderFn = config.getDataSourceProviderFn();
    List<String> filesList = config.getFilesList();
    String database = config.getDatabase();
    String schema = config.getSchema();
    String table = config.getTable();
    String query = config.getQuery();
    SnowflakeTableSchema tableSchema = config.getTableSchema();
    CreateDisposition createDisposition = config.getCreateDisposition();
    WriteDisposition writeDisposition = config.getWriteDisposition();
    String storageIntegrationName = config.getStorageIntegrationName();
    String stagingBucketDir = config.getStagingBucketDir();

    String source;
    if (query != null) {
      // Query must be surrounded with brackets
      source = String.format("(%s)", query);
    } else {
      source = String.format("'%s'", stagingBucketDir);
    }

    filesList = filesList.stream().map(e -> String.format("'%s'", e)).collect(Collectors.toList());
    String files = String.join(", ", filesList);
    files = files.replaceAll(stagingBucketDir, "");
    DataSource dataSource = dataSourceProviderFn.apply(null);

    prepareTableAccordingCreateDisposition(
        dataSource, database, schema, table, tableSchema, createDisposition);
    prepareTableAccordingWriteDisposition(dataSource, table, writeDisposition);

    if (!storageIntegrationName.isEmpty()) {
      query =
          String.format(
              "COPY INTO %s FROM %s FILES=(%s) FILE_FORMAT=(TYPE=CSV FIELD_OPTIONALLY_ENCLOSED_BY='%s' ESCAPE='\\\\' COMPRESSION=GZIP) STORAGE_INTEGRATION=%s;",
              getTablePath(database, schema, table),
              getProperBucketDir(source),
              files,
              getASCIICharRepresentation(config.getQuotationMark()),
              storageIntegrationName);
    } else {
      query =
          String.format(
              "COPY INTO %s FROM %s FILES=(%s) FILE_FORMAT=(TYPE=CSV FIELD_OPTIONALLY_ENCLOSED_BY='%s' ESCAPE='\\\\' COMPRESSION=GZIP);",
              table, source, files, getASCIICharRepresentation(config.getQuotationMark()));
    }

    runStatement(query, dataSource.getConnection(), null);
  }

  private void truncateTable(DataSource dataSource, String tablePath) throws SQLException {
    String query = String.format("TRUNCATE %s;", tablePath);
    runConnectionWithStatement(dataSource, query, null);
  }

  private static void checkIfTableIsEmpty(DataSource dataSource, String tablePath)
      throws SQLException {
    String selectQuery = String.format("SELECT count(*) FROM %s LIMIT 1;", tablePath);
    runConnectionWithStatement(
        dataSource,
        selectQuery,
        resultSet -> {
          assert resultSet != null;
          checkIfTableIsEmpty(resultSet);
        });
  }

  private static void checkIfTableIsEmpty(ResultSet resultSet) {
    int columnId = 1;
    try {
      if (!resultSet.next() || !checkIfTableIsEmpty(resultSet, columnId)) {
        throw new RuntimeException("Table is not empty. Aborting COPY with disposition EMPTY");
      }
    } catch (SQLException e) {
      throw new RuntimeException("Unable run pipeline with EMPTY disposition.", e);
    }
  }

  private static boolean checkIfTableIsEmpty(ResultSet resultSet, int columnId)
      throws SQLException {
    int rowCount = resultSet.getInt(columnId);
    if (rowCount >= 1) {
      return false;
    }
    return true;
  }

  private void prepareTableAccordingCreateDisposition(
      DataSource dataSource,
      String database,
      String schema,
      String table,
      SnowflakeTableSchema tableSchema,
      CreateDisposition createDisposition)
      throws SQLException {
    switch (createDisposition) {
      case CREATE_NEVER:
        break;
      case CREATE_IF_NEEDED:
        createTableIfNotExists(dataSource, database, schema, table, tableSchema);
        break;
    }
  }

  private void prepareTableAccordingWriteDisposition(
      DataSource dataSource, String table, WriteDisposition writeDisposition) throws SQLException {
    switch (writeDisposition) {
      case TRUNCATE:
        truncateTable(dataSource, table);
        break;
      case EMPTY:
        checkIfTableIsEmpty(dataSource, table);
        break;
      case APPEND:
      default:
        break;
    }
  }

  private void createTableIfNotExists(
      DataSource dataSource,
      String database,
      String schema,
      String table,
      SnowflakeTableSchema tableSchema)
      throws SQLException {
    String query =
        String.format(
            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_catalog = '%s' AND table_schema = '%s' AND table_name = '%s');",
            database.toUpperCase(), schema.toUpperCase(), table.toUpperCase());

    runConnectionWithStatement(
        dataSource,
        query,
        resultSet -> {
          assert resultSet != null;
          if (!checkResultIfTableExists(resultSet)) {
            try {
              createTable(dataSource, table, tableSchema);
            } catch (SQLException e) {
              throw new RuntimeException("Unable to create table.", e);
            }
          }
        });
  }

  private static boolean checkResultIfTableExists(ResultSet resultSet) {
    try {
      if (resultSet.next()) {
        return checkIfResultIsTrue(resultSet);
      } else {
        throw new RuntimeException("Unable run pipeline with CREATE IF NEEDED - no response.");
      }
    } catch (SQLException e) {
      throw new RuntimeException("Unable run pipeline with CREATE IF NEEDED disposition.", e);
    }
  }

  private void createTable(DataSource dataSource, String table, SnowflakeTableSchema tableSchema)
      throws SQLException {
    checkArgument(
        tableSchema != null,
        "The CREATE_IF_NEEDED disposition requires schema if table doesn't exists");
    String query = String.format("CREATE TABLE %s (%s);", table, tableSchema.sql());
    runConnectionWithStatement(dataSource, query, null);
  }

  private static boolean checkIfResultIsTrue(ResultSet resultSet) throws SQLException {
    int columnId = 1;
    return resultSet.getBoolean(columnId);
  }

  private static void runConnectionWithStatement(
      DataSource dataSource, String query, Consumer<ResultSet> resultSetMethod)
      throws SQLException {
    Connection connection = dataSource.getConnection();
    runStatement(query, connection, resultSetMethod);
    connection.close();
  }

  private static void runStatement(
      String query, Connection connection, Consumer<ResultSet> resultSetMethod)
      throws SQLException {
    PreparedStatement statement = connection.prepareStatement(query);
    try {
      if (resultSetMethod != null) {
        ResultSet resultSet = statement.executeQuery();
        resultSetMethod.accept(resultSet);
      } else {
        statement.execute();
      }
    } finally {
      statement.close();
      connection.close();
    }
  }

  private Connection getConnection(SerializableFunction<Void, DataSource> dataSourceProviderFn)
      throws SQLException {
    DataSource dataSource = dataSourceProviderFn.apply(null);
    return dataSource.getConnection();
  }

  // Snowflake is expecting "gcs://" prefix for GCS and Beam "gs://"
  private String getProperBucketDir(String bucketDir) {
    if (bucketDir.contains(GCS_PREFIX)) {
      return bucketDir.replace(GCS_PREFIX, SNOWFLAKE_GCS_PREFIX);
    }
    return bucketDir;
  }

  private String getTablePath(String database, String schema, String table) {
    return String.format("%s.%s.%s", database, schema, table);
  }
}
