/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * license agreements; and to You under the Apache License, version 2.0:
 *
 *   https://www.apache.org/licenses/LICENSE-2.0
 *
 * This file is part of the Apache Pekko project, which was derived from Akka.
 */

/*
 * Copyright (C) 2014 - 2016 Softwaremill <https://softwaremill.com>
 * Copyright (C) 2016 - 2020 Lightbend Inc. <https://www.lightbend.com>
 */

package org.apache.pekko.kafka.javadsl

import java.util.concurrent.{ CompletionStage, Executor }

import org.apache.pekko
import pekko.actor.ActorRef
import pekko.annotation.ApiMayChange
import pekko.dispatch.ExecutionContexts
import pekko.japi.Pair
import pekko.kafka.ConsumerMessage.{ CommittableMessage, CommittableOffset }
import pekko.kafka._
import pekko.kafka.internal.{ ConsumerControlAsJava, SourceWithOffsetContext }
import pekko.stream.javadsl.{ Source, SourceWithContext }
import pekko.{ Done, NotUsed }
import pekko.util.FutureConverters._
import pekko.util.ccompat.JavaConverters._
import org.apache.kafka.clients.consumer.ConsumerRecord
import org.apache.kafka.common.{ Metric, MetricName, TopicPartition }

import scala.concurrent.duration.FiniteDuration

/**
 * Apache Pekko Stream connector for subscribing to Kafka topics.
 */
object Consumer {

  /**
   * Materialized value of the consumer `Source`.
   */
  trait Control {

    /**
     * Stop producing messages from the `Source`. This does not stop underlying kafka consumer
     * and does not unsubscribe from any topics/partitions.
     *
     * Call [[#shutdown]] to close consumer
     */
    def stop(): CompletionStage[Done]

    /**
     * Shutdown the consumer `Source`. It will wait for outstanding offset
     * commit requests before shutting down.
     */
    def shutdown(): CompletionStage[Done]

    /**
     * Stop producing messages from the `Source`, wait for stream completion
     * and shut down the consumer `Source` so that all consumed messages
     * reach the end of the stream.
     * Failures in stream completion will be propagated, the source will be shut down anyway.
     */
    def drainAndShutdown[T](streamCompletion: CompletionStage[T], ec: Executor): CompletionStage[T]

    /**
     * Shutdown status. The `CompletionStage` will be completed when the stage has been shut down
     * and the underlying `KafkaConsumer` has been closed. Shutdown can be triggered
     * from downstream cancellation, errors, or [[#shutdown]].
     */
    def isShutdown: CompletionStage[Done]

    /**
     * Exposes underlying consumer or producer metrics (as reported by underlying Kafka client library)
     */
    def getMetrics: CompletionStage[java.util.Map[MetricName, Metric]]

  }

  /**
   * Combine control and a stream completion signal materialized values into
   * one, so that the stream can be stopped in a controlled way without losing
   * commits.
   */
  final class DrainingControl[T] private[javadsl] (control: Control, val streamCompletion: CompletionStage[T])
      extends Control {

    override def stop(): CompletionStage[Done] = control.stop()

    override def shutdown(): CompletionStage[Done] = control.shutdown()

    override def drainAndShutdown[S](streamCompletion: CompletionStage[S], ec: Executor): CompletionStage[S] =
      control.drainAndShutdown(streamCompletion, ec)

    /**
     * Stop producing messages from the `Source`, wait for stream completion
     * and shut down the consumer `Source`. It will wait for outstanding offset
     * commit requests to finish before shutting down.
     */
    def drainAndShutdown(ec: Executor): CompletionStage[T] =
      control.drainAndShutdown(streamCompletion, ec)

    override def isShutdown: CompletionStage[Done] = control.isShutdown

    override def getMetrics: CompletionStage[java.util.Map[MetricName, Metric]] = control.getMetrics
  }

  /**
   * Combine the consumer control and a stream completion signal materialized values into
   * one, so that the stream can be stopped in a controlled way without losing
   * commits.
   *
   * For use in `mapMaterializedValue`.
   */
  def createDrainingControl[T](pair: Pair[Control, CompletionStage[T]]) =
    new DrainingControl[T](pair.first, pair.second)

  /**
   * Combine the consumer control and a stream completion signal materialized values into
   * one, so that the stream can be stopped in a controlled way without losing
   * commits.
   *
   * For use in the `toMat` combination of materialized values.
   */
  def createDrainingControl[T](c: Control, mat: CompletionStage[T]): DrainingControl[T] = new DrainingControl[T](c, mat)

