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

package io.airbyte.integrations.destination.kinesis;

import io.airbyte.commons.json.Jsons;
import io.airbyte.integrations.base.AirbyteStreamNameNamespacePair;
import io.airbyte.integrations.base.FailureTrackingAirbyteMessageConsumer;
import io.airbyte.protocol.models.AirbyteMessage;
import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
import java.time.Instant;
import java.util.Map;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * KinesisMessageConsumer class for handling incoming Airbyte messages.
 */
public class KinesisMessageConsumer extends FailureTrackingAirbyteMessageConsumer {

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

  private final Consumer<AirbyteMessage> outputRecordCollector;

  private final KinesisStream kinesisStream;

  private final Map<AirbyteStreamNameNamespacePair, KinesisStreamConfig> kinesisStreams;

  public KinesisMessageConsumer(final ConfiguredAirbyteCatalog configuredCatalog,
                                final KinesisStream kinesisStream,
                                final Consumer<AirbyteMessage> outputRecordCollector) {
    this.outputRecordCollector = outputRecordCollector;
    this.kinesisStream = kinesisStream;
    var nameTransformer = new KinesisNameTransformer();
    this.kinesisStreams = configuredCatalog.getStreams().stream()
        .collect(Collectors.toUnmodifiableMap(
            AirbyteStreamNameNamespacePair::fromConfiguredAirbyteSteam,
            k -> new KinesisStreamConfig(
                nameTransformer.streamName(k.getStream().getNamespace(), k.getStream().getName()),
                k.getDestinationSyncMode())));
  }

  /**
   * Start tracking the incoming Airbyte streams by creating the needed Kinesis streams.
   */
  @Override
  protected void startTracked() {
    kinesisStreams.forEach((k, v) -> kinesisStream.createStream(v.getStreamName()));
  }

  /**
   * Handle an incoming Airbyte message by serializing it to the appropriate Kinesis structure and
   * sending it to the stream.
   *
   * @param message received from the Airbyte source.
   */
  @Override
  protected void acceptTracked(final AirbyteMessage message) {
    if (message.getType() == AirbyteMessage.Type.RECORD) {
      var messageRecord = message.getRecord();

      var streamConfig =
          kinesisStreams.get(AirbyteStreamNameNamespacePair.fromRecordMessage(messageRecord));

      if (streamConfig == null) {
        throw new IllegalArgumentException("Unrecognized destination stream");
      }

      var partitionKey = KinesisUtils.buildPartitionKey();

      var data = Jsons.jsonNode(Map.of(
          KinesisRecord.COLUMN_NAME_AB_ID, partitionKey,
          KinesisRecord.COLUMN_NAME_DATA, Jsons.serialize(messageRecord.getData()),
          KinesisRecord.COLUMN_NAME_EMITTED_AT, Instant.now()));

      var streamName = streamConfig.getStreamName();
      kinesisStream.putRecord(streamName, partitionKey, Jsons.serialize(data), e -> {
        LOGGER.error("Error while streaming data to Kinesis", e);
        // throw exception and end sync?
      });
    } else if (message.getType() == AirbyteMessage.Type.STATE) {
      outputRecordCollector.accept(message);
    } else {
      LOGGER.warn("Unsupported airbyte message type: {}", message.getType());
    }
  }

  /**
   * Flush the Kinesis stream if there are any remaining messages to be sent and close the client as a
   * terminal operation.
   *
   * @param hasFailed flag for indicating if the operation has failed.
   */
  @Override
  protected void close(final boolean hasFailed) {
    try {
      if (!hasFailed) {
        kinesisStream.flush(e -> {
          LOGGER.error("Error while streaming data to Kinesis", e);
        });
      }
    } finally {
      kinesisStream.close();
    }
  }

}
