package net.corda.node.services.messaging

import co.paralleluniverse.fibers.Suspendable
import com.codahale.metrics.MetricRegistry
import net.corda.core.crypto.toStringShort
import net.corda.core.flows.StateMachineRunId
import net.corda.core.identity.CordaX500Name
import net.corda.core.internal.NamedCacheFactory
import net.corda.core.internal.ThreadBox
import net.corda.core.messaging.CordaRPCOps
import net.corda.core.messaging.MessageRecipients
import net.corda.core.messaging.SingleMessageRecipient
import net.corda.core.node.NodeInfo
import net.corda.core.node.services.NetworkMapCache
import net.corda.core.node.services.PartyInfo
import net.corda.core.serialization.SerializationDefaults
import net.corda.core.serialization.SingletonSerializeAsToken
import net.corda.core.serialization.deserialize
import net.corda.core.serialization.internal.nodeSerializationEnv
import net.corda.core.serialization.serialize
import net.corda.core.utilities.ByteSequence
import net.corda.core.utilities.NetworkHostAndPort
import net.corda.core.utilities.OpaqueBytes
import net.corda.core.utilities.contextLogger
import net.corda.core.utilities.debug
import net.corda.core.utilities.loggerFor
import net.corda.core.utilities.minutes
import net.corda.core.utilities.seconds
import net.corda.core.utilities.trace
import net.corda.node.VersionInfo
import net.corda.node.internal.LifecycleSupport
import net.corda.node.internal.artemis.ReactiveArtemisConsumer
import net.corda.node.internal.artemis.ReactiveArtemisConsumer.Companion.multiplex
import net.corda.node.services.api.NetworkMapCacheInternal
import net.corda.node.services.config.NodeConfiguration
import net.corda.node.services.statemachine.DeduplicationId
import net.corda.node.services.statemachine.ExternalEvent
import net.corda.node.services.statemachine.SenderDeduplicationId
import net.corda.node.utilities.AffinityExecutor
import net.corda.node.utilities.errorAndTerminate
import net.corda.nodeapi.internal.ArtemisMessagingComponent
import net.corda.nodeapi.internal.ArtemisMessagingComponent.ArtemisAddress
import net.corda.nodeapi.internal.ArtemisMessagingComponent.Companion.BRIDGE_CONTROL
import net.corda.nodeapi.internal.ArtemisMessagingComponent.Companion.BRIDGE_NOTIFY
import net.corda.nodeapi.internal.ArtemisMessagingComponent.Companion.JOURNAL_HEADER_SIZE
import net.corda.nodeapi.internal.ArtemisMessagingComponent.Companion.P2PMessagingHeaders
import net.corda.nodeapi.internal.ArtemisMessagingComponent.Companion.PEERS_PREFIX
import net.corda.nodeapi.internal.ArtemisMessagingComponent.NodeAddress
import net.corda.nodeapi.internal.ArtemisMessagingComponent.RemoteInboxAddress
import net.corda.nodeapi.internal.ArtemisMessagingComponent.ServiceAddress
import net.corda.nodeapi.internal.ArtemisTcpTransport.Companion.p2pConnectorTcpTransport
import net.corda.nodeapi.internal.bridging.BridgeControl
import net.corda.nodeapi.internal.bridging.BridgeEntry
import net.corda.nodeapi.internal.lifecycle.ServiceStateHelper
import net.corda.nodeapi.internal.lifecycle.ServiceStateSupport
import net.corda.nodeapi.internal.persistence.CordaPersistence
import net.corda.nodeapi.internal.requireMessageSize
import org.apache.activemq.artemis.api.config.ActiveMQDefaultConfiguration
import org.apache.activemq.artemis.api.core.ActiveMQObjectClosedException
import org.apache.activemq.artemis.api.core.Message.HDR_DUPLICATE_DETECTION_ID
import org.apache.activemq.artemis.api.core.Message.HDR_VALIDATED_USER
import org.apache.activemq.artemis.api.core.QueueConfiguration
import org.apache.activemq.artemis.api.core.RoutingType
import org.apache.activemq.artemis.api.core.SimpleString
import org.apache.activemq.artemis.api.core.client.ActiveMQClient
import org.apache.activemq.artemis.api.core.client.ClientConsumer
import org.apache.activemq.artemis.api.core.client.ClientMessage
import org.apache.activemq.artemis.api.core.client.ClientProducer
import org.apache.activemq.artemis.api.core.client.ClientSession
import org.apache.activemq.artemis.api.core.client.ClientSessionFactory
import org.apache.activemq.artemis.api.core.client.FailoverEventType
import org.apache.activemq.artemis.api.core.client.ServerLocator
import rx.Observable
import rx.Subscription
import rx.subjects.PublishSubject
import java.security.PublicKey
import java.time.Duration
import java.time.Instant
import java.util.Collections
import java.util.Timer
import java.util.UUID
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.CountDownLatch
import javax.annotation.concurrent.ThreadSafe
import kotlin.concurrent.timer

