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

package io.airbyte.integrations.destination.redshift.copiers;

import static java.util.Comparator.comparing;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;

import com.amazonaws.services.s3.AmazonS3Client;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.Lists;
import io.airbyte.db.jdbc.JdbcDatabase;
import io.airbyte.integrations.destination.ExtendedNameTransformer;
import io.airbyte.integrations.destination.jdbc.SqlOperations;
import io.airbyte.integrations.destination.jdbc.copy.s3.S3CopyConfig;
import io.airbyte.integrations.destination.s3.S3DestinationConfig;
import io.airbyte.protocol.models.AirbyteStream;
import io.airbyte.protocol.models.ConfiguredAirbyteStream;
import io.airbyte.protocol.models.DestinationSyncMode;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.time.Instant;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicReference;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

class RedshiftStreamCopierTest {

  private static final Logger LOGGER = LoggerFactory.getLogger(RedshiftStreamCopierTest.class);

  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

  // The full path would be something like
  // "fake-namespace/fake_stream/2021_12_09_1639077474000_e549e712-b89c-4272-9496-9690ba7f973e.csv"
  // The namespace and stream have their hyphens replaced by underscores. Not super clear that that's
  // actually required.
  // 2021_12_09_1639077474000 is generated from the timestamp. It's followed by a random UUID, in case
  // we need to create multiple files.
  private static final String EXPECTED_OBJECT_BEGINNING = "fake-bucketPath/fake_namespace/fake_stream/2021_12_09_1639077474000_";
  private static final String EXPECTED_OBJECT_ENDING = ".csv";

  // equivalent to Thu, 09 Dec 2021 19:17:54 GMT
  private static final Timestamp UPLOAD_TIME = Timestamp.from(Instant.ofEpochMilli(1639077474000L));

  private AmazonS3Client s3Client;
  private JdbcDatabase db;
  private SqlOperations sqlOperations;
  private RedshiftStreamCopier copier;

  @BeforeEach
  public void setup() {
    s3Client = mock(AmazonS3Client.class, RETURNS_DEEP_STUBS);
    db = mock(JdbcDatabase.class);
    sqlOperations = mock(SqlOperations.class);

    final S3DestinationConfig s3Config = S3DestinationConfig.create(
        "fake-bucket",
        "fake-bucketPath",
        "fake-region")
        .withEndpoint("fake-endpoint")
        .withAccessKeyCredential("fake-access-key-id", "fake-secret-access-key")
        .get();

    copier = new RedshiftStreamCopier(
        // In reality, this is normally a UUID - see CopyConsumerFactory#createWriteConfigs
        "fake-staging-folder",
        "fake-schema",
        s3Client,
        db,
        new S3CopyConfig(true, s3Config),
        new ExtendedNameTransformer(),
        sqlOperations,
        UPLOAD_TIME,
        new ConfiguredAirbyteStream()
            .withDestinationSyncMode(DestinationSyncMode.APPEND)
            .withStream(new AirbyteStream()
                .withName("fake-stream")
                .withNamespace("fake-namespace")));
  }

  @Test
  public void copiesCorrectFilesToTable() throws SQLException {
    // Generate two files
    final String file1 = copier.prepareStagingFile();
    for (int i = 0; i < RedshiftStreamCopier.MAX_PARTS_PER_FILE - 1; i++) {
      copier.prepareStagingFile();
    }
    final String file2 = copier.prepareStagingFile();
    final List<String> expectedFiles = List.of(file1, file2).stream().sorted().toList();

    copier.copyStagingFileToTemporaryTable();

    final AtomicReference<String> manifestUuid = new AtomicReference<>();
    verify(s3Client).putObject(
        eq("fake-bucket"),
        argThat(path -> {
          final boolean startsCorrectly = path.startsWith("fake-bucketPath/fake-staging-folder/fake-schema/");
          final boolean endsCorrectly = path.endsWith(".manifest");
          // Make sure that we have a valid UUID
          manifestUuid.set(path.replaceFirst("^fake-bucketPath/fake-staging-folder/fake-schema/", "").replaceFirst(".manifest$", ""));
          UUID.fromString(manifestUuid.get());

          return startsCorrectly && endsCorrectly;
        }),
        (String) argThat(manifestStr -> {
          try {
            final JsonNode manifest = OBJECT_MAPPER.readTree((String) manifestStr);
            final List<JsonNode> entries = Lists.newArrayList(manifest.get("entries").elements()).stream()
                .sorted(comparing(entry -> entry.get("url").asText())).toList();

            boolean entriesAreCorrect = true;
            for (int i = 0; i < 2; i++) {
              final String expectedFilename = expectedFiles.get(i);
              final JsonNode manifestEntry = entries.get(i);
              entriesAreCorrect &= isManifestEntryCorrect(manifestEntry, expectedFilename);
              if (!entriesAreCorrect) {
                LOGGER.error("Invalid entry: {}", manifestEntry);
              }
            }

            return entriesAreCorrect && entries.size() == 2;
          } catch (final JsonProcessingException e) {
            throw new RuntimeException(e);
          }
        }));

    verify(db).execute(String.format(
        """
        COPY fake-schema.%s FROM 's3://fake-bucket/fake-bucketPath/fake-staging-folder/fake-schema/%s.manifest'
        CREDENTIALS 'aws_access_key_id=fake-access-key-id;aws_secret_access_key=fake-secret-access-key'
        CSV REGION 'fake-region' TIMEFORMAT 'auto'
        STATUPDATE OFF
        MANIFEST;""",
        copier.getTmpTableName(),
        manifestUuid.get()));
  }

  private static boolean isManifestEntryCorrect(final JsonNode entry, final String expectedFilename) {
    final String url = entry.get("url").asText();
    final boolean mandatory = entry.get("mandatory").asBoolean();

    return ("s3://fake-bucket/" + expectedFilename).equals(url) && mandatory;
  }

}
