/*
 * Copyright (C) 2014-2018 Lightbend Inc. <https://www.lightbend.com>
 */

package akka.actor

import java.util.concurrent.{ CountDownLatch, TimeUnit }

import scala.concurrent.Await
import scala.concurrent.duration.Duration
import scala.concurrent.duration._

object BenchmarkActors {

  val timeout = 30.seconds

  case object Message
  case object Stop

  class PingPong(val messagesPerPair: Int, latch: CountDownLatch) extends Actor {
    var left = messagesPerPair / 2
    def receive = {
      case Message ⇒

        if (left == 0) {
          latch.countDown()
          context stop self
        }

        sender() ! Message
        left -= 1
    }
  }

  object PingPong {
    def props(messagesPerPair: Int, latch: CountDownLatch) = Props(new PingPong(messagesPerPair, latch))
  }

  class Echo extends Actor {
    def receive = {
      case Message ⇒
        sender() ! Message
    }
  }

  object EchoSender {
    def props(messagesPerPair: Int, latch: CountDownLatch, batchSize: Int): Props =
      Props(new EchoSender(messagesPerPair, latch, batchSize))
  }

  class EchoSender(messagesPerPair: Int, latch: CountDownLatch, batchSize: Int) extends Actor {
    private val echo = context.actorOf(Props[Echo].withDispatcher(context.props.dispatcher), "echo")

    private var left = messagesPerPair / 2
    private var batch = 0

    def receive = {
      case Message ⇒
        batch -= 1
        if (batch <= 0) {
          if (!sendBatch()) {
            latch.countDown()
            context.stop(self)
          }
        }
    }

    private def sendBatch(): Boolean = {
      if (left > 0) {
        var i = 0
        while (i < batchSize) {
          echo ! Message
          i += 1
        }
        left -= batchSize
        batch = batchSize
        true
      } else
        false
    }
  }

  class Pipe(next: Option[ActorRef]) extends Actor {
    def receive = {
      case Message ⇒
        if (next.isDefined) next.get forward Message
      case Stop ⇒
        context stop self
        if (next.isDefined) next.get forward Stop
    }
  }

  object Pipe {
    def props(next: Option[ActorRef]) = Props(new Pipe(next))
  }

  private def startPingPongActorPairs(messagesPerPair: Int, numPairs: Int, dispatcher: String)(implicit system: ActorSystem) = {
    val fullPathToDispatcher = "akka.actor." + dispatcher
    val latch = new CountDownLatch(numPairs * 2)
    val actors = for {
      i ← (1 to numPairs).toVector
    } yield {
      val ping = system.actorOf(PingPong.props(messagesPerPair, latch).withDispatcher(fullPathToDispatcher))
      val pong = system.actorOf(PingPong.props(messagesPerPair, latch).withDispatcher(fullPathToDispatcher))
      (ping, pong)
    }
    (actors, latch)
  }

  private def initiatePingPongForPairs(refs: Vector[(ActorRef, ActorRef)], inFlight: Int) = {
    for {
      (ping, pong) ← refs
      _ ← 1 to inFlight
    } {
      ping.tell(Message, pong)
    }
  }

  private def startEchoActorPairs(messagesPerPair: Int, numPairs: Int, dispatcher: String,
                                  batchSize: Int)(implicit system: ActorSystem) = {

    val fullPathToDispatcher = "akka.actor." + dispatcher
    val latch = new CountDownLatch(numPairs)
    val actors = (1 to numPairs).map { _ ⇒
      system.actorOf(EchoSender.props(messagesPerPair, latch, batchSize).withDispatcher(fullPathToDispatcher))
    }.toVector
    (actors, latch)
  }

  private def initiateEchoPairs(refs: Vector[ActorRef]) = {
    refs.foreach(_ ! Message)
  }

  def printProgress(totalMessages: Long, numActors: Int, startNanoTime: Long) = {
    val durationMicros = (System.nanoTime() - startNanoTime) / 1000
    println(f"  $totalMessages messages by $numActors actors took ${durationMicros / 1000} ms, " +
      f"${totalMessages.toDouble / durationMicros}%,.2f M msg/s")
  }

  def requireRightNumberOfCores(numCores: Int) =
    require(
      Runtime.getRuntime.availableProcessors == numCores,
      s"Update the cores constant to ${Runtime.getRuntime.availableProcessors}"
    )

  def benchmarkPingPongActors(numMessagesPerActorPair: Int, numActors: Int, dispatcher: String, throughPut: Int, shutdownTimeout: Duration)(implicit system: ActorSystem): Unit = {
    val numPairs = numActors / 2
    val totalNumMessages = numPairs * numMessagesPerActorPair
    val (actors, latch) = startPingPongActorPairs(numMessagesPerActorPair, numPairs, dispatcher)
    val startNanoTime = System.nanoTime()
    initiatePingPongForPairs(actors, inFlight = throughPut * 2)
    latch.await(shutdownTimeout.toSeconds, TimeUnit.SECONDS)
    printProgress(totalNumMessages, numActors, startNanoTime)
  }

  def benchmarkEchoActors(numMessagesPerActorPair: Int, numActors: Int, dispatcher: String, batchSize: Int, shutdownTimeout: Duration)(implicit system: ActorSystem): Unit = {
    val numPairs = numActors / 2
    val totalNumMessages = numPairs * numMessagesPerActorPair
    val (actors, latch) = startEchoActorPairs(numMessagesPerActorPair, numPairs, dispatcher, batchSize)
    val startNanoTime = System.nanoTime()
    initiateEchoPairs(actors)
    latch.await(shutdownTimeout.toSeconds, TimeUnit.SECONDS)
    printProgress(totalNumMessages, numActors, startNanoTime)
  }

  def tearDownSystem()(implicit system: ActorSystem): Unit = {
    system.terminate()
    Await.ready(system.whenTerminated, timeout)
  }

}

