package com.r3corda.demos

import co.paralleluniverse.fibers.Suspendable
import com.google.common.net.HostAndPort
import com.google.common.util.concurrent.ListenableFuture
import com.r3corda.contracts.CommercialPaper
import com.r3corda.contracts.asset.DUMMY_CASH_ISSUER
import com.r3corda.contracts.asset.cashBalances
import com.r3corda.contracts.testing.fillWithSomeTestCash
import com.r3corda.core.contracts.*
import com.r3corda.core.crypto.Party
import com.r3corda.core.crypto.SecureHash
import com.r3corda.core.crypto.generateKeyPair
import com.r3corda.core.days
import com.r3corda.core.logElapsedTime
import com.r3corda.core.node.NodeInfo
import com.r3corda.core.node.services.ServiceInfo
import com.r3corda.core.protocols.ProtocolLogic
import com.r3corda.core.seconds
import com.r3corda.core.success
import com.r3corda.core.transactions.SignedTransaction
import com.r3corda.core.utilities.Emoji
import com.r3corda.core.utilities.LogHelper
import com.r3corda.core.utilities.ProgressTracker
import com.r3corda.node.internal.Node
import com.r3corda.node.services.config.ConfigHelper
import com.r3corda.node.services.config.FullNodeConfiguration
import com.r3corda.node.services.messaging.NodeMessagingClient
import com.r3corda.node.services.network.NetworkMapService
import com.r3corda.node.services.persistence.NodeAttachmentService
import com.r3corda.node.services.transactions.ValidatingNotaryService
import com.r3corda.node.utilities.databaseTransaction
import com.r3corda.protocols.NotaryProtocol
import com.r3corda.protocols.TwoPartyTradeProtocol
import joptsimple.OptionParser
import org.slf4j.Logger
import org.slf4j.LoggerFactory
import java.nio.file.Path
import java.nio.file.Paths
import java.security.PublicKey
import java.time.Instant
import java.util.*
import kotlin.concurrent.thread
import kotlin.system.exitProcess
import kotlin.test.assertEquals

// TRADING DEMO
//
// Please see docs/build/html/running-the-demos.html
//
// This program is a simple driver for exercising the two party trading protocol. Until Corda has a unified node server
// programs like this are required to wire up the pieces and run a demo scenario end to end.
//
// If you are creating a new scenario, you can use this program as a template for creating your own driver. Make sure to
// copy/paste the right parts of the build.gradle file to make sure it gets a script to run it deposited in
// build/install/r3prototyping/bin
//
// In this scenario, a buyer wants to purchase some commercial paper by swapping his cash for the CP. The seller learns
// that the buyer exists, and sends them a message to kick off the trade. The seller, having obtained his CP, then quits
// and the buyer goes back to waiting. The buyer will sell as much CP as he can!
//
// The different roles in the scenario this program can adopt are:

enum class Role {
    BUYER,
    SELLER
}

// And this is the directory under the current working directory where each node will create its own server directory,
// which holds things like checkpoints, keys, databases, message logs etc.
val DEFAULT_BASE_DIRECTORY = "./build/trader-demo"

private val log: Logger = LoggerFactory.getLogger("TraderDemo")

