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

package io.airbyte.integrations.destination.kinesis;

import io.airbyte.commons.json.Jsons;
import java.io.Closeable;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.core.BytesWrapper;
import software.amazon.awssdk.core.SdkBytes;
import software.amazon.awssdk.services.kinesis.KinesisClient;
import software.amazon.awssdk.services.kinesis.model.DescribeStreamResponse;
import software.amazon.awssdk.services.kinesis.model.PutRecordsRequestEntry;
import software.amazon.awssdk.services.kinesis.model.Record;
import software.amazon.awssdk.services.kinesis.model.ResourceInUseException;
import software.amazon.awssdk.services.kinesis.model.ResourceNotFoundException;
import software.amazon.awssdk.services.kinesis.model.Shard;
import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
import software.amazon.awssdk.services.kinesis.model.StreamStatus;

/**
 * KinesisStream class for performing various operations on a Kinesis stream.
 */
public class KinesisStream implements Closeable {

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

  private final KinesisClient kinesisClient;

  private final KinesisConfig kinesisConfig;

  private final int bufferSize;

  // k:v tuples of <streamName:<partitionKey,data>>
  private final List<Tuple<String, Tuple<String, String>>> buffer;

  public KinesisStream(KinesisConfig kinesisConfig) {
    this.kinesisConfig = kinesisConfig;
    this.kinesisClient = KinesisClientPool.initClient(kinesisConfig);
    this.bufferSize = kinesisConfig.getBufferSize();
    this.buffer = new ArrayList<>(bufferSize);
  }

  /**
   * Creates a stream specified via its name and with the provided shard count. The method will block
   * and retry every 2s until it verifies that the stream is active and can be written to. If the
   * stream is already created it will only wait until it is active.
   *
   * @param streamName name of the stream to be created.
   */
  public void createStream(String streamName) {
    try {
      kinesisClient.createStream(b -> b.streamName(streamName).shardCount(kinesisConfig.getShardCount()));
    } catch (ResourceInUseException e) {
      LOGGER.info("Stream with name {} has already been created", streamName);
    }
    // block/wait until stream is active
    for (;;) {
      DescribeStreamResponse describeStream = kinesisClient.describeStream(b -> b.streamName(streamName));
      if (describeStream.streamDescription().streamStatus() == StreamStatus.ACTIVE) {
        return;
      }
      try {
        Thread.sleep(2000);
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        throw KinesisUtils.buildKinesisException("Thread interrupted while waiting for stream to be active", e);
      }
    }
  }

  /**
   * Delete the stream specified via its name. The method will block and retry every 2s until it
   * verifies that the stream is deleted by receiving the appropriate exception.
   *
   * @param streamName name of the stream to be deleted.
   */
  public void deleteStream(String streamName) {
    kinesisClient.deleteStream(b -> b.streamName(streamName));
    // block/wait until stream is deleted
    for (;;) {
      try {
        kinesisClient.describeStream(b -> b.streamName(streamName));
        Thread.sleep(2000);
      } catch (ResourceNotFoundException e) {
        return;
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        throw KinesisUtils.buildKinesisException("Thread interrupted while waiting for stream to be deleted", e);
      }
    }
  }

  /**
   * Deletes all streams in the Kinesis service, waiting/blocking until all of them are deleted.
   */
  public void deleteAllStreams() {
    kinesisClient.listStreams().streamNames().forEach(this::deleteStream);
  }

  /**
   * Sends a record to the Kinesis stream specified via its name. To improve performance the records
   * are buffered until the buffer limit is reached after which they are flushed to its destination
   * stream.
   *
   * @param streamName name of the stream where the record should be sent
   * @param partitionKey to determine the destination shard
   * @param data actual data to be streamed
   * @param exceptionConsumer for handling errors related to flushing data per stream
   */
  public void putRecord(String streamName, String partitionKey, String data, Consumer<Exception> exceptionConsumer) {
    buffer.add(Tuple.of(streamName, Tuple.of(partitionKey, data)));
    if (buffer.size() == bufferSize) {
      flush(exceptionConsumer);
    }
  }

  /**
   * Iterates over all the shards for a given streams and retrieves the records which are combined and
   * deserialized to a {@link io.airbyte.integrations.destination.kinesis.KinesisRecord} objects.
   *
   * @param streamName from where to retrieve the records.
   * @return List of KinesisRecord objects retrieved from the stream.
   */
  public List<KinesisRecord> getRecords(String streamName) {
    DescribeStreamResponse describeStream;
    List<Shard> shards = new ArrayList<>();
    do {

      describeStream = kinesisClient.describeStream(b -> b.streamName(streamName));

      shards.addAll(describeStream.streamDescription().shards());

    } while (describeStream.streamDescription().hasMoreShards());

    // iterate over stream shards and retrieve records
    return shards.stream()
        .map(Shard::shardId)
        .map(sh -> kinesisClient.getShardIterator(b -> b.streamName(streamName)
            .shardIteratorType(ShardIteratorType.TRIM_HORIZON)
            .shardId(sh))
            .shardIterator())
        .flatMap(it -> kinesisClient.getRecords(b -> b.shardIterator(it)).records().stream())
        .map(Record::data)
        .map(BytesWrapper::asUtf8String)
        .map(str -> Jsons.deserialize(str, KinesisRecord.class))
        .collect(Collectors.toList());
  }

  /**
   * Flush all records previously buffered to increase throughput and performance. Records are grouped
   * by stream name and are sent for each stream separately.
   *
   * @param exceptionConsumer for handling errors related to flushing data per stream, rethrowing an
   *        exception in the consumer will stop the sync and clear the cache
   */
  public void flush(Consumer<Exception> exceptionConsumer) {
    try {
      buffer.stream()
          .collect(Collectors.groupingBy(Tuple::value1, Collectors.mapping(Tuple::value2, Collectors.toList())))
          .forEach((k, v) -> {
            var records = v.stream().map(entry -> PutRecordsRequestEntry.builder()
                // partition key used to determine stream shard.
                .partitionKey(entry.value1())
                .data(SdkBytes.fromUtf8String(entry.value2()))
                .build())
                .collect(Collectors.toList());
            try {
              kinesisClient.putRecords(b -> b.streamName(k).records(records));
            } catch (Exception e) {
              exceptionConsumer.accept(e);
            }
          });
    } finally {
      buffer.clear();
    }
  }

  /**
   * Return the kinesis client to the pool to be closed if no longer used.
   */
  @Override
  public void close() {
    KinesisClientPool.closeClient(kinesisConfig);
  }

}