  /**
   * An implementation of Control to be used as an empty value, all methods return
   * a failed `CompletionStage`.
   */
  def createNoopControl(): Control = new ConsumerControlAsJava(scaladsl.Consumer.NoopControl)

  /**
   * The `plainSource` emits `ConsumerRecord` elements (as received from the underlying `KafkaConsumer`).
   * It has no support for committing offsets to Kafka. It can be used when the offset is stored externally
   * or with auto-commit (note that auto-commit is by default disabled).
   *
   * The consumer application doesn't need to use Kafka's built-in offset storage and can store offsets in a store of its own
   * choosing. The primary use case for this is allowing the application to store both the offset and the results of the
   * consumption in the same system in a way that both the results and offsets are stored atomically. This is not always
   * possible, but when it is, it will make the consumption fully atomic and give "exactly once" semantics that are
   * stronger than the "at-least once" semantics you get with Kafka's offset commit functionality.
   */
  def plainSource[K, V](settings: ConsumerSettings[K, V],
      subscription: Subscription): Source[ConsumerRecord[K, V], Control] =
    scaladsl.Consumer
      .plainSource(settings, subscription)
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * The `committableSource` makes it possible to commit offset positions to Kafka.
   * This is useful when "at-least once delivery" is desired, as each message will likely be
   * delivered one time but in failure cases could be duplicated.
   *
   * If you commit the offset before processing the message you get "at-most once delivery" semantics,
   * and for that there is a [[#atMostOnceSource]].
   *
   * Compared to auto-commit, this gives exact control over when a message is considered consumed.
   *
   * If you need to store offsets in anything other than Kafka, [[#plainSource]] should be used
   * instead of this API.
   */
  def committableSource[K, V](settings: ConsumerSettings[K, V],
      subscription: Subscription): Source[CommittableMessage[K, V], Control] =
    scaladsl.Consumer
      .committableSource(settings, subscription)
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * API MAY CHANGE
   *
   * This source emits `ConsumerRecord` together with the offset position as flow context, thus makes it possible
   * to commit offset positions to Kafka.
   * This is useful when "at-least once delivery" is desired, as each message will likely be
   * delivered one time but in failure cases could be duplicated.
   *
   * It is intended to be used with Apache Pekko's [flow with context](https://pekko.apache.org/docs/pekko/current/stream/operators/Flow/asFlowWithContext.html)
   * and [[Producer.flowWithContext]].
   */
  @ApiMayChange
  def sourceWithOffsetContext[K, V](
      settings: ConsumerSettings[K, V],
      subscription: Subscription): SourceWithContext[ConsumerRecord[K, V], CommittableOffset, Control] =
    // TODO this could use `scaladsl committableSourceWithContext` but `mapMaterializedValue` is not available, yet
    // See https://github.com/akka/akka/issues/26836
    pekko.stream.scaladsl.Source
      .fromGraph(new SourceWithOffsetContext[K, V](settings, subscription))
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asSourceWithContext(_._2)
      .map(_._1)
      .asJava

  /**
   * API MAY CHANGE
   *
   * This source emits `ConsumerRecord` together with the offset position as flow context, thus makes it possible
   * to commit offset positions to Kafka.
   * This is useful when "at-least once delivery" is desired, as each message will likely be
   * delivered one time but in failure cases could be duplicated.
   *
   * It is intended to be used with Apache Pekko's [flow with context](https://pekko.apache.org/docs/pekko/current/stream/operators/Flow/asFlowWithContext.html)
   * and [[Producer.flowWithContext]].
   *
   * This variant makes it possible to add additional metadata (in the form of a string)
   * when an offset is committed based on the record. This can be useful (for example) to store information about which
   * node made the commit, what time the commit was made, the timestamp of the record etc.
   */
  @ApiMayChange
  def sourceWithOffsetContext[K, V](
      settings: ConsumerSettings[K, V],
      subscription: Subscription,
      metadataFromRecord: java.util.function.Function[ConsumerRecord[K, V], String])
      : SourceWithContext[ConsumerRecord[K, V], CommittableOffset, Control] =
    // TODO this could use `scaladsl committableSourceWithContext` but `mapMaterializedValue` is not available, yet
    // See https://github.com/akka/akka/issues/26836
    pekko.stream.scaladsl.Source
      .fromGraph(
        new SourceWithOffsetContext[K, V](settings,
          subscription,
          (record: ConsumerRecord[K, V]) => metadataFromRecord(record)))
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asSourceWithContext(_._2)
      .map(_._1)
      .asJava

  /**
   * The `commitWithMetadataSource` makes it possible to add additional metadata (in the form of a string)
   * when an offset is committed based on the record. This can be useful (for example) to store information about which
   * node made the commit, what time the commit was made, the timestamp of the record etc.
   */
  def commitWithMetadataSource[K, V](
      settings: ConsumerSettings[K, V],
      subscription: Subscription,
      metadataFromRecord: java.util.function.Function[ConsumerRecord[K, V], String])
      : Source[CommittableMessage[K, V], Control] =
    scaladsl.Consumer
      .commitWithMetadataSource(settings, subscription, (record: ConsumerRecord[K, V]) => metadataFromRecord(record))
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * Convenience for "at-most once delivery" semantics. The offset of each message is committed to Kafka
   * before being emitted downstream.
   */
  def atMostOnceSource[K, V](settings: ConsumerSettings[K, V],
      subscription: Subscription): Source[ConsumerRecord[K, V], Control] =
    scaladsl.Consumer
      .atMostOnceSource(settings, subscription)
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * The `plainPartitionedSource` is a way to track automatic partition assignment from kafka.
   * When a topic-partition is assigned to a consumer, this source will emit pairs with the assigned topic-partition and a corresponding
   * source of `ConsumerRecord`s.
   * When a topic-partition is revoked, the corresponding source completes.
   */
  def plainPartitionedSource[K, V](
      settings: ConsumerSettings[K, V],
      subscription: AutoSubscription): Source[Pair[TopicPartition, Source[ConsumerRecord[K, V], NotUsed]], Control] =
    scaladsl.Consumer
      .plainPartitionedSource(settings, subscription)
      .map {
        case (tp, source) => Pair(tp, source.asJava)
      }
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * The `plainPartitionedManualOffsetSource` is similar to [[#plainPartitionedSource]] but allows the use of an offset store outside
   * of Kafka, while retaining the automatic partition assignment. When a topic-partition is assigned to a consumer, the `getOffsetsOnAssign`
   * function will be called to retrieve the offset, followed by a seek to the correct spot in the partition.
   */
  def plainPartitionedManualOffsetSource[K, V](
      settings: ConsumerSettings[K, V],
      subscription: AutoSubscription,
      getOffsetsOnAssign: java.util.function.Function[java.util.Set[TopicPartition],
        CompletionStage[
          java.util.Map[TopicPartition, Long]]])
      : Source[Pair[TopicPartition, Source[ConsumerRecord[K, V], NotUsed]], Control] =
    scaladsl.Consumer
      .plainPartitionedManualOffsetSource(
        settings,
        subscription,
        (tps: Set[TopicPartition]) =>
          getOffsetsOnAssign(tps.asJava).asScala.map(_.asScala.toMap)(ExecutionContexts.parasitic),
        _ => ())
      .map {
        case (tp, source) => Pair(tp, source.asJava)
      }
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * The `plainPartitionedManualOffsetSource` is similar to [[#plainPartitionedSource]] but allows the use of an offset store outside
   * of Kafka, while retaining the automatic partition assignment. When a topic-partition is assigned to a consumer, the `getOffsetsOnAssign`
   * function will be called to retrieve the offset, followed by a seek to the correct spot in the partition.
   *
   * The `onRevoke` function gives the consumer a chance to store any uncommitted offsets, and do any other cleanup
   * that is required. Also allows the user access to the `onPartitionsRevoked` hook, useful for cleaning up any
   * partition-specific resources being used by the consumer.
   */
  def plainPartitionedManualOffsetSource[K, V](
      settings: ConsumerSettings[K, V],
      subscription: AutoSubscription,
      getOffsetsOnAssign: java.util.function.Function[java.util.Set[TopicPartition],
        CompletionStage[
          java.util.Map[TopicPartition, Long]]],
      onRevoke: java.util.function.Consumer[java.util.Set[TopicPartition]])
      : Source[Pair[TopicPartition, Source[ConsumerRecord[K, V], NotUsed]], Control] =
    scaladsl.Consumer
      .plainPartitionedManualOffsetSource(
        settings,
        subscription,
        (tps: Set[TopicPartition]) =>
          getOffsetsOnAssign(tps.asJava).asScala.map(_.asScala.toMap)(ExecutionContexts.parasitic),
        (tps: Set[TopicPartition]) => onRevoke.accept(tps.asJava))
      .map {
        case (tp, source) => Pair(tp, source.asJava)
      }
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * The same as [[#plainPartitionedSource]] but with offset commit support.
   */
  def committablePartitionedSource[K, V](
      settings: ConsumerSettings[K, V],
      subscription: AutoSubscription)
      : Source[Pair[TopicPartition, Source[CommittableMessage[K, V], NotUsed]], Control] =
    scaladsl.Consumer
      .committablePartitionedSource(settings, subscription)
      .map {
        case (tp, source) => Pair(tp, source.asJava)
      }
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * The same as [[#plainPartitionedManualOffsetSource]] but with offset commit support.
   */
  def committablePartitionedManualOffsetSource[K, V](
      settings: ConsumerSettings[K, V],
      subscription: AutoSubscription,
      getOffsetsOnAssign: java.util.function.Function[java.util.Set[TopicPartition],
        CompletionStage[
          java.util.Map[TopicPartition, Long]]])
      : Source[Pair[TopicPartition, Source[CommittableMessage[K, V], NotUsed]], Control] =
    scaladsl.Consumer
      .committablePartitionedManualOffsetSource(
        settings,
        subscription,
        (tps: Set[TopicPartition]) =>
          getOffsetsOnAssign(tps.asJava).asScala.map(_.asScala.toMap)(ExecutionContexts.parasitic),
        _ => ())
      .map {
        case (tp, source) => Pair(tp, source.asJava)
      }
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * The same as [[#plainPartitionedManualOffsetSource]] but with offset commit support.
   */
  def committablePartitionedManualOffsetSource[K, V](
      settings: ConsumerSettings[K, V],
      subscription: AutoSubscription,
      getOffsetsOnAssign: java.util.function.Function[java.util.Set[TopicPartition],
        CompletionStage[
          java.util.Map[TopicPartition, Long]]],
      onRevoke: java.util.function.Consumer[java.util.Set[TopicPartition]])
      : Source[Pair[TopicPartition, Source[CommittableMessage[K, V], NotUsed]], Control] =
    scaladsl.Consumer
      .committablePartitionedManualOffsetSource(
        settings,
        subscription,
        (tps: Set[TopicPartition]) =>
          getOffsetsOnAssign(tps.asJava).asScala.map(_.asScala.toMap)(ExecutionContexts.parasitic),
        (tps: Set[TopicPartition]) => onRevoke.accept(tps.asJava))
      .map {
        case (tp, source) => Pair(tp, source.asJava)
      }
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * The same as [[#plainPartitionedSource]] but with offset commit with metadata support.
   */
  def commitWithMetadataPartitionedSource[K, V](
      settings: ConsumerSettings[K, V],
      subscription: AutoSubscription,
      metadataFromRecord: java.util.function.Function[ConsumerRecord[K, V], String])
      : Source[Pair[TopicPartition, Source[CommittableMessage[K, V], NotUsed]], Control] =
    scaladsl.Consumer
      .commitWithMetadataPartitionedSource(settings,
        subscription,
        (record: ConsumerRecord[K, V]) => metadataFromRecord(record))
      .map {
        case (tp, source) => Pair(tp, source.asJava)
      }
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * Special source that can use an external `KafkaAsyncConsumer`. This is useful when you have
   * a lot of manually assigned topic-partitions and want to keep only one kafka consumer.
   */
  def plainExternalSource[K, V](consumer: ActorRef,
      subscription: ManualSubscription): Source[ConsumerRecord[K, V], Control] =
    scaladsl.Consumer
      .plainExternalSource(consumer, subscription)
      .mapMaterializedValue(ConsumerControlAsJava.apply)
      .asJava

  /**
   * The same as [[#plainExternalSource]] but with offset commit support.
   */
  def committableExternalSource[K, V](consumer: ActorRef,
      subscription: ManualSubscription,
      groupId: String,
      commitTimeout: FiniteDuration): Source[CommittableMessage[K, V], Control] =
    scaladsl.Consumer
      .committableExternalSource(consumer, subscription, groupId, commitTimeout)
      .mapMaterializedValue(new ConsumerControlAsJava(_))
      .asJava
      .asInstanceOf[Source[CommittableMessage[K, V], Control]]

}
