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

package akka.remote.artery

import java.util.concurrent.atomic.AtomicReference

import scala.annotation.tailrec
import akka.actor.Address
import akka.remote.transport.ThrottlerTransportAdapter.Direction
import akka.stream.Attributes
import akka.stream.FlowShape
import akka.stream.Inlet
import akka.stream.Outlet
import akka.stream.stage._
import akka.util.OptionVal
import akka.event.Logging

object TestManagementCommands {
  /** INTERNAL API */
  @SerialVersionUID(1L)
  final case class FailInboundStreamOnce(ex: Throwable)
}

/**
 * INTERNAL API: Thread safe mutable state that is shared among
 * the test operators.
 */
private[remote] class SharedTestState {

  private val state = new AtomicReference[TestState](TestState(Map.empty, None))

  def isBlackhole(from: Address, to: Address): Boolean =
    state.get.blackholes.get(from) match {
      case Some(destinations) ⇒ destinations(to)
      case None               ⇒ false
    }

  /** Enable blackholing between given address in given direction */
  def blackhole(a: Address, b: Address, direction: Direction): Unit =
    direction match {
      case Direction.Send    ⇒ addBlackhole(a, b)
      case Direction.Receive ⇒ addBlackhole(b, a)
      case Direction.Both ⇒
        addBlackhole(a, b)
        addBlackhole(b, a)
    }

  /**
   * Cause the inbound stream to fail with the given exception.
   * Can be used to test inbound stream restart / recovery.
   */
  @tailrec final def failInboundStreamOnce(ex: Throwable): Unit = {
    val current = state.get
    if (state.compareAndSet(current, current.copy(failInboundStream = Some(ex)))) ()
    else failInboundStreamOnce(ex)
  }
  /**
   * Get the exception to fail the inbound stream with and immediately reset the state to not-failed.
   * This is used to simulate a single failure on the stream, where a successful restart recovers operations.
   */
  @tailrec final def getInboundFailureOnce: Option[Throwable] = {
    val current = state.get()
    if (state.compareAndSet(current, current.copy(failInboundStream = None))) current.failInboundStream
    else getInboundFailureOnce
  }

  @tailrec private def addBlackhole(from: Address, to: Address): Unit = {
    val current = state.get
    val newState = current.blackholes.get(from) match {
      case Some(destinations) ⇒ current.copy(blackholes = current.blackholes.updated(from, destinations + to))
      case None               ⇒ current.copy(blackholes = current.blackholes.updated(from, Set(to)))
    }
    if (!state.compareAndSet(current, newState))
      addBlackhole(from, to)
  }

  def passThrough(a: Address, b: Address, direction: Direction): Unit =
    direction match {
      case Direction.Send    ⇒ removeBlackhole(a, b)
      case Direction.Receive ⇒ removeBlackhole(b, a)
      case Direction.Both ⇒
        removeBlackhole(a, b)
        removeBlackhole(b, a)
    }

  @tailrec private def removeBlackhole(from: Address, to: Address): Unit = {
    val current = state.get
    val newState = current.blackholes.get(from) match {
      case Some(destinations) ⇒ current.copy(blackholes = current.blackholes.updated(from, destinations - to))
      case None               ⇒ current
    }
    if (!state.compareAndSet(current, newState))
      removeBlackhole(from, to)
  }

}

/**
 * INTERNAL API
 */
private[remote] final case class TestState(
  blackholes:        Map[Address, Set[Address]],
  failInboundStream: Option[Throwable])

/**
 * INTERNAL API
 */
private[remote] class OutboundTestStage(outboundContext: OutboundContext, state: SharedTestState)
  extends GraphStage[FlowShape[OutboundEnvelope, OutboundEnvelope]] {
  val in: Inlet[OutboundEnvelope] = Inlet("OutboundTestStage.in")
  val out: Outlet[OutboundEnvelope] = Outlet("OutboundTestStage.out")
  override val shape: FlowShape[OutboundEnvelope, OutboundEnvelope] = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes) =
    new TimerGraphStageLogic(shape) with InHandler with OutHandler with StageLogging {

      // InHandler
      override def onPush(): Unit = {
        val env = grab(in)
        if (state.isBlackhole(outboundContext.localAddress.address, outboundContext.remoteAddress)) {
          log.debug(
            "dropping outbound message [{}] to [{}] because of blackhole",
            Logging.messageClassName(env.message), outboundContext.remoteAddress)
          pull(in) // drop message
        } else
          push(out, env)
      }

      // OutHandler
      override def onPull(): Unit = pull(in)

      setHandlers(in, out, this)
    }

}

/**
 * INTERNAL API
 */
private[remote] class InboundTestStage(inboundContext: InboundContext, state: SharedTestState)
  extends GraphStage[FlowShape[InboundEnvelope, InboundEnvelope]] {
  val in: Inlet[InboundEnvelope] = Inlet("InboundTestStage.in")
  val out: Outlet[InboundEnvelope] = Outlet("InboundTestStage.out")
  override val shape: FlowShape[InboundEnvelope, InboundEnvelope] = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes) =
    new TimerGraphStageLogic(shape) with InHandler with OutHandler with StageLogging {

      // InHandler
      override def onPush(): Unit = {
        state.getInboundFailureOnce match {
          case Some(shouldFailEx) ⇒
            log.info("Fail inbound stream from [{}]: {}", classOf[InboundTestStage].getName, shouldFailEx.getMessage)
            failStage(shouldFailEx)
          case _ ⇒
            val env = grab(in)
            env.association match {
              case OptionVal.None ⇒
                // unknown, handshake not completed
                push(out, env)
              case OptionVal.Some(association) ⇒
                if (state.isBlackhole(inboundContext.localAddress.address, association.remoteAddress)) {
                  log.debug(
                    "dropping inbound message [{}] from [{}] with UID [{}] because of blackhole",
                    Logging.messageClassName(env.message), association.remoteAddress, env.originUid)
                  pull(in) // drop message
                } else
                  push(out, env)
            }
        }
      }

      // OutHandler
      override def onPull(): Unit = pull(in)

      setHandlers(in, out, this)
    }

}

