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

package io.airbyte.integrations.standardtest.source.performancetest;

import static org.junit.jupiter.api.Assertions.fail;

import com.fasterxml.jackson.databind.JsonNode;
import com.google.common.collect.Lists;
import io.airbyte.integrations.standardtest.source.TestDestinationEnv;
import io.airbyte.protocol.models.AirbyteStream;
import io.airbyte.protocol.models.CatalogHelpers;
import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
import io.airbyte.protocol.models.ConfiguredAirbyteStream;
import io.airbyte.protocol.models.DestinationSyncMode;
import io.airbyte.protocol.models.Field;
import io.airbyte.protocol.models.JsonSchemaType;
import io.airbyte.protocol.models.SyncMode;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * This abstract class contains common methods for Performance tests.
 */
public abstract class AbstractSourcePerformanceTest extends AbstractSourceBasePerformanceTest {

  protected static final Logger c = LoggerFactory.getLogger(AbstractSourcePerformanceTest.class);
  private static final String ID_COLUMN_NAME = "id";
  protected JsonNode config;

  /**
   * Setup the test database. All tables and data described in the registered tests will be put there.
   *
   * @throws Exception - might throw any exception during initialization.
   */
  protected abstract void setupDatabase(String dbName) throws Exception;

  @Override
  protected JsonNode getConfig() {
    return config;
  }

  @Override
  protected void tearDown(final TestDestinationEnv testEnv) {}

  /**
   * This is a data provider for performance tests, Each argument's group would be ran as a separate
   * test. Set the "testArgs" in test class of your DB in @BeforeTest method.
   *
   * 1st arg - a name of DB that will be used in jdbc connection string. 2nd arg - a schemaName that
   * will be ised as a NameSpace in Configured Airbyte Catalog. 3rd arg - a number of expected records
   * retrieved in each stream. 4th arg - a number of columns in each stream\table that will be used
   * for Airbyte Cataloq configuration 5th arg - a number of streams to read in configured airbyte
   * Catalog. Each stream\table in DB should be names like "test_0", "test_1",..., test_n.
   *
   * Example: Stream.of( Arguments.of("test1000tables240columns200recordsDb", "dbo", 200, 240, 1000),
   * Arguments.of("test5000tables240columns200recordsDb", "dbo", 200, 240, 1000),
   * Arguments.of("newregular25tables50000records", "dbo", 50052, 8, 25),
   * Arguments.of("newsmall1000tableswith10000rows", "dbo", 10011, 8, 1000) );
   */
  protected abstract Stream<Arguments> provideParameters();

  @ParameterizedTest
  @MethodSource("provideParameters")
  public void testPerformance(final String dbName,
                              final String schemaName,
                              final int numberOfDummyRecords,
                              final int numberOfColumns,
                              final int numberOfStreams)
      throws Exception {

    setupDatabase(dbName);

    final ConfiguredAirbyteCatalog catalog = getConfiguredCatalog(schemaName, numberOfStreams,
        numberOfColumns);
    final Map<String, Integer> mapOfExpectedRecordsCount = prepareMapWithExpectedRecords(
        numberOfStreams, numberOfDummyRecords);
    final Map<String, Integer> checkStatusMap = runReadVerifyNumberOfReceivedMsgs(catalog, null,
        mapOfExpectedRecordsCount);
    validateNumberOfReceivedMsgs(checkStatusMap);

  }

  /**
   * The column name will be used for a PK column in the test tables. Override it if default name is
   * not valid for your source.
   *
   * @return Id column name
   */
  protected String getIdColumnName() {
    return ID_COLUMN_NAME;
  }

  protected void validateNumberOfReceivedMsgs(final Map<String, Integer> checkStatusMap) {
    // Iterate through all streams map and check for streams where
    final Map<String, Integer> failedStreamsMap = checkStatusMap.entrySet().stream()
        .filter(el -> el.getValue() != 0).collect(Collectors.toMap(Entry::getKey, Entry::getValue));

    if (!failedStreamsMap.isEmpty()) {
      fail("Non all messages were delivered. " + failedStreamsMap.toString());
    }
    c.info("Finished all checks, no issues found for {} of streams", checkStatusMap.size());
  }

  protected Map<String, Integer> prepareMapWithExpectedRecords(final int streamNumber,
                                                               final int expectedRecordsNumberInEachStream) {
    final Map<String, Integer> resultMap = new HashMap<>(); // streamName&expected records in stream

    for (int currentStream = 0; currentStream < streamNumber; currentStream++) {
      final String streamName = String.format(getTestStreamNameTemplate(), currentStream);
      resultMap.put(streamName, expectedRecordsNumberInEachStream);
    }
    return resultMap;
  }

  /**
   * Configures streams for all registered data type tests.
   *
   * @return configured catalog
   */
  protected ConfiguredAirbyteCatalog getConfiguredCatalog(final String nameSpace,
                                                          final int numberOfStreams,
                                                          final int numberOfColumns) {
    final List<ConfiguredAirbyteStream> streams = new ArrayList<>();

    for (int currentStream = 0; currentStream < numberOfStreams; currentStream++) {

      // CREATE TABLE test.test_1_int(id INTEGER PRIMARY KEY, test_column int)
      final List<Field> fields = new ArrayList<>();

      fields.add(Field.of(getIdColumnName(), JsonSchemaType.NUMBER));
      for (int currentColumnNumber = 0;
          currentColumnNumber < numberOfColumns;
          currentColumnNumber++) {
        fields.add(Field.of(getTestColumnName() + currentColumnNumber, JsonSchemaType.STRING));
      }

      final AirbyteStream airbyteStream = CatalogHelpers
          .createAirbyteStream(String.format(getTestStreamNameTemplate(), currentStream),
              nameSpace, fields)
          .withSourceDefinedCursor(true)
          .withSourceDefinedPrimaryKey(List.of(List.of(getIdColumnName())))
          .withSupportedSyncModes(
              Lists.newArrayList(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL));

      final ConfiguredAirbyteStream configuredAirbyteStream = new ConfiguredAirbyteStream()
          .withSyncMode(SyncMode.INCREMENTAL)
          .withCursorField(Lists.newArrayList(getIdColumnName()))
          .withDestinationSyncMode(DestinationSyncMode.APPEND)
          .withStream(airbyteStream);

      streams.add(configuredAirbyteStream);

    }

    return new ConfiguredAirbyteCatalog().withStreams(streams);
  }

}