fun main(args: Array<String>) {
    val parser = OptionParser()

    val roleArg = parser.accepts("role").withRequiredArg().ofType(Role::class.java).required()
    val myNetworkAddress = parser.accepts("network-address").withRequiredArg().defaultsTo("localhost")
    val theirNetworkAddress = parser.accepts("other-network-address").withRequiredArg().defaultsTo("localhost")
    val apiNetworkAddress = parser.accepts("api-address").withRequiredArg().defaultsTo("localhost")
    val baseDirectoryArg = parser.accepts("base-directory").withRequiredArg().defaultsTo(DEFAULT_BASE_DIRECTORY)
    val h2PortArg = parser.accepts("h2-port").withRequiredArg().ofType(Int::class.java).defaultsTo(-1)
    val options = try {
        parser.parse(*args)
    } catch (e: Exception) {
        log.error(e.message)
        printHelp(parser)
        exitProcess(1)
    }

    val role = options.valueOf(roleArg)!!

    val myNetAddr = HostAndPort.fromString(options.valueOf(myNetworkAddress)).withDefaultPort(
            when (role) {
                Role.BUYER -> 31337
                Role.SELLER -> 31340
            }
    )
    val theirNetAddr = HostAndPort.fromString(options.valueOf(theirNetworkAddress)).withDefaultPort(
            when (role) {
                Role.BUYER -> 31340
                Role.SELLER -> 31337
            }
    )
    val apiNetAddr = HostAndPort.fromString(options.valueOf(apiNetworkAddress)).withDefaultPort(myNetAddr.port + 1)
    val h2Port = if (options.valueOf(h2PortArg) < 0) {
        myNetAddr.port + 2
    } else options.valueOf(h2PortArg)

    val baseDirectory = options.valueOf(baseDirectoryArg)!!

    // Suppress the Artemis MQ noise, and activate the demo logging.
    //
    // The first two strings correspond to the first argument to StateMachineManager.add() but the way we handle logging
    // for protocols will change in future.
    LogHelper.setLevel("+demo.buyer", "+demo.seller", "-org.apache.activemq")

    val directory = Paths.get(baseDirectory, role.name.toLowerCase())
    log.info("Using base demo directory $directory")

    // Override the default config file (which you can find in the file "reference.conf") to give each node a name.
    val config = run {
        val myLegalName = when (role) {
            Role.BUYER -> "Bank A"
            Role.SELLER -> "Bank B"
        }
        val configOverrides = mapOf(
                "myLegalName" to myLegalName,
                "artemisAddress" to myNetAddr.toString(),
                "webAddress" to apiNetAddr.toString(),
                "h2port" to h2Port.toString()
        )
        FullNodeConfiguration(ConfigHelper.loadConfig(directory, allowMissingConfig = true, configOverrides = configOverrides))
    }

    // Which services will this instance of the node provide to the network?
    val advertisedServices: Set<ServiceInfo>

    // One of the two servers needs to run the network map and notary services. In such a trivial two-node network
    // the map is not very helpful, but we need one anyway. So just make the buyer side run the network map as it's
    // the side that sticks around waiting for the seller.
    val networkMapId = if (role == Role.BUYER) {
        advertisedServices = setOf(ServiceInfo(NetworkMapService.type), ServiceInfo(ValidatingNotaryService.type))
        null
    } else {
        advertisedServices = emptySet()
        NodeMessagingClient.makeNetworkMapAddress(theirNetAddr)
    }

    // And now construct then start the node object. It takes a little while.
    val node = logElapsedTime("Node startup", log) {
        Node(config, networkMapId, advertisedServices).setup().start()
    }

    // What happens next depends on the role. The buyer sits around waiting for a trade to start. The seller role
    // will contact the buyer and actually make something happen.
    val amount = 1000.DOLLARS
    if (role == Role.BUYER) {
        runBuyer(node, amount)
    } else {
        node.networkMapRegistrationFuture.success {
            val party = node.netMapCache.getNodeByLegalName("Bank A")?.legalIdentity ?: throw IllegalStateException("Cannot find other node?!")
            runSeller(node, amount, party)
        }
    }

    node.run()
}

private fun runSeller(node: Node, amount: Amount<Currency>, otherSide: Party) {
    // The seller will sell some commercial paper to the buyer, who will pay with (self issued) cash.
    //
    // The CP sale transaction comes with a prospectus PDF, which will tag along for the ride in an
    // attachment. Make sure we have the transaction prospectus attachment loaded into our store.
    //
    // This can also be done via an HTTP upload, but here we short-circuit and do it from code.
    if (node.storage.attachments.openAttachment(TraderDemoProtocolSeller.PROSPECTUS_HASH) == null) {
        TraderDemoProtocolSeller::class.java.getResourceAsStream("bank-of-london-cp.jar").use {
            val id = node.storage.attachments.importAttachment(it)
            assertEquals(TraderDemoProtocolSeller.PROSPECTUS_HASH, id)
        }
    }

    val tradeTX: ListenableFuture<SignedTransaction>
    if (node.isPreviousCheckpointsPresent) {
        tradeTX = node.smm.findStateMachines(TraderDemoProtocolSeller::class.java).single().second
    } else {
        val seller = TraderDemoProtocolSeller(otherSide, amount)
        tradeTX = node.services.startProtocol(seller)
    }

    tradeTX.success {
        log.info("Sale completed - we have a happy customer!\n\nFinal transaction is:\n\n${Emoji.renderIfSupported(it.tx)}")
        thread {
            node.stop()
        }
    }
}

private fun runBuyer(node: Node, amount: Amount<Currency>) {
    // Buyer will fetch the attachment from the seller automatically when it resolves the transaction.
    // For demo purposes just extract attachment jars when saved to disk, so the user can explore them.
    val attachmentsPath = (node.storage.attachments as NodeAttachmentService).let {
        it.automaticallyExtractAttachments = true
        it.storePath
    }

    // Self issue some cash.
    //
    // TODO: At some point this demo should be extended to have a central bank node.
    databaseTransaction(node.database) {
        node.services.fillWithSomeTestCash(300000.DOLLARS,
                outputNotary = node.info.notaryIdentity, // In this demo, the buyer and notary are on the same node, but need to use right key.
                ownedBy = node.info.legalIdentity.owningKey)
    }

    // Wait around until a node asks to start a trade with us. In a real system, this part would happen out of band
    // via some other system like an exchange or maybe even a manual messaging system like Bloomberg. But for the
    // next stage in our building site, we will just auto-generate fake trades to give our nodes something to do.
    //
    // As the seller initiates the two-party trade protocol, here, we will be the buyer.
    node.services.registerProtocolInitiator(TraderDemoProtocolSeller::class) { otherParty ->
        TraderDemoProtocolBuyer(otherParty, attachmentsPath, amount)
    }
}

// We create a couple of ad-hoc test protocols that wrap the two party trade protocol, to give us the demo logic.

private class TraderDemoProtocolBuyer(val otherSide: Party,
                                      private val attachmentsPath: Path,
                                      val amount: Amount<Currency>,
                                      override val progressTracker: ProgressTracker = ProgressTracker(STARTING_BUY)) : ProtocolLogic<Unit>() {

    object STARTING_BUY : ProgressTracker.Step("Seller connected, purchasing commercial paper asset")

    @Suspendable
    override fun call() {
        progressTracker.currentStep = STARTING_BUY

        val notary: NodeInfo = serviceHub.networkMapCache.notaryNodes[0]
        val buyer = TwoPartyTradeProtocol.Buyer(
                otherSide,
                notary.notaryIdentity,
                amount,
                CommercialPaper.State::class.java)

        // This invokes the trading protocol and out pops our finished transaction.
        val tradeTX: SignedTransaction = subProtocol(buyer, shareParentSessions = true)
        // TODO: This should be moved into the protocol itself.
        serviceHub.recordTransactions(listOf(tradeTX))

        log.info("Purchase complete - we are a happy customer! Final transaction is: " +
                "\n\n${Emoji.renderIfSupported(tradeTX.tx)}")

        logIssuanceAttachment(tradeTX)
        logBalance()
    }

    private fun logBalance() {
        val balances = serviceHub.vaultService.currentVault.cashBalances.entries.map { "${it.key.currencyCode} ${it.value}" }
        logger.info("Remaining balance: ${balances.joinToString()}")
    }

    private fun logIssuanceAttachment(tradeTX: SignedTransaction) {
        // Find the original CP issuance.
        val search = TransactionGraphSearch(serviceHub.storageService.validatedTransactions, listOf(tradeTX.tx))
        search.query = TransactionGraphSearch.Query(withCommandOfType = CommercialPaper.Commands.Issue::class.java,
                followInputsOfType = CommercialPaper.State::class.java)
        val cpIssuance = search.call().single()

        cpIssuance.attachments.first().let {
            val p = attachmentsPath.toAbsolutePath().resolve("$it.jar")
            log.info("""

The issuance of the commercial paper came with an attachment. You can find it expanded in this directory:
$p

${Emoji.renderIfSupported(cpIssuance)}""")
        }
    }
}