/**
 * This class implements the [MessagingService] API using Apache Artemis, the successor to their ActiveMQ product.
 * Artemis is a message queue broker and here we run a client connecting to the specified broker instance
 * [ArtemisMessagingServer]. It's primarily concerned with peer-to-peer messaging.
 *
 * Message handlers are run on the provided [AffinityExecutor] synchronously, that is, the Artemis callback threads
 * are blocked until the handler is scheduled and completed. This allows backpressure to propagate from the given
 * executor through into Artemis and from there, back through to senders.
 *
 * An implementation of [CordaRPCOps] can be provided. If given, clients using the CordaMQClient RPC library can
 * invoke methods on the provided implementation. There is more documentation on this in the doc-site and the
 * CordaRPCClient class.
 *
 * @param config The configuration of the node, which is used for controlling the message redelivery options.
 * @param versionInfo All messages from the node carry the version info and received messages are checked against this for compatibility.
 * @param serverAddress The host and port of the Artemis broker.
 * @param nodeExecutor The received messages are marshalled onto the server executor to prevent Netty buffers leaking during fiber suspends.
 * @param database The node's database, which is used to deduplicate messages.
 * @param terminateOnConnectionError whether the process should be terminated forcibly if connection with the broker fails.
 */
@Suppress("LongParameterList")
@ThreadSafe
class P2PMessagingClient(val config: NodeConfiguration,
                         private val versionInfo: VersionInfo,
                         val serverAddress: NetworkHostAndPort,
                         private val nodeExecutor: AffinityExecutor.ServiceAffinityExecutor,
                         private val database: CordaPersistence,
                         private val networkMap: NetworkMapCacheInternal,
                         @Suppress("UNUSED")
                         private val metricRegistry: MetricRegistry,
                         cacheFactory: NamedCacheFactory,
                         private val isDrainingModeOn: () -> Boolean,
                         private val drainingModeWasChangedEvents: Observable<Pair<Boolean, Boolean>>,
                         private val threadPoolName: String = "P2PClient",
                         private val stateHelper: ServiceStateHelper = ServiceStateHelper(log),
                         private val terminateOnConnectionError: Boolean = true,
                         private val timeoutConfig: TimeoutConfig = TimeoutConfig.default()
) : SingletonSerializeAsToken(), MessagingService, AddressToArtemisQueueResolver, ServiceStateSupport by stateHelper {
    companion object {
        private val log = contextLogger()
    }

    private class NodeClientMessage(override val topic: String,
                                    override val data: ByteSequence,
                                    override val uniqueMessageId: DeduplicationId,
                                    override val senderUUID: String?,
                                    override val additionalHeaders: Map<String, String>) : Message {
        override val debugTimestamp: Instant = Instant.now()
        override fun toString() = "$topic#${String(data.bytes)}"
    }

    private class InnerState {
        var started = false
        var running = false
        var eventsSubscription: Subscription? = null
        var p2pConsumer: P2PMessagingConsumer? = null
        var locator: ServerLocator? = null
        var executorProducer: ClientProducer? = null
        var executorSession: ClientSession? = null
        var producer: ClientProducer? = null
        var producerSession: ClientSession? = null
        var bridgeSession: ClientSession? = null
        var bridgeNotifyConsumer: ClientConsumer? = null
        var networkChangeSubscription: Subscription? = null
        var sessionFactory: ClientSessionFactory? = null

        fun sendMessage(address: String, message: ClientMessage) = producer!!.send(address, message)
    }

    /**
     * @property callTimeout the time a blocking call (e.g. message send) from a client waits for a response until it times out.
     * @property serverConnectionTtl the time the server waits for a packet/heartbeat from a client before it announces the connection dead and cleans it up.
     * @property clientConnectionTtl the time the client waits for a packet/heartbeat from a client before it announces the connection dead and cleans it up.
     */
    data class TimeoutConfig(val callTimeout: Duration, val serverConnectionTtl: Duration, val clientConnectionTtl: Duration) {
        companion object {
            /**
             * Some sensible defaults, aligned with defaults of Artemis
             */
            @Suppress("MagicNumber")
            fun default() = TimeoutConfig(30.seconds, 60.seconds, 30.seconds)
        }
    }

    /** A registration to handle messages of different types */
    data class HandlerRegistration(val topic: String, val callback: Any) : MessageHandlerRegistration

    private lateinit var myIdentity: PublicKey
    private var serviceIdentity: PublicKey? = null
    private lateinit var advertisedAddress: NetworkHostAndPort
    private var maxMessageSize: Int = -1

    override val myAddress: SingleMessageRecipient get() = NodeAddress(myIdentity)
    override val ourSenderUUID = UUID.randomUUID().toString()

    private val state = ThreadBox(InnerState())
    private val knownQueues = Collections.newSetFromMap(ConcurrentHashMap<String, Boolean>())
    private val delayStartQueues = Collections.newSetFromMap(ConcurrentHashMap<String, Boolean>())

    private val handlers = ConcurrentHashMap<String, MessageHandler>()
    private val handlersChangedSignal = Object()

    private val deduplicator = P2PMessageDeduplicator(cacheFactory, database)
    internal var messagingExecutor: MessagingExecutor? = null

    /**
     * @param myIdentity The primary identity of the node, which defines the messaging address for externally received messages.
     * It is also used to construct the myAddress field, which is ultimately advertised in the network map.
     * @param serviceIdentity An optional second identity if the node is also part of a group address, for example a notary.
     * @param advertisedAddress The externally advertised version of the Artemis broker address used to construct myAddress and included
     * in the network map data.
     * @param maxMessageSize A bound applied to the message size.
     */
    fun start(myIdentity: PublicKey, serviceIdentity: PublicKey?, maxMessageSize: Int, advertisedAddress: NetworkHostAndPort = serverAddress) {
        this.myIdentity = myIdentity
        this.serviceIdentity = serviceIdentity
        this.advertisedAddress = advertisedAddress
        this.maxMessageSize = maxMessageSize
        state.locked {
            started = true
            log.info("Connecting to message broker: $serverAddress")
            // TODO Add broker CN to config for host verification in case the embedded broker isn't used
            val tcpTransport = p2pConnectorTcpTransport(serverAddress, config.p2pSslOptions, threadPoolName = threadPoolName)
            locator = ActiveMQClient.createServerLocatorWithoutHA(tcpTransport).apply {
                callTimeout = timeoutConfig.callTimeout.toMillis()
                connectionTTL = timeoutConfig.serverConnectionTtl.toMillis()
                clientFailureCheckPeriod = timeoutConfig.clientConnectionTtl.toMillis()
                minLargeMessageSize = maxMessageSize + JOURNAL_HEADER_SIZE
                isUseGlobalPools = nodeSerializationEnv != null
            }

            sessionFactory = if (terminateOnConnectionError) {
                locator!!.createSessionFactory().addFailoverListener(::failoverCallback)
            } else {
                locator!!.createSessionFactory()
            }
            // Login using the node username. The broker will authenticate us as its node (as opposed to another peer)
            // using our TLS certificate.
            // Note that the acknowledgement of messages is not flushed to the Artemis journal until the default buffer
            // size of 1MB is acknowledged.
            val createNewSession = { sessionFactory!!.createSession(ArtemisMessagingComponent.NODE_P2P_USER, ArtemisMessagingComponent.NODE_P2P_USER, false, true, true, false, ActiveMQClient.DEFAULT_ACK_BATCH_SIZE) }

            executorSession = createNewSession()
            producerSession = createNewSession()
            bridgeSession = createNewSession()
            executorSession!!.start()
            producerSession!!.start()
            bridgeSession!!.start()

            val inboxes = mutableSetOf<String>()
            // Create a queue, consumer and producer for handling P2P network messages.
            // Create a general purpose producer.
            producer = producerSession!!.createProducer()
            executorProducer = executorSession!!.createProducer()

            inboxes += RemoteInboxAddress(myIdentity).queueName
            serviceIdentity?.let {
                inboxes += RemoteInboxAddress(it).queueName
            }

            inboxes.forEach { createQueueIfAbsent(it, producerSession!!, exclusive = true, isServiceAddress = false) }

            p2pConsumer = P2PMessagingConsumer(inboxes, createNewSession, isDrainingModeOn, drainingModeWasChangedEvents, metricRegistry)

            messagingExecutor = MessagingExecutor(
                    executorSession!!,
                    executorProducer!!,
                    versionInfo,
                    this@P2PMessagingClient,
                    ourSenderUUID = ourSenderUUID
            )

            registerBridgeControl(bridgeSession!!, inboxes.toList())
            enumerateBridges(bridgeSession!!, inboxes.toList())
        }
    }

    private fun failoverCallback(event: FailoverEventType) {
        when (event) {
            FailoverEventType.FAILURE_DETECTED -> {
                errorAndTerminate("Connection to the broker was lost. Node is shutting down.", null)
            }
            FailoverEventType.FAILOVER_FAILED -> state.locked {
                if (running) {
                    errorAndTerminate("Could not reconnect to the broker. Node is shutting down.", null)
                }
            }
            else -> {
                log.warn("Cannot handle event $event.")
            }
        }
    }

    private fun InnerState.registerBridgeControl(session: ClientSession, inboxes: List<String>) {
        val bridgeNotifyQueue = "$BRIDGE_NOTIFY.${myIdentity.toStringShort()}"
        if (!session.queueQuery(SimpleString(bridgeNotifyQueue)).isExists) {
            session.createQueue(QueueConfiguration(bridgeNotifyQueue).setAddress(BRIDGE_NOTIFY).setRoutingType(RoutingType.MULTICAST)
                    .setTemporary(true).setDurable(false))
        }
        val bridgeConsumer = session.createConsumer(bridgeNotifyQueue)
        bridgeNotifyConsumer = bridgeConsumer
        bridgeConsumer.setMessageHandler { msg ->
            state.locked {
                val data: ByteArray = ByteArray(msg.bodySize).apply { msg.bodyBuffer.readBytes(this) }
                val notifyMessage = data.deserialize<BridgeControl>(context = SerializationDefaults.P2P_CONTEXT)
                log.info(notifyMessage.toString())
                when (notifyMessage) {
                    is BridgeControl.BridgeToNodeSnapshotRequest -> enumerateBridges(session, inboxes)
                    else -> log.error("Unexpected Bridge Control message type on notify topic $notifyMessage")
                }
                msg.acknowledge()
            }
        }
        networkChangeSubscription = networkMap.changed.subscribe { updateBridgesOnNetworkChange(it) }
    }

    private fun sendBridgeControl(message: BridgeControl) {
        state.locked {
            val controlPacket = message.serialize(context = SerializationDefaults.P2P_CONTEXT).bytes
            val artemisMessage = producerSession!!.createMessage(false)
            artemisMessage.writeBodyBufferBytes(controlPacket)
            sendMessage(BRIDGE_CONTROL, artemisMessage)
        }
    }

    private fun updateBridgesOnNetworkChange(change: NetworkMapCache.MapChange) {
        log.info("Updating bridges on network map change: ${change::class.simpleName} ${change.node}")
        fun gatherAddresses(node: NodeInfo): Sequence<BridgeEntry> {
            return state.locked {
                node.legalIdentitiesAndCerts.map { partyAndCertificate ->
                    val messagingAddress = NodeAddress(partyAndCertificate.party.owningKey)
                    BridgeEntry(messagingAddress.queueName, node.addresses, node.legalIdentities.map { it.name }, serviceAddress = false)
                }.filter { producerSession!!.queueQuery(SimpleString(it.queueName)).isExists }.asSequence()
            }
        }

        fun deployBridges(node: NodeInfo) {
            gatherAddresses(node)
                    .forEach {
                        sendBridgeControl(BridgeControl.Create(config.myLegalName.toString(), it))
                    }
        }

        fun destroyBridges(node: NodeInfo) {
            gatherAddresses(node)
                    .forEach {
                        sendBridgeControl(BridgeControl.Delete(config.myLegalName.toString(), it))
                    }
        }

        when (change) {
            is NetworkMapCache.MapChange.Added -> {
                deployBridges(change.node)
            }
            is NetworkMapCache.MapChange.Removed -> {
                destroyBridges(change.node)
            }
            is NetworkMapCache.MapChange.Modified -> {
                destroyBridges(change.previousNode)
                deployBridges(change.node)
            }
        }
    }

    private fun enumerateBridges(session: ClientSession, inboxes: List<String>) {
        val requiredBridges = mutableListOf<BridgeEntry>()
        fun createBridgeEntry(queueName: SimpleString) {
            val keyHash = queueName.substring(PEERS_PREFIX.length)
            val peers = networkMap.getNodesByOwningKeyIndex(keyHash)
            for (node in peers) {
                val bridge = BridgeEntry(queueName.toString(), node.addresses, node.legalIdentities.map { it.name }, serviceAddress = false)
                requiredBridges += bridge
                knownQueues += queueName.toString()
            }
        }

        val queues = session.addressQuery(SimpleString("$PEERS_PREFIX#")).queueNames
        knownQueues.clear()
        for (queue in queues) {
            val queueQuery = session.queueQuery(queue)
            if (!config.lazyBridgeStart || queueQuery.messageCount > 0) {
                createBridgeEntry(queue)
            } else {
                delayStartQueues += queue.toString()
            }
        }
        val startupMessage = BridgeControl.NodeToBridgeSnapshot(config.myLegalName.toString(), inboxes, requiredBridges)
        sendBridgeControl(startupMessage)
    }

    private val shutdownLatch = CountDownLatch(1)

    /**
     * Starts the p2p event loop: this method only returns once [stop] has been called.
     */
    override fun start() {
        val latch = CountDownLatch(1)
        try {
            synchronized(handlersChangedSignal) {
                while (handlers.isEmpty() && state.locked { (p2pConsumer != null) }) {
                    handlersChangedSignal.wait()
                }
            }
            val consumer = state.locked {
                check(started) { "start must be called first" }
                check(!running) { "run can't be called twice" }
                running = true
                // If it's null, it means we already called stop, so return immediately.
                if (p2pConsumer == null) {
                    return
                }
                eventsSubscription = p2pConsumer!!.messages
                        // this `run()` method is semantically meant to block until the message consumption runs, hence the latch here
                        .doOnCompleted(latch::countDown)
                        .subscribe({ message -> deliver(message) }, { error -> throw error })
                p2pConsumer!!
            }
            consumer.start()
            log.debug("Signalling active")
            stateHelper.active = true
            log.debug("Awaiting on latch")
            latch.await()
        } finally {
            shutdownLatch.countDown()
        }
    }

    private fun artemisToCordaMessage(message: ClientMessage): ReceivedMessage? {
        try {
            requireMessageSize(message.bodySize, maxMessageSize)
            val topic = message.required(P2PMessagingHeaders.topicProperty) { getStringProperty(it) }
            val user = requireNotNull(message.getStringProperty(HDR_VALIDATED_USER)) { "Message is not authenticated" }
            val platformVersion = message.required(P2PMessagingHeaders.platformVersionProperty) { getIntProperty(it) }
            // Use the magic deduplication property built into Artemis as our message identity too
            val uniqueMessageId = message.required(HDR_DUPLICATE_DETECTION_ID) { DeduplicationId(message.getStringProperty(it)) }
            val receivedSenderUUID = message.getStringProperty(P2PMessagingHeaders.senderUUID)
            val receivedSenderSeqNo = if (message.containsProperty(P2PMessagingHeaders.senderSeqNo)) message.getLongProperty(P2PMessagingHeaders.senderSeqNo) else null
            val isSessionInit = message.getStringProperty(P2PMessagingHeaders.Type.KEY) == P2PMessagingHeaders.Type.SESSION_INIT_VALUE
            log.trace { "Received message from: ${message.address} user: $user topic: $topic id: $uniqueMessageId senderUUID: $receivedSenderUUID senderSeqNo: $receivedSenderSeqNo isSessionInit: $isSessionInit" }

            return ArtemisReceivedMessage(topic, CordaX500Name.parse(user), platformVersion, uniqueMessageId, receivedSenderUUID, receivedSenderSeqNo, isSessionInit, message)
        } catch (e: Exception) {
            log.error("Unable to process message, ignoring it: $message", e)
            return null
        }
    }

    private inline fun <T> ClientMessage.required(key: SimpleString, extractor: ClientMessage.(SimpleString) -> T): T {
        require(containsProperty(key)) { "Missing $key" }
        return extractor(key)
    }

    private class ArtemisReceivedMessage(override val topic: String,
                                         override val peer: CordaX500Name,
                                         override val platformVersion: Int,
                                         override val uniqueMessageId: DeduplicationId,
                                         override val senderUUID: String?,
                                         override val senderSeqNo: Long?,
                                         override val isSessionInit: Boolean,
                                         private val message: ClientMessage) : ReceivedMessage {
        override val data: ByteSequence by lazy { OpaqueBytes(ByteArray(message.bodySize).apply { message.bodyBuffer.readBytes(this) }) }
        override val debugTimestamp: Instant get() = Instant.ofEpochMilli(message.timestamp)
        override val additionalHeaders: Map<String, String> = emptyMap()
        override fun toString() = "$topic#$data"
    }

    internal fun deliver(artemisMessage: ClientMessage) {
        artemisToCordaMessage(artemisMessage)?.let { cordaMessage ->
            val outcome = deduplicator.checkDuplicate(cordaMessage)
            if (outcome == P2PMessageDeduplicator.Outcome.NEW) {
                deduplicator.signalMessageProcessStart(cordaMessage)
                deliver(cordaMessage, artemisMessage)
            } else if (outcome == P2PMessageDeduplicator.Outcome.DUPLICATE) {
                log.debug { "Acknowledge duplicate message id: ${cordaMessage.uniqueMessageId} senderUUID: ${cordaMessage.senderUUID} senderSeqNo: ${cordaMessage.senderSeqNo} isSessionInit: ${cordaMessage.isSessionInit}" }
                messagingExecutor!!.acknowledge(artemisMessage)
            } else {
                log.debug { "Discard in-flight message id: ${cordaMessage.uniqueMessageId} senderUUID: ${cordaMessage.senderUUID} senderSeqNo: ${cordaMessage.senderSeqNo} isSessionInit: ${cordaMessage.isSessionInit}" }
            }
        }
    }

    private fun deliver(msg: ReceivedMessage, artemisMessage: ClientMessage) {
        state.checkNotLocked()
        val deliverTo = handlers[msg.topic]
        if (deliverTo != null) {
            try {
                deliverTo(msg, HandlerRegistration(msg.topic, deliverTo), MessageDeduplicationHandler(artemisMessage, msg))
            } catch (e: Exception) {
                log.error("Caught exception whilst executing message handler for ${msg.topic}", e)
            }
        } else {
            log.warn("Received message ${msg.uniqueMessageId} for ${msg.topic} that doesn't have any registered handlers yet")
        }
    }

    private inner class MessageDeduplicationHandler(val artemisMessage: ClientMessage, override val receivedMessage: ReceivedMessage) : DeduplicationHandler, ExternalEvent.ExternalMessageEvent {
        override val externalCause: ExternalEvent
            get() = this
        override val flowId: StateMachineRunId by lazy { StateMachineRunId.createRandom() }
        override val deduplicationHandler: MessageDeduplicationHandler
            get() = this

        override fun insideDatabaseTransaction() {
            deduplicator.persistDeduplicationId(receivedMessage.uniqueMessageId)
        }

        override fun afterDatabaseTransaction() {
            deduplicator.signalMessageProcessFinish(receivedMessage.uniqueMessageId)
            messagingExecutor!!.acknowledge(artemisMessage)
        }

        override fun toString(): String {
            return "${javaClass.simpleName}(${receivedMessage.uniqueMessageId})"
        }
    }

    /**
     * Initiates shutdown: if called from a thread that isn't controlled by the executor passed to the constructor
     * then this will block until all in-flight messages have finished being handled and acknowledged. If called
     * from a thread that's a part of the [net.corda.node.utilities.AffinityExecutor] given to the constructor,
     * it returns immediately and shutdown is asynchronous.
     */
    override fun stop() {
        val running = state.locked {
            // We allow stop() to be called without a run() in between, but it must have at least been started.
            check(started)
            val prevRunning = running
            running = false
            stateHelper.active = false
            networkChangeSubscription?.unsubscribe()
            require(p2pConsumer != null) { "stop can't be called twice" }
            require(producer != null) { "stop can't be called twice" }

            close(p2pConsumer)
            p2pConsumer = null

            close(producer)
            producer = null
            producerSession!!.commit()

            close(executorProducer)
            executorProducer = null
            executorSession!!.commit()

            close(bridgeNotifyConsumer)
            knownQueues.clear()
            eventsSubscription?.unsubscribe()
            eventsSubscription = null
            prevRunning
        }
        synchronized(handlersChangedSignal) {
            handlersChangedSignal.notifyAll()
        }
        if (running && !nodeExecutor.isOnThread) {
            // Wait for the main loop to notice the consumer has gone and finish up.
            shutdownLatch.await()
        }

        // Only first caller to gets running true to protect against double stop, which seems to happen in some integration tests.
        state.locked {
            sessionFactory?.close()
            locator?.close()
        }
    }

    private fun close(target: AutoCloseable?) {
        try {
            target?.close()
        } catch (ignored: ActiveMQObjectClosedException) {
            // swallow
        }
    }

    @Suspendable
    override fun send(message: Message, target: MessageRecipients, sequenceKey: Any) {
        requireMessageSize(message.data.size, maxMessageSize)
        messagingExecutor!!.send(message, target)
    }

    @Suspendable
    override fun sendAll(addressedMessages: List<MessagingService.AddressedMessage>) {
        for ((message, target, sequenceKey) in addressedMessages) {
            send(message, target, sequenceKey)
        }
    }

    override fun resolveTargetToArtemisQueue(address: MessageRecipients): String {
        return if (address == myAddress) {
            // If we are sending to ourselves then route the message directly to our P2P queue.
            RemoteInboxAddress(myIdentity).queueName
        } else {
            // Otherwise, we send the message to an internal queue for the target residing on our broker. It's then the
            // broker's job to route the message to the target's P2P queue.
            val internalTargetQueue = (address as? ArtemisAddress)?.queueName
                    ?: throw IllegalArgumentException("Not an Artemis address")
            state.locked {
                val serviceAddress = address is ServiceAddress
                createQueueIfAbsent(internalTargetQueue, producerSession!!, exclusive = !serviceAddress, isServiceAddress = serviceAddress)
            }
            internalTargetQueue
        }
    }

    /** Attempts to create a durable queue on the broker which is bound to an address of the same name. */
    private fun createQueueIfAbsent(queueName: String, session: ClientSession, exclusive: Boolean, isServiceAddress: Boolean) {
        fun sendBridgeCreateMessage() {
            val keyHash = queueName.substring(PEERS_PREFIX.length)
            val peers = networkMap.getNodesByOwningKeyIndex(keyHash)
            for (node in peers) {
                val bridge = BridgeEntry(queueName, node.addresses, node.legalIdentities.map { it.name }, isServiceAddress)
                val createBridgeMessage = BridgeControl.Create(config.myLegalName.toString(), bridge)
                sendBridgeControl(createBridgeMessage)
            }
        }
        if (!knownQueues.contains(queueName)) {
            if (delayStartQueues.contains(queueName)) {
                log.info("Start bridge for previously empty queue $queueName")
                sendBridgeCreateMessage()
                delayStartQueues -= queueName
            } else {
                val queueQuery = session.queueQuery(SimpleString(queueName))
                if (!queueQuery.isExists) {
                    log.info("Create fresh queue $queueName bound on same address")
                    session.createQueue(QueueConfiguration(queueName).setRoutingType(RoutingType.ANYCAST).setAddress(queueName)
                            .setDurable(true).setAutoCreated(false)
                            .setMaxConsumers(ActiveMQDefaultConfiguration.getDefaultMaxQueueConsumers())
                            .setPurgeOnNoConsumers(ActiveMQDefaultConfiguration.getDefaultPurgeOnNoConsumers())
                            .setExclusive(exclusive)
                            .setLastValue(null)
                    )
                    sendBridgeCreateMessage()
                }
            }
            knownQueues += queueName
        }
    }

    override fun addMessageHandler(topic: String, callback: MessageHandler): MessageHandlerRegistration {
        require(topic.isNotBlank()) { "Topic must not be blank, as the empty topic is a special case." }
        handlers.compute(topic) { _, handler ->
            if (handler != null) {
                throw IllegalStateException("Cannot add another acking handler for $topic, there is already an acking one")
            }
            callback
        }
        synchronized(handlersChangedSignal) {
            handlersChangedSignal.notifyAll()
        }
        return HandlerRegistration(topic, callback)
    }

    override fun removeMessageHandler(registration: MessageHandlerRegistration) {
        registration as HandlerRegistration
        handlers.remove(registration.topic)
    }

    override fun createMessage(topic: String, data: ByteArray, deduplicationId: SenderDeduplicationId, additionalHeaders: Map<String, String>): Message {
        return NodeClientMessage(topic, OpaqueBytes(data), deduplicationId.deduplicationId, deduplicationId.senderUUID, additionalHeaders)
    }

    override fun getAddressOfParty(partyInfo: PartyInfo): MessageRecipients {
        return when (partyInfo) {
            is PartyInfo.SingleNode -> NodeAddress(partyInfo.party.owningKey)
            is PartyInfo.DistributedNode -> ServiceAddress(partyInfo.party.owningKey)
        }
    }
}

