/*
 * 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) 2015-2022 Lightbend Inc. <https://www.lightbend.com>
 */

package docs.stream

import org.apache.pekko.NotUsed
import org.apache.pekko.stream.{ KillSwitches, UniqueKillSwitch }
import org.apache.pekko.stream.scaladsl._
import org.apache.pekko.testkit.PekkoSpec
import docs.CompileOnlySpec

import scala.concurrent.duration._
import org.apache.pekko.stream.ThrottleMode

class HubsDocSpec extends PekkoSpec with CompileOnlySpec {

  "Hubs" must {

    "demonstrate creating a dynamic merge" in {
      def println(s: String) = testActor ! s

      // #merge-hub
      // A simple consumer that will print to the console for now
      val consumer = Sink.foreach(println)

      // Attach a MergeHub Source to the consumer. This will materialize to a
      // corresponding Sink.
      val runnableGraph: RunnableGraph[Sink[String, NotUsed]] =
        MergeHub.source[String](perProducerBufferSize = 16).to(consumer)

      // By running/materializing the consumer we get back a Sink, and hence
      // now have access to feed elements into it. This Sink can be materialized
      // any number of times, and every element that enters the Sink will
      // be consumed by our consumer.
      val toConsumer: Sink[String, NotUsed] = runnableGraph.run()

      // Feeding two independent sources into the hub.
      Source.single("Hello!").runWith(toConsumer)
      Source.single("Hub!").runWith(toConsumer)
      // #merge-hub

      expectMsgAllOf("Hello!", "Hub!")
    }

    "demonstrate creating a dynamic broadcast" in compileOnlySpec {
      // #broadcast-hub
      // A simple producer that publishes a new "message" every second
      val producer = Source.tick(1.second, 1.second, "New message")

      // Attach a BroadcastHub Sink to the producer. This will materialize to a
      // corresponding Source.
      // (We need to use toMat and Keep.right since by default the materialized
      // value to the left is used)
      val runnableGraph: RunnableGraph[Source[String, NotUsed]] =
        producer.toMat(BroadcastHub.sink(bufferSize = 256))(Keep.right)

      // By running/materializing the producer, we get back a Source, which
      // gives us access to the elements published by the producer.
      val fromProducer: Source[String, NotUsed] = runnableGraph.run()

      // Print out messages from the producer in two independent consumers
      fromProducer.runForeach(msg => println("consumer1: " + msg))
      fromProducer.runForeach(msg => println("consumer2: " + msg))
      // #broadcast-hub
    }

    "demonstrate combination" in {
      def println(s: String) = testActor ! s

      // #pub-sub-1
      // Obtain a Sink and Source which will publish and receive from the "bus" respectively.
      val (sink, source) =
        MergeHub.source[String](perProducerBufferSize = 16).toMat(BroadcastHub.sink(bufferSize = 256))(Keep.both).run()
      // #pub-sub-1

      // #pub-sub-2
      // Ensure that the Broadcast output is dropped if there are no listening parties.
      // If this dropping Sink is not attached, then the broadcast hub will not drop any
      // elements itself when there are no subscribers, backpressuring the producer instead.
      source.runWith(Sink.ignore)
      // #pub-sub-2

      // #pub-sub-3
      // We create now a Flow that represents a publish-subscribe channel using the above
      // started stream as its "topic". We add two more features, external cancellation of
      // the registration and automatic cleanup for very slow subscribers.
      val busFlow: Flow[String, String, UniqueKillSwitch] =
        Flow
          .fromSinkAndSource(sink, source)
          .joinMat(KillSwitches.singleBidi[String, String])(Keep.right)
          .backpressureTimeout(3.seconds)
      // #pub-sub-3

      // #pub-sub-4
      val switch: UniqueKillSwitch =
        Source.repeat("Hello world!").viaMat(busFlow)(Keep.right).to(Sink.foreach(println)).run()

      // Shut down externally
      switch.shutdown()
      // #pub-sub-4
    }

    "demonstrate creating a dynamic partition hub" in compileOnlySpec {
      // #partition-hub
      // A simple producer that publishes a new "message-" every second
      val producer = Source.tick(1.second, 1.second, "message").zipWith(Source(1 to 100))((a, b) => s"$a-$b")

      // Attach a PartitionHub Sink to the producer. This will materialize to a
      // corresponding Source.
      // (We need to use toMat and Keep.right since by default the materialized
      // value to the left is used)
      val runnableGraph: RunnableGraph[Source[String, NotUsed]] =
        producer.toMat(
          PartitionHub.sink(
            (size, elem) => math.abs(elem.hashCode % size),
            startAfterNrOfConsumers = 2,
            bufferSize = 256))(Keep.right)

      // By running/materializing the producer, we get back a Source, which
      // gives us access to the elements published by the producer.
      val fromProducer: Source[String, NotUsed] = runnableGraph.run()

      // Print out messages from the producer in two independent consumers
      fromProducer.runForeach(msg => println("consumer1: " + msg))
      fromProducer.runForeach(msg => println("consumer2: " + msg))
      // #partition-hub
    }

    "demonstrate creating a dynamic stateful partition hub" in compileOnlySpec {
      // #partition-hub-stateful
      // A simple producer that publishes a new "message-" every second
      val producer = Source.tick(1.second, 1.second, "message").zipWith(Source(1 to 100))((a, b) => s"$a-$b")

      // New instance of the partitioner function and its state is created
      // for each materialization of the PartitionHub.
      def roundRobin(): (PartitionHub.ConsumerInfo, String) => Long = {
        var i = -1L

        (info, elem) => {
          i += 1
          info.consumerIdByIdx((i % info.size).toInt)
        }
      }

      // Attach a PartitionHub Sink to the producer. This will materialize to a
      // corresponding Source.
      // (We need to use toMat and Keep.right since by default the materialized
      // value to the left is used)
      val runnableGraph: RunnableGraph[Source[String, NotUsed]] =
        producer.toMat(PartitionHub.statefulSink(() => roundRobin(), startAfterNrOfConsumers = 2, bufferSize = 256))(
          Keep.right)

      // By running/materializing the producer, we get back a Source, which
      // gives us access to the elements published by the producer.
      val fromProducer: Source[String, NotUsed] = runnableGraph.run()

      // Print out messages from the producer in two independent consumers
      fromProducer.runForeach(msg => println("consumer1: " + msg))
      fromProducer.runForeach(msg => println("consumer2: " + msg))
      // #partition-hub-stateful
    }

    "demonstrate creating a dynamic partition hub routing to fastest consumer" in compileOnlySpec {
      // #partition-hub-fastest
      val producer = Source(0 until 100)

      // ConsumerInfo.queueSize is the approximate number of buffered elements for a consumer.
      // Note that this is a moving target since the elements are consumed concurrently.
      val runnableGraph: RunnableGraph[Source[Int, NotUsed]] =
        producer.toMat(
          PartitionHub.statefulSink(
            () => (info, elem) => info.consumerIds.minBy(id => info.queueSize(id)),
            startAfterNrOfConsumers = 2,
            bufferSize = 16))(Keep.right)

      val fromProducer: Source[Int, NotUsed] = runnableGraph.run()

      fromProducer.runForeach(msg => println("consumer1: " + msg))
      fromProducer.throttle(10, 100.millis).runForeach(msg => println("consumer2: " + msg))
      // #partition-hub-fastest
    }

  }

}
