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

package io.airbyte.integrations.source.mysql_strict_encrypt;

import static io.airbyte.integrations.source.mysql.MySqlSource.SSL_PARAMETERS;
import static org.junit.jupiter.api.Assertions.assertEquals;

import com.fasterxml.jackson.databind.JsonNode;
import com.google.common.collect.ImmutableMap;
import io.airbyte.commons.json.Jsons;
import io.airbyte.commons.resources.MoreResources;
import io.airbyte.commons.string.Strings;
import io.airbyte.db.Database;
import io.airbyte.db.factory.DSLContextFactory;
import io.airbyte.db.factory.DatabaseDriver;
import io.airbyte.db.jdbc.JdbcUtils;
import io.airbyte.integrations.base.Source;
import io.airbyte.integrations.base.ssh.SshHelpers;
import io.airbyte.integrations.source.jdbc.test.JdbcSourceAcceptanceTest;
import io.airbyte.integrations.source.mysql.MySqlSource;
import io.airbyte.integrations.source.relationaldb.models.DbStreamState;
import io.airbyte.protocol.models.AirbyteCatalog;
import io.airbyte.protocol.models.AirbyteMessage;
import io.airbyte.protocol.models.AirbyteRecordMessage;
import io.airbyte.protocol.models.CatalogHelpers;
import io.airbyte.protocol.models.ConnectorSpecification;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

import io.airbyte.protocol.models.Field;
import io.airbyte.protocol.models.JsonSchemaType;
import io.airbyte.protocol.models.SyncMode;
import org.jooq.DSLContext;
import org.jooq.SQLDialect;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.MySQLContainer;

class MySqlStrictEncryptJdbcSourceAcceptanceTest extends JdbcSourceAcceptanceTest {

  protected static final String TEST_USER = "test";
  protected static final String TEST_PASSWORD = "test";
  protected static MySQLContainer<?> container;

  protected Database database;
  protected DSLContext dslContext;

  @BeforeAll
  static void init() throws SQLException {
    container = new MySQLContainer<>("mysql:8.0")
        .withUsername(TEST_USER)
        .withPassword(TEST_PASSWORD)
        .withEnv("MYSQL_ROOT_HOST", "%")
        .withEnv("MYSQL_ROOT_PASSWORD", TEST_PASSWORD);
    container.start();
    final Connection connection = DriverManager.getConnection(container.getJdbcUrl(), "root", container.getPassword());
    connection.createStatement().execute("GRANT ALL PRIVILEGES ON *.* TO '" + TEST_USER + "'@'%';\n");
  }

  @BeforeEach
  public void setup() throws Exception {
    config = Jsons.jsonNode(ImmutableMap.builder()
        .put(JdbcUtils.HOST_KEY, container.getHost())
        .put(JdbcUtils.PORT_KEY, container.getFirstMappedPort())
        .put(JdbcUtils.DATABASE_KEY, Strings.addRandomSuffix("db", "_", 10))
        .put(JdbcUtils.USERNAME_KEY, container.getUsername())
        .put(JdbcUtils.PASSWORD_KEY, container.getPassword())
        .build());

    dslContext = DSLContextFactory.create(
        config.get(JdbcUtils.USERNAME_KEY).asText(),
        config.get(JdbcUtils.PASSWORD_KEY).asText(),
        DatabaseDriver.MYSQL.getDriverClassName(),
        String.format("jdbc:mysql://%s:%s?%s",
            config.get(JdbcUtils.HOST_KEY).asText(),
            config.get(JdbcUtils.PORT_KEY).asText(),
            String.join("&", SSL_PARAMETERS)),
        SQLDialect.MYSQL);
    database = new Database(dslContext);

    database.query(ctx -> {
      ctx.fetch("CREATE DATABASE " + config.get(JdbcUtils.DATABASE_KEY).asText());
      return null;
    });

    super.setup();
  }

  @AfterEach
  void tearDownMySql() throws Exception {
    dslContext.close();
    super.tearDown();
  }

  @AfterAll
  static void cleanUp() {
    container.close();
  }

  // MySql does not support schemas in the way most dbs do. Instead we namespace by db name.
  @Override
  public boolean supportsSchemas() {
    return false;
  }

  @Override
  public MySqlSource getJdbcSource() {
    return new MySqlSource();
  }

  @Override
  public Source getSource() {
    return new MySqlStrictEncryptSource();
  }

  @Override
  public String getDriverClass() {
    return MySqlSource.DRIVER_CLASS;
  }

  @Override
  public JsonNode getConfig() {
    return Jsons.clone(config);
  }

  @Test
  void testSpec() throws Exception {
    final ConnectorSpecification actual = source.spec();
    final ConnectorSpecification expected =
        SshHelpers.injectSshIntoSpec(Jsons.deserialize(MoreResources.readResource("expected_spec.json"), ConnectorSpecification.class));
    assertEquals(expected, actual);
  }

