content: /*
   * Copyright (c) 2023 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.destination.mongodb;
  
  import static com.mongodb.client.model.Projections.excludeId;
  
  import com.fasterxml.jackson.core.type.TypeReference;
  import com.fasterxml.jackson.databind.ObjectMapper;
  import com.mongodb.client.MongoCursor;
  import io.airbyte.cdk.db.mongodb.MongoDatabase;
  import io.airbyte.cdk.integrations.base.AirbyteMessageConsumer;
  import io.airbyte.cdk.integrations.base.FailureTrackingAirbyteMessageConsumer;
  import io.airbyte.commons.json.Jsons;
  import io.airbyte.commons.lang.Exceptions;
  import io.airbyte.protocol.models.v0.AirbyteMessage;
  import io.airbyte.protocol.models.v0.AirbyteRecordMessage;
  import io.airbyte.protocol.models.v0.AirbyteStreamNameNamespacePair;
  import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog;
  import java.nio.charset.Charset;
  import java.time.LocalDateTime;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.Map;
  import java.util.UUID;
  import java.util.function.Consumer;
  import org.apache.commons.codec.digest.DigestUtils;
  import org.bson.Document;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  public class MongodbRecordConsumer extends FailureTrackingAirbyteMessageConsumer implements AirbyteMessageConsumer {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(MongodbRecordConsumer.class);
  
    private static final String AIRBYTE_DATA = "_airbyte_data";
    private static final String AIRBYTE_DATA_HASH = "_airbyte_data_hash";
    private static final String AIRBYTE_EMITTED_AT = "_airbyte_emitted_at";
  
    private final Map<AirbyteStreamNameNamespacePair, MongodbWriteConfig> writeConfigs;
    private final MongoDatabase mongoDatabase;
    private final ConfiguredAirbyteCatalog catalog;
    private final Consumer<AirbyteMessage> outputRecordCollector;
    private final ObjectMapper objectMapper;
  
    private AirbyteMessage lastStateMessage = null;
  
    public MongodbRecordConsumer(final Map<AirbyteStreamNameNamespacePair, MongodbWriteConfig> writeConfigs,
                                 final MongoDatabase mongoDatabase,
                                 final ConfiguredAirbyteCatalog catalog,
                                 final Consumer<AirbyteMessage> outputRecordCollector) {
      this.writeConfigs = writeConfigs;
      this.mongoDatabase = mongoDatabase;
      this.catalog = catalog;
      this.outputRecordCollector = outputRecordCollector;
      this.objectMapper = new ObjectMapper();
    }
  
    @Override
    protected void startTracked() {
      // todo (cgardens) - move contents of #write into this method.
    }
  
    @Override
    protected void acceptTracked(final AirbyteMessage message) {
      if (message.getType() == AirbyteMessage.Type.STATE) {
        outputRecordCollector.accept(message);
        lastStateMessage = message;
      } else if (message.getType() == AirbyteMessage.Type.RECORD) {
        final AirbyteRecordMessage recordMessage = message.getRecord();
        final AirbyteStreamNameNamespacePair pair = AirbyteStreamNameNamespacePair.fromRecordMessage(recordMessage);
  
        if (!writeConfigs.containsKey(pair)) {
          LOGGER.error("Message contained record from a stream that was not in the catalog. catalog: {}, message: {}",
              Jsons.serialize(catalog), Jsons.serialize(recordMessage));
          throw new IllegalArgumentException(
              String.format("Message contained record from a stream that was not in the catalog. \ncatalog: %s , \nmessage: %s",
                  Jsons.serialize(catalog), Jsons.serialize(recordMessage)));
        }
        final var writeConfig = writeConfigs.get(pair);
        insertRecordToTmpCollection(writeConfig, message);
      }
    }
  
    @Override
    protected void close(final boolean hasFailed) {
      try {
        if (!hasFailed) {
          LOGGER.info("Migration finished with no explicit errors. Copying data from tmp tables to permanent");
          writeConfigs.values().forEach(mongodbWriteConfig -> Exceptions.toRuntime(() -> {
            try {
              copyTable(mongoDatabase, mongodbWriteConfig.getCollectionName(), mongodbWriteConfig.getTmpCollectionName());
            } catch (final RuntimeException e) {
              LOGGER.error("Failed to process a message for Streams numbers: {}, SyncMode: {}, CollectionName: {}, TmpCollectionName: {}",
                  catalog.getStreams().size(), mongodbWriteConfig.getSyncMode(), mongodbWriteConfig.getCollectionName(),
                  mongodbWriteConfig.getTmpCollectionName());
              LOGGER.error("Failed with exception: {}", e.getMessage());
              throw new RuntimeException(e);
            }
          }));
          outputRecordCollector.accept(lastStateMessage);
        } else {
          LOGGER.error("Had errors while migrations");
        }
      } finally {
        LOGGER.info("Removing tmp collections...");
        writeConfigs.values()
            .forEach(mongodbWriteConfig -> mongoDatabase.getCollection(mongodbWriteConfig.getTmpCollectionName()).drop());
        LOGGER.info("Finishing destination process...completed");
      }
    }
  
    /* Helpers */
  
    private void insertRecordToTmpCollection(final MongodbWriteConfig writeConfig,
                                             final AirbyteMessage message) {
      try {
        final AirbyteRecordMessage recordMessage = message.getRecord();
        final Map<String, Object> result = objectMapper.convertValue(recordMessage.getData(), new TypeReference<>() {});
        final var newDocumentDataHashCode = UUID.nameUUIDFromBytes(DigestUtils.sha256Hex(Jsons.toBytes(recordMessage.getData())).getBytes(
            Charset.defaultCharset())).toString();
        final var newDocument = new Document();
        newDocument.put(AIRBYTE_DATA, new Document(result));
        newDocument.put(AIRBYTE_DATA_HASH, newDocumentDataHashCode);
        newDocument.put(AIRBYTE_EMITTED_AT, LocalDateTime.now().toString());
  
        final var collection = writeConfig.getCollection();
  
        final var documentsHash = writeConfig.getDocumentsHash();
        if (!documentsHash.contains(newDocumentDataHashCode)) {
          collection.insertOne(newDocument);
          documentsHash.add(newDocumentDataHashCode);
        } else {
          LOGGER.info("Object with hashCode = {} already exist in table {}.", newDocumentDataHashCode, writeConfig.getCollectionName());
        }
      } catch (final RuntimeException e) {
        LOGGER.error("Got an error while writing message:" + e.getMessage());
        LOGGER.error(String.format(
            "Failed to process a message for Streams numbers: %s, SyncMode: %s, CollectionName: %s, TmpCollectionName: %s, AirbyteMessage: %s",
            catalog.getStreams().size(), writeConfig.getSyncMode(), writeConfig.getCollectionName(), writeConfig.getTmpCollectionName(), message));
        throw new RuntimeException(e);
      }
    }
  
    private static void copyTable(final MongoDatabase mongoDatabase, final String collectionName, final String tmpCollectionName) {
  
      final var tempCollection = mongoDatabase.getOrCreateNewCollection(tmpCollectionName);
      final var collection = mongoDatabase.getOrCreateNewCollection(collectionName);
      final List<Document> documents = new ArrayList<>();
      try (final MongoCursor<Document> cursor = tempCollection.find().projection(excludeId()).iterator()) {
        while (cursor.hasNext()) {
          documents.add(cursor.next());
        }
      }
      if (!documents.isEmpty()) {
        collection.insertMany(documents);
      }
    }
  
  }
  