/*
 * SPDX-License-Identifier: Apache-2.0
 * Copyright 2018-2020 The Feast Authors
 *
 * Licensed 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
 *
 *     https://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 feast.common.it;

import io.prometheus.client.CollectorRegistry;
import java.sql.*;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.junit.jupiter.api.*;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.springframework.kafka.config.ConcurrentKafkaListenerContainerFactory;
import org.springframework.kafka.config.KafkaListenerContainerFactory;
import org.springframework.kafka.core.ConsumerFactory;
import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
import org.springframework.kafka.listener.ConcurrentMessageListenerContainer;
import org.springframework.test.annotation.DirtiesContext;
import org.springframework.test.context.ActiveProfiles;
import org.springframework.test.context.DynamicPropertyRegistry;
import org.springframework.test.context.DynamicPropertySource;
import org.testcontainers.containers.KafkaContainer;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;

/**
 * Base Integration Test class. Setups postgres and kafka containers. Configures related properties
 * and beans. Provides DB related clean up between tests.
 */
@SpringBootTest
@ActiveProfiles("it")
@Testcontainers
@DirtiesContext(classMode = DirtiesContext.ClassMode.AFTER_CLASS)
public class BaseIT {

  @Container public static PostgreSQLContainer<?> postgreSQLContainer = new PostgreSQLContainer<>();

  @Container public static KafkaContainer kafka = new KafkaContainer();

  /**
   * Configure Spring Application to use postgres and kafka rolled out in containers
   *
   * @param registry
   */
  @DynamicPropertySource
  static void properties(DynamicPropertyRegistry registry) {

    registry.add("spring.datasource.url", postgreSQLContainer::getJdbcUrl);
    registry.add("spring.datasource.username", postgreSQLContainer::getUsername);
    registry.add("spring.datasource.password", postgreSQLContainer::getPassword);
    registry.add("spring.jpa.hibernate.ddl-auto", () -> "none");

    registry.add("feast.stream.options.bootstrapServers", kafka::getBootstrapServers);
  }

  /**
   * SequentialFlow is base class that is supposed to be inherited by @Nested test classes that
   * wants to preserve context between test cases. For SequentialFlow databases is being truncated
   * only once after all tests passed.
   */
  @TestInstance(TestInstance.Lifecycle.PER_CLASS)
  public class SequentialFlow {
    @AfterAll
    public void tearDown() throws Exception {
      cleanTables();
    }
  }

  /**
   * This class must be inherited inside IT Class and annotated with {@link
   * org.springframework.boot.test.context.TestConfiguration}. It provides configuration needed to
   * communicate with Feast via Kafka
   */
  public static class BaseTestConfig {
    @Bean
    public KafkaListenerContainerFactory<ConcurrentMessageListenerContainer<String, byte[]>>
        testListenerContainerFactory(ConsumerFactory<String, byte[]> consumerFactory) {
      ConcurrentKafkaListenerContainerFactory<String, byte[]> factory =
          new ConcurrentKafkaListenerContainerFactory<>();
      factory.setConsumerFactory(consumerFactory);
      return factory;
    }

    @Bean
    public ConsumerFactory<String, byte[]> testConsumerFactory() {
      Map<String, Object> props = new HashMap<>();

      props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, kafka.getBootstrapServers());
      props.put(ConsumerConfig.GROUP_ID_CONFIG, this.getClass().getName());

      return new DefaultKafkaConsumerFactory<>(
          props, new StringDeserializer(), new ByteArrayDeserializer());
    }
  }

  /**
   * Truncates all tables in Database (between tests or flows). Retries on deadlock
   *
   * @throws SQLException
   */
  public static void cleanTables() throws SQLException {
    Connection connection =
        DriverManager.getConnection(
            postgreSQLContainer.getJdbcUrl(),
            postgreSQLContainer.getUsername(),
            postgreSQLContainer.getPassword());

    List<String> tableNames = new ArrayList<>();
    Statement statement = connection.createStatement();
    ResultSet rs =
        statement.executeQuery(
            "SELECT table_name FROM information_schema.tables WHERE table_schema='public'");
    while (rs.next()) {
      tableNames.add(rs.getString(1));
    }

    if (tableNames.isEmpty()) {
      return;
    }

    // retries are needed since truncate require exclusive lock
    // and that often leads to Deadlock
    // since SpringApp is still running in another thread
    int num_retries = 5;
    for (int i = 1; i <= num_retries; i++) {
      try {
        statement = connection.createStatement();
        statement.execute(String.format("truncate %s cascade", String.join(", ", tableNames)));
      } catch (SQLException e) {
        if (i == num_retries) {
          throw e;
        }
        continue;
      }

      break;
    }
  }

  /** Used to determine SequentialFlows */
  public Boolean isSequentialTest(TestInfo testInfo) {
    try {
      testInfo.getTestClass().get().asSubclass(SequentialFlow.class);
    } catch (ClassCastException e) {
      return false;
    }
    return true;
  }

  @AfterEach
  public void tearDown(TestInfo testInfo) throws Exception {
    CollectorRegistry.defaultRegistry.clear();

    if (!isSequentialTest(testInfo)) {
      cleanTables();
    }
  }
}
