package net.corda.node.flows

import co.paralleluniverse.fibers.Suspendable
import net.corda.core.flows.FlowLogic
import net.corda.core.flows.FlowSession
import net.corda.core.flows.HospitalizeFlowException
import net.corda.core.flows.InitiatedBy
import net.corda.core.flows.InitiatingFlow
import net.corda.core.flows.StartableByRPC
import net.corda.core.flows.StateMachineRunId
import net.corda.core.identity.Party
import net.corda.core.internal.FlowIORequest
import net.corda.core.internal.IdempotentFlow
import net.corda.core.internal.TimedFlow
import net.corda.core.internal.concurrent.transpose
import net.corda.core.internal.mapToSet
import net.corda.core.messaging.StateMachineTransactionMapping
import net.corda.core.messaging.startFlow
import net.corda.core.utilities.OpaqueBytes
import net.corda.core.utilities.getOrThrow
import net.corda.core.utilities.seconds
import net.corda.core.utilities.unwrap
import net.corda.finance.DOLLARS
import net.corda.finance.flows.CashIssueAndPaymentFlow
import net.corda.node.services.config.NodeConfiguration
import net.corda.node.services.statemachine.FlowStateMachineImpl
import net.corda.node.services.statemachine.FlowTimeoutException
import net.corda.node.services.statemachine.StaffedFlowHospital
import net.corda.testing.core.ALICE_NAME
import net.corda.testing.core.BOB_NAME
import net.corda.testing.core.singleIdentity
import net.corda.testing.driver.DriverParameters
import net.corda.testing.driver.driver
import net.corda.testing.node.internal.FINANCE_CORDAPPS
import net.corda.testing.node.internal.enclosedCordapp
import org.junit.Test
import java.sql.SQLTransientConnectionException
import java.time.Duration
import java.util.concurrent.ConcurrentLinkedQueue
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import kotlin.test.assertEquals
import kotlin.test.assertTrue

class FlowReloadAfterCheckpointTest {

    private companion object {
        private val DEFAULT_TIMEOUT = Duration.ofSeconds(10)
        val cordapps = listOf(enclosedCordapp())
    }

    @Test(timeout = 300_000)
    fun `flow will reload from its checkpoint after suspending when reloadCheckpointAfterSuspend is true`() {
        val reloads = ConcurrentLinkedQueue<StateMachineRunId>()
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { id ->
            reloads.add(id)
        }
        driver(DriverParameters(startNodesInProcess = true, notarySpecs = emptyList(), cordappsForAllNodes = cordapps)) {

            val (alice, bob) = listOf(ALICE_NAME, BOB_NAME)
                .map {
                    startNode(
                        providedName = it,
                        customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
                    )
                }
                .transpose()
                .getOrThrow()

            val handle = alice.rpc.startFlow(::ReloadFromCheckpointFlow, bob.nodeInfo.singleIdentity(), false, false, false)
            val flowStartedByAlice = handle.id
            handle.returnValue.getOrThrow()
            assertEquals(5, reloads.count { it == flowStartedByAlice })
            assertEquals(6, reloads.count { it == ReloadFromCheckpointResponder.flowId })
        }
    }

