/*
 * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
 */

package io.airbyte.db.jdbc;

import com.google.errorprone.annotations.MustBeClosed;
import io.airbyte.commons.functional.CheckedFunction;
import io.airbyte.db.JdbcCompatibleSourceOperations;
import io.airbyte.db.jdbc.streaming.JdbcStreamingQueryConfig;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Spliterator;
import java.util.Spliterators;
import java.util.function.Consumer;
import java.util.function.Supplier;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;
import javax.sql.DataSource;

/**
 * This database allows a developer to specify a {@link JdbcStreamingQueryConfig}. This allows the
 * developer to specify the correct configuration in order for a {@link PreparedStatement} to
 * execute as in a streaming / chunked manner.
 */
public class StreamingJdbcDatabase extends DefaultJdbcDatabase {

  private final Supplier<JdbcStreamingQueryConfig> streamingQueryConfigProvider;

  public StreamingJdbcDatabase(final DataSource dataSource,
                               final JdbcCompatibleSourceOperations<?> sourceOperations,
                               final Supplier<JdbcStreamingQueryConfig> streamingQueryConfigProvider) {
    super(dataSource, sourceOperations);
    this.streamingQueryConfigProvider = streamingQueryConfigProvider;
  }

  /**
   * Assuming that the {@link JdbcStreamingQueryConfig} is configured correctly for the JDBC driver
   * being used, this method will return data in streaming / chunked fashion. Review the provided
   * {@link JdbcStreamingQueryConfig} to understand the size of these chunks. If the entire stream is
   * consumed the database connection will be closed automatically and the caller need not call close
   * on the returned stream. This query (and the first chunk) are fetched immediately. Subsequent
   * chunks will not be pulled until the first chunk is consumed.
   *
   * @param statementCreator create a {@link PreparedStatement} from a {@link Connection}.
   * @param recordTransform transform each record of that result set into the desired type. do NOT
   *        just pass the {@link ResultSet} through. it is a stateful object will not be accessible if
   *        returned from recordTransform.
   * @param <T> type that each record will be mapped to.
   * @return Result of the query mapped to a stream. This stream must be closed!
   * @throws SQLException SQL related exceptions.
   */
  @Override
  @MustBeClosed
  public <T> Stream<T> unsafeQuery(final CheckedFunction<Connection, PreparedStatement, SQLException> statementCreator,
                                   final CheckedFunction<ResultSet, T, SQLException> recordTransform)
      throws SQLException {
    try {
      final Connection connection = dataSource.getConnection();
      final PreparedStatement statement = statementCreator.apply(connection);
      final JdbcStreamingQueryConfig streamingConfig = streamingQueryConfigProvider.get();
      streamingConfig.initialize(connection, statement);
      return toUnsafeStream(statement.executeQuery(), recordTransform, streamingConfig)
          .onClose(() -> {
            try {
              connection.setAutoCommit(true);
              connection.close();
            } catch (final SQLException e) {
              throw new RuntimeException(e);
            }
          });
    } catch (final SQLException e) {
      throw new RuntimeException(e);
    }
  }

  /**
   * This method differs from {@link DefaultJdbcDatabase#toUnsafeStream} in that it takes a streaming
   * config that adjusts the fetch size dynamically according to sampled row size.
   */
  protected static <T> Stream<T> toUnsafeStream(final ResultSet resultSet,
                                                final CheckedFunction<ResultSet, T, SQLException> mapper,
                                                final JdbcStreamingQueryConfig streamingConfig) {
    return StreamSupport.stream(new Spliterators.AbstractSpliterator<>(Long.MAX_VALUE, Spliterator.ORDERED) {

      @Override
      public boolean tryAdvance(final Consumer<? super T> action) {
        try {
          if (!resultSet.next()) {
            resultSet.close();
            return false;
          }
          final T dataRow = mapper.apply(resultSet);
          streamingConfig.accept(resultSet, dataRow);
          action.accept(dataRow);
          return true;
        } catch (final SQLException e) {
          throw new RuntimeException(e);
        }
      }

    }, false);
  }

}