private class TraderDemoProtocolSeller(val otherSide: Party,
                                       val amount: Amount<Currency>,
                                       override val progressTracker: ProgressTracker = TraderDemoProtocolSeller.tracker()) : ProtocolLogic<SignedTransaction>() {
    companion object {
        val PROSPECTUS_HASH = SecureHash.parse("decd098666b9657314870e192ced0c3519c2c9d395507a238338f8d003929de9")

        object SELF_ISSUING : ProgressTracker.Step("Got session ID back, issuing and timestamping some commercial paper")

        object TRADING : ProgressTracker.Step("Starting the trade protocol") {
            override fun childProgressTracker(): ProgressTracker = TwoPartyTradeProtocol.Seller.tracker()
        }

        // We vend a progress tracker that already knows there's going to be a TwoPartyTradingProtocol involved at some
        // point: by setting up the tracker in advance, the user can see what's coming in more detail, instead of being
        // surprised when it appears as a new set of tasks below the current one.
        fun tracker() = ProgressTracker(SELF_ISSUING, TRADING)
    }

    @Suspendable
    override fun call(): SignedTransaction {
        progressTracker.currentStep = SELF_ISSUING

        val notary: NodeInfo = serviceHub.networkMapCache.notaryNodes[0]
        val cpOwnerKey = serviceHub.legalIdentityKey
        val commercialPaper = selfIssueSomeCommercialPaper(cpOwnerKey.public, notary)

        progressTracker.currentStep = TRADING

        val seller = TwoPartyTradeProtocol.Seller(
                otherSide,
                notary,
                commercialPaper,
                amount,
                cpOwnerKey,
                progressTracker.getChildProgressTracker(TRADING)!!)
        val tradeTX: SignedTransaction = subProtocol(seller, shareParentSessions = true)
        serviceHub.recordTransactions(listOf(tradeTX))

        return tradeTX
    }

    @Suspendable
    fun selfIssueSomeCommercialPaper(ownedBy: PublicKey, notaryNode: NodeInfo): StateAndRef<CommercialPaper.State> {
        // Make a fake company that's issued its own paper.
        val keyPair = generateKeyPair()
        val party = Party("Bank of London", keyPair.public)

        val issuance: SignedTransaction = run {
            val tx = CommercialPaper().generateIssue(party.ref(1, 2, 3), 1100.DOLLARS `issued by` DUMMY_CASH_ISSUER,
                    Instant.now() + 10.days, notaryNode.notaryIdentity)

            // TODO: Consider moving these two steps below into generateIssue.

            // Attach the prospectus.
            tx.addAttachment(serviceHub.storageService.attachments.openAttachment(PROSPECTUS_HASH)!!.id)

            // Requesting timestamping, all CP must be timestamped.
            tx.setTime(Instant.now(), 30.seconds)

            // Sign it as ourselves.
            tx.signWith(keyPair)

            // Get the notary to sign the timestamp
            val notarySig = subProtocol(NotaryProtocol.Client(tx.toSignedTransaction(false)))
            tx.addSignatureUnchecked(notarySig)

            // Commit it to local storage.
            val stx = tx.toSignedTransaction(true)
            serviceHub.recordTransactions(listOf(stx))

            stx
        }

        // Now make a dummy transaction that moves it to a new key, just to show that resolving dependencies works.
        val move: SignedTransaction = run {
            val builder = TransactionType.General.Builder(notaryNode.notaryIdentity)
            CommercialPaper().generateMove(builder, issuance.tx.outRef(0), ownedBy)
            builder.signWith(keyPair)
            val notarySignature = subProtocol(NotaryProtocol.Client(builder.toSignedTransaction(false)))
            builder.addSignatureUnchecked(notarySignature)
            val tx = builder.toSignedTransaction(true)
            serviceHub.recordTransactions(listOf(tx))
            tx
        }

        return move.tx.outRef(0)
    }

}

private fun printHelp(parser: OptionParser) {
    println("""
    Usage: trader-demo --role [BUYER|SELLER] [options]
    Please refer to the documentation in docs/build/index.html for more info.

    """.trimIndent())
    parser.printHelpOn(System.out)
}