    @Test(timeout = 300_000)
    fun `flow will not reload from its checkpoint after suspending when reloadCheckpointAfterSuspend is false`() {
        val reloads = ConcurrentLinkedQueue<StateMachineRunId>()
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { id ->
            reloads.add(id)
        }
        driver(DriverParameters(startNodesInProcess = true, notarySpecs = emptyList(), cordappsForAllNodes = cordapps)) {

            val (alice, bob) = listOf(ALICE_NAME, BOB_NAME)
                .map {
                    startNode(
                        providedName = it,
                        customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to false)
                    )
                }
                .transpose()
                .getOrThrow()

            val handle = alice.rpc.startFlow(::ReloadFromCheckpointFlow, bob.nodeInfo.singleIdentity(), false, false, false)
            handle.returnValue.getOrThrow()
            assertEquals(0, reloads.size)
        }
    }

    @Test(timeout = 300_000)
    fun `flow will reload from its checkpoint after suspending when reloadCheckpointAfterSuspend is true and be kept for observation due to failed deserialization`() {
        val observations = QueueWithCountdown<StateMachineRunId>(1)
        val reloads = QueueWithCountdown<StateMachineRunId>(8)
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { id ->
            reloads.add(id)
        }
        StaffedFlowHospital.onFlowKeptForOvernightObservation.add { id, _ ->
            observations.add(id)
        }
        driver(DriverParameters(startNodesInProcess = true, notarySpecs = emptyList(), cordappsForAllNodes = cordapps)) {

            val (alice, bob) = listOf(ALICE_NAME, BOB_NAME)
                .map {
                    startNode(
                        providedName = it,
                        customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
                    )
                }
                .transpose()
                .getOrThrow()

            val handle = alice.rpc.startFlow(::ReloadFromCheckpointFlow, bob.nodeInfo.singleIdentity(), true, false, false)
            val flowStartedByAlice: StateMachineRunId = handle.id

            // We can't wait on the flow ending, because it breaks, so we need to wait on internal status changes instead
            observations.await(DEFAULT_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)
            reloads.await(DEFAULT_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)
            assertEquals(flowStartedByAlice, observations.singleOrNull())
            assertEquals(4, reloads.count { it == flowStartedByAlice })
            assertEquals(4, reloads.count { it == ReloadFromCheckpointResponder.flowId })
        }
    }

    @Test(timeout = 300_000)
    fun `flow will reload from a previous checkpoint after calling suspending function and skipping the persisting the current checkpoint when reloadCheckpointAfterSuspend is true`() {
        val reloads = QueueWithCountdown<StateMachineRunId>(11)
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { id ->
            reloads.add(id)
        }
        driver(DriverParameters(startNodesInProcess = true, notarySpecs = emptyList(), cordappsForAllNodes = cordapps)) {

            val (alice, bob) = listOf(ALICE_NAME, BOB_NAME)
                .map {
                    startNode(
                        providedName = it,
                        customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
                    )
                }
                .transpose()
                .getOrThrow()

            val handle = alice.rpc.startFlow(::ReloadFromCheckpointFlow, bob.nodeInfo.singleIdentity(), false, false, true)
            val flowStartedByAlice = handle.id
            handle.returnValue.getOrThrow()
            reloads.await(DEFAULT_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)
            assertEquals(5, reloads.count { it == flowStartedByAlice })
            assertEquals(6, reloads.count { it == ReloadFromCheckpointResponder.flowId })
        }
    }

    @Test(timeout = 300_000)
    fun `idempotent flow will reload from initial checkpoint after calling a suspending function when reloadCheckpointAfterSuspend is true`() {
        var reloadCount = 0
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { _ -> reloadCount += 1 }
        driver(DriverParameters(startNodesInProcess = true, notarySpecs = emptyList(), cordappsForAllNodes = cordapps)) {

            val alice = startNode(
                providedName = ALICE_NAME,
                customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
            ).getOrThrow()

            alice.rpc.startFlow(::MyIdempotentFlow, false).returnValue.getOrThrow()
            assertEquals(5, reloadCount)
        }
    }

    @Test(timeout = 300_000)
    fun `idempotent flow will reload from initial checkpoint after calling a suspending function when reloadCheckpointAfterSuspend is true but can't throw deserialization error from objects in the call function`() {
        var reloadCount = 0
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { _ -> reloadCount += 1 }
        driver(DriverParameters(startNodesInProcess = true, notarySpecs = emptyList(), cordappsForAllNodes = cordapps)) {

            val alice = startNode(
                providedName = ALICE_NAME,
                customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
            ).getOrThrow()

            alice.rpc.startFlow(::MyIdempotentFlow, true).returnValue.getOrThrow()
            assertEquals(5, reloadCount)
        }
    }

    @Test(timeout = 300_000)
    fun `timed flow will reload from initial checkpoint after calling a suspending function when reloadCheckpointAfterSuspend is true`() {
        val reloads = ConcurrentLinkedQueue<StateMachineRunId>()
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { runId -> reloads.add(runId) }
        driver(DriverParameters(startNodesInProcess = true, notarySpecs = emptyList(), cordappsForAllNodes = cordapps)) {

            val alice = startNode(
                providedName = ALICE_NAME,
                customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
            ).getOrThrow()

            alice.rpc.startFlow(::MyTimedFlow).returnValue.getOrThrow()
            assertEquals(5, reloads.size)
        }
    }

    @Test(timeout = 300_000)
    fun `flow will correctly retry after an error when reloadCheckpointAfterSuspend is true`() {
        val reloads = ConcurrentLinkedQueue<StateMachineRunId>()
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { runId -> reloads.add(runId) }
        var timesDischarged = 0
        StaffedFlowHospital.onFlowDischarged.add { _, _ -> timesDischarged += 1 }
        driver(DriverParameters(startNodesInProcess = true, notarySpecs = emptyList(), cordappsForAllNodes = cordapps)) {

            val alice = startNode(
                providedName = ALICE_NAME,
                customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
            ).getOrThrow()

            alice.rpc.startFlow(::TransientConnectionFailureFlow).returnValue.getOrThrow()
            assertEquals(5, reloads.size)
            assertEquals(3, timesDischarged)
        }
    }

    @Test(timeout = 300_000)
    fun `flow continues reloading from checkpoints after node restart when reloadCheckpointAfterSuspend is true`() {
        val reloads = QueueWithCountdown<StateMachineRunId>(5)
        val firstLatch = CountDownLatch(2)
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { runId ->
            reloads.add(runId)
            firstLatch.countDown()
        }
        driver(
            DriverParameters(
                inMemoryDB = false,
                startNodesInProcess = true,
                notarySpecs = emptyList(),
                cordappsForAllNodes = cordapps
            )
        ) {

            val alice = startNode(
                providedName = ALICE_NAME,
                customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
            ).getOrThrow()

            alice.rpc.startFlow(::MyHospitalizingFlow)
            assertTrue { firstLatch.await(10, TimeUnit.SECONDS) }
            alice.stop()
            assertEquals(2, reloads.size)

            // Set up a new latch
            startNode(
                providedName = ALICE_NAME,
                customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
            ).getOrThrow()

            reloads.await(DEFAULT_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)
            assertEquals(5, reloads.size)
        }
    }

    @Test(timeout = 300_000)
    fun `idempotent flow continues reloading from checkpoints after node restart when reloadCheckpointAfterSuspend is true`() {
        // restarts completely from the beginning and forgets the in-memory reload count therefore
        // it reloads an extra 2 times for checkpoints it had already reloaded before the node shutdown
        val reloads = QueueWithCountdown<StateMachineRunId>(7)
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { runId ->
            reloads.add(runId)
        }
        driver(
            DriverParameters(
                inMemoryDB = false,
                startNodesInProcess = true,
                notarySpecs = emptyList(),
                cordappsForAllNodes = cordapps
            )
        ) {

            val alice = startNode(
                providedName = ALICE_NAME,
                customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
            ).getOrThrow()

            alice.rpc.startFlow(::IdempotentHospitalizingFlow)
            Thread.sleep(10.seconds.toMillis())

            alice.stop()

            startNode(
                providedName = ALICE_NAME,
                customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
            ).getOrThrow()

            // restarts completely from the beginning and forgets the in-memory reload count therefore
            // it reloads an extra 2 times for checkpoints it had already reloaded before the node shutdown
            reloads.await(DEFAULT_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)
            assertEquals(7, reloads.size)
        }
    }

    @Test(timeout = 300_000)
    fun `more complicated flow will reload from its checkpoint after suspending when reloadCheckpointAfterSuspend is true`() {
        val reloads = QueueWithCountdown<StateMachineRunId>(13)
        FlowStateMachineImpl.onReloadFlowFromCheckpoint = { id ->
            reloads.add(id)
        }
        driver(DriverParameters(startNodesInProcess = true, cordappsForAllNodes = FINANCE_CORDAPPS)) {

            val (alice, bob) = listOf(ALICE_NAME, BOB_NAME)
                .map {
                    startNode(
                        providedName = it,
                        customOverrides = mapOf(NodeConfiguration::reloadCheckpointAfterSuspend.name to true)
                    )
                }
                .transpose()
                .getOrThrow()

            val handle = alice.rpc.startFlow(
                ::CashIssueAndPaymentFlow,
                500.DOLLARS,
                OpaqueBytes.of(0x01),
                bob.nodeInfo.singleIdentity(),
                false,
                defaultNotaryIdentity
            )
            val flowStartedByAlice = handle.id
            handle.returnValue.getOrThrow(30.seconds)
            val flowStartedByBob = bob.rpc.stateMachineRecordedTransactionMappingSnapshot()
                .mapToSet(StateMachineTransactionMapping::stateMachineRunId)
                .single()
            reloads.await(DEFAULT_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)
            assertEquals(8, reloads.filter { it == flowStartedByAlice }.size)
            assertEquals(7, reloads.filter { it == flowStartedByBob }.size)
        }
    }

    /**
     * Has 4 suspension points inside the flow and 1 in [FlowStateMachineImpl.run] totaling 5.
     * Therefore this flow should reload 5 times when completed without errors or restarts.
     */
    @StartableByRPC
    @InitiatingFlow
    class ReloadFromCheckpointFlow(
        private val party: Party,
        private val shouldHaveDeserializationError: Boolean,
        private val counterPartyHasDeserializationError: Boolean,
        private val skipCheckpoints: Boolean
    ) : FlowLogic<Unit>() {

        @Suspendable
        override fun call() {
            val session = initiateFlow(party)
            session.send(counterPartyHasDeserializationError, skipCheckpoints)
            session.receive(String::class.java, skipCheckpoints).unwrap { it }
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, skipCheckpoints)
            val map = if (shouldHaveDeserializationError) {
                BrokenMap(mutableMapOf("i dont want" to "this to work"))
            } else {
                mapOf("i dont want" to "this to work")
            }
            logger.info("I need to use my variable to pass the build!: $map")
            session.sendAndReceive<String>("hey I made it this far")
        }
    }

    /**
     * Has 5 suspension points inside the flow and 1 in [FlowStateMachineImpl.run] totaling 6.
     * Therefore this flow should reload 6 times when completed without errors or restarts.
     */
    @InitiatedBy(ReloadFromCheckpointFlow::class)
    class ReloadFromCheckpointResponder(private val session: FlowSession) : FlowLogic<Unit>() {

        companion object {
            var flowId: StateMachineRunId? = null
        }

        @Suspendable
        override fun call() {
            flowId = runId
            val counterPartyHasDeserializationError = session.receive<Boolean>().unwrap { it }
            session.send("hello there 12312311")
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            val map = if (counterPartyHasDeserializationError) {
                BrokenMap(mutableMapOf("i dont want" to "this to work"))
            } else {
                mapOf("i dont want" to "this to work")
            }
            logger.info("I need to use my variable to pass the build!: $map")
            session.receive<String>().unwrap { it }
            session.send("sending back a message")
        }
    }

    /**
     * Has 4 suspension points inside the flow and 1 in [FlowStateMachineImpl.run] totaling 5.
     * Therefore this flow should reload 5 times when completed without errors or restarts.
     */
    @StartableByRPC
    @InitiatingFlow
    class MyIdempotentFlow(private val shouldHaveDeserializationError: Boolean) : FlowLogic<Unit>(), IdempotentFlow {

        @Suspendable
        override fun call() {
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            val map = if (shouldHaveDeserializationError) {
                BrokenMap(mutableMapOf("i dont want" to "this to work"))
            } else {
                mapOf("i dont want" to "this to work")
            }
            logger.info("I need to use my variable to pass the build!: $map")
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
        }
    }

    /**
     * Has 4 suspension points inside the flow and 1 in [FlowStateMachineImpl.run] totaling 5.
     * Therefore this flow should reload 5 times when completed without errors or restarts.
     */
    @StartableByRPC
    @InitiatingFlow
    class MyTimedFlow : FlowLogic<Unit>(), TimedFlow {

        companion object {
            var thrown = false
        }

        override val isTimeoutEnabled: Boolean = true

        @Suspendable
        override fun call() {
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            if (!thrown) {
                thrown = true
                throw FlowTimeoutException()
            }
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
        }
    }

    @StartableByRPC
    @InitiatingFlow
    class TransientConnectionFailureFlow : FlowLogic<Unit>() {

        companion object {
            var retryCount = 0
        }

        @Suspendable
        override fun call() {
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            if (retryCount < 3) {
                retryCount += 1
                throw SQLTransientConnectionException("Connection is not available")

            }
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
        }
    }

    /**
     * Has 4 suspension points inside the flow and 1 in [FlowStateMachineImpl.run] totaling 5.
     * Therefore this flow should reload 5 times when completed without errors or restarts.
     */
    @StartableByRPC
    @InitiatingFlow
    class MyHospitalizingFlow : FlowLogic<Unit>() {

        companion object {
            var thrown = false
        }

        @Suspendable
        override fun call() {
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            if (!thrown) {
                thrown = true
                throw HospitalizeFlowException("i want to try again")
            }
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
        }
    }

    /**
     * Has 4 suspension points inside the flow and 1 in [FlowStateMachineImpl.run] totaling 5.
     * Therefore this flow should reload 5 times when completed without errors or restarts.
     */
    @StartableByRPC
    @InitiatingFlow
    class IdempotentHospitalizingFlow : FlowLogic<Unit>(), IdempotentFlow {

        companion object {
            var thrown = false
        }

        @Suspendable
        override fun call() {
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            if (!thrown) {
                thrown = true
                throw HospitalizeFlowException("i want to try again")
            }
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
            stateMachine.suspend(FlowIORequest.ForceCheckpoint, false)
        }
    }
}

internal class BrokenMap<K, V>(delegate: MutableMap<K, V> = mutableMapOf()) : MutableMap<K, V> by delegate {
    override fun put(key: K, value: V): V = throw IllegalStateException("Broken on purpose")
}