private class P2PMessagingConsumer(
        queueNames: Set<String>,
        createSession: () -> ClientSession,
        private val isDrainingModeOn: () -> Boolean,
        private val drainingModeWasChangedEvents: Observable<Pair<Boolean, Boolean>>,
        private val metricsRegistry : MetricRegistry) : LifecycleSupport {

    private companion object {
        private const val initialSessionMessages = "${P2PMessagingHeaders.Type.KEY}<>'${P2PMessagingHeaders.Type.SESSION_INIT_VALUE}'"
        private val logger by lazy { loggerFor<P2PMessagingClient>() }
    }

    private var startedFlag = false

    val messages: PublishSubject<ClientMessage> = PublishSubject.create<ClientMessage>()

    private val existingOnlyConsumer = multiplex(queueNames, createSession, initialSessionMessages)
    private val initialAndExistingConsumer = multiplex(queueNames, createSession)
    private val subscriptions = mutableSetOf<Subscription>()

    private var notificationTimer : Timer? = null
    private fun scheduleDrainNotificationTimer() {
        notificationTimer =  timer("DrainNotificationTimer", true, 10.seconds.toMillis(), 1.minutes.toMillis()) {
            logger.warn("Node is currently in draining mode, new flows will not be processed! Flows in flight: ${metricsRegistry.gauges["Flows.InFlight"]?.value}")
        }
    }

    override fun start() {

        synchronized(this) {
            require(!startedFlag){"Must not already be started"}
            drainingModeWasChangedEvents.filter { change -> change.switchedOn() }.doOnNext {
                initialAndExistingConsumer.switchTo(existingOnlyConsumer)
                scheduleDrainNotificationTimer()
            }.subscribe()
            drainingModeWasChangedEvents.filter { change -> change.switchedOff() }.doOnNext {
                existingOnlyConsumer.switchTo(initialAndExistingConsumer)
                notificationTimer?.cancel()
            }.subscribe()
            subscriptions += existingOnlyConsumer.messages.doOnNext(messages::onNext).subscribe()
            subscriptions += initialAndExistingConsumer.messages.doOnNext(messages::onNext).subscribe()
            if (isDrainingModeOn()) {
                existingOnlyConsumer.start()
                scheduleDrainNotificationTimer()
            } else {
                initialAndExistingConsumer.start()
            }
            startedFlag = true
        }
    }

    override fun stop() {

        synchronized(this) {
            if (startedFlag) {
                existingOnlyConsumer.stop()
                initialAndExistingConsumer.stop()
                subscriptions.forEach(Subscription::unsubscribe)
                subscriptions.clear()
                startedFlag = false
            }
            messages.onCompleted()
        }
    }

    override val started: Boolean
        get() = startedFlag

    private fun Pair<Boolean, Boolean>.switchedOff() = first && !second

    private fun Pair<Boolean, Boolean>.switchedOn() = !first && second
}

private fun ReactiveArtemisConsumer.switchTo(other: ReactiveArtemisConsumer) {

    disconnect()
    when {
        !other.started -> other.start()
        !other.connected -> other.connect()
    }
}