  @Override
  protected AirbyteCatalog getCatalog(final String defaultNamespace) {
    return new AirbyteCatalog().withStreams(List.of(
            CatalogHelpers.createAirbyteStream(
                            TABLE_NAME,
                            defaultNamespace,
                            Field.of(COL_ID, JsonSchemaType.INTEGER),
                            Field.of(COL_NAME, JsonSchemaType.STRING),
                            Field.of(COL_UPDATED_AT, JsonSchemaType.STRING_DATE))
                    .withSupportedSyncModes(List.of(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL))
                    .withSourceDefinedPrimaryKey(List.of(List.of(COL_ID))),
            CatalogHelpers.createAirbyteStream(
                            TABLE_NAME_WITHOUT_PK,
                            defaultNamespace,
                            Field.of(COL_ID, JsonSchemaType.INTEGER),
                            Field.of(COL_NAME, JsonSchemaType.STRING),
                            Field.of(COL_UPDATED_AT, JsonSchemaType.STRING_DATE))
                    .withSupportedSyncModes(List.of(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL))
                    .withSourceDefinedPrimaryKey(Collections.emptyList()),
            CatalogHelpers.createAirbyteStream(
                            TABLE_NAME_COMPOSITE_PK,
                            defaultNamespace,
                            Field.of(COL_FIRST_NAME, JsonSchemaType.STRING),
                            Field.of(COL_LAST_NAME, JsonSchemaType.STRING),
                            Field.of(COL_UPDATED_AT, JsonSchemaType.STRING_DATE))
                    .withSupportedSyncModes(List.of(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL))
                    .withSourceDefinedPrimaryKey(
                            List.of(List.of(COL_FIRST_NAME), List.of(COL_LAST_NAME)))));
  }

  @Override
  protected void incrementalDateCheck() throws Exception {
    incrementalCursorCheck(
            COL_UPDATED_AT,
            "2005-10-18",
            "2006-10-19",
            List.of(getTestMessages().get(1), getTestMessages().get(2)));
  }

  @Override
  protected List<AirbyteMessage> getTestMessages() {
    return List.of(
            new AirbyteMessage().withType(AirbyteMessage.Type.RECORD)
                    .withRecord(new AirbyteRecordMessage().withStream(streamName).withNamespace(getDefaultNamespace())
                            .withData(Jsons.jsonNode(Map
                                    .of(COL_ID, ID_VALUE_1,
                                            COL_NAME, "picard",
                                            COL_UPDATED_AT, "2004-10-19")))),
            new AirbyteMessage().withType(AirbyteMessage.Type.RECORD)
                    .withRecord(new AirbyteRecordMessage().withStream(streamName).withNamespace(getDefaultNamespace())
                            .withData(Jsons.jsonNode(Map
                                    .of(COL_ID, ID_VALUE_2,
                                            COL_NAME, "crusher",
                                            COL_UPDATED_AT,
                                            "2005-10-19")))),
            new AirbyteMessage().withType(AirbyteMessage.Type.RECORD)
                    .withRecord(new AirbyteRecordMessage().withStream(streamName).withNamespace(getDefaultNamespace())
                            .withData(Jsons.jsonNode(Map
                                    .of(COL_ID, ID_VALUE_3,
                                            COL_NAME, "vash",
                                            COL_UPDATED_AT, "2006-10-19")))));
  }

  @Override
  protected List<AirbyteMessage> getExpectedAirbyteMessagesSecondSync(String namespace) {
    final List<AirbyteMessage> expectedMessages = new ArrayList<>();
    expectedMessages.add(new AirbyteMessage().withType(AirbyteMessage.Type.RECORD)
            .withRecord(new AirbyteRecordMessage().withStream(streamName).withNamespace(namespace)
                    .withData(Jsons.jsonNode(Map
                            .of(COL_ID, ID_VALUE_4,
                                    COL_NAME, "riker",
                                    COL_UPDATED_AT, "2006-10-19")))));
    expectedMessages.add(new AirbyteMessage().withType(AirbyteMessage.Type.RECORD)
            .withRecord(new AirbyteRecordMessage().withStream(streamName).withNamespace(namespace)
                    .withData(Jsons.jsonNode(Map
                            .of(COL_ID, ID_VALUE_5,
                                    COL_NAME, "data",
                                    COL_UPDATED_AT, "2006-10-19")))));
    final DbStreamState state = new DbStreamState()
            .withStreamName(streamName)
            .withStreamNamespace(namespace)
            .withCursorField(List.of(COL_ID))
            .withCursor("5");
    expectedMessages.addAll(createExpectedTestMessages(List.of(state)));
    return expectedMessages;
  }

}
