package net.corda.notary.jpa

import com.google.common.collect.Queues
import net.corda.core.concurrent.CordaFuture
import net.corda.core.contracts.StateRef
import net.corda.core.contracts.TimeWindow
import net.corda.core.crypto.SecureHash
import net.corda.core.flows.NotarisationRequestSignature
import net.corda.core.flows.NotaryError
import net.corda.core.flows.StateConsumptionDetails
import net.corda.core.identity.CordaX500Name
import net.corda.core.identity.Party
import net.corda.core.internal.concurrent.OpenFuture
import net.corda.core.internal.concurrent.openFuture
import net.corda.core.internal.mapToSet
import net.corda.core.internal.notary.NotaryInternalException
import net.corda.core.internal.notary.UniquenessProvider
import net.corda.core.internal.notary.isConsumedByTheSameTx
import net.corda.core.internal.notary.validateTimeWindow
import net.corda.core.schemas.PersistentStateRef
import net.corda.core.serialization.CordaSerializable
import net.corda.core.serialization.SerializationDefaults
import net.corda.core.serialization.SingletonSerializeAsToken
import net.corda.core.serialization.serialize
import net.corda.core.utilities.contextLogger
import net.corda.core.utilities.debug
import net.corda.node.services.vault.toStateRef
import net.corda.nodeapi.internal.persistence.CordaPersistence
import net.corda.nodeapi.internal.persistence.NODE_DATABASE_PREFIX
import net.corda.notary.common.BatchSigningFunction
import net.corda.notary.common.InternalResult
import net.corda.serialization.internal.CordaSerializationEncoding
import org.hibernate.Session
import java.sql.SQLException
import java.time.Clock
import java.time.Instant
import java.util.LinkedList
import java.util.UUID
import java.util.concurrent.LinkedBlockingQueue
import java.util.concurrent.TimeUnit
import javax.annotation.concurrent.ThreadSafe
import javax.persistence.Column
import javax.persistence.EmbeddedId
import javax.persistence.Entity
import javax.persistence.Id
import javax.persistence.Lob
import javax.persistence.NamedQuery
import kotlin.concurrent.thread

/** A JPA backed Uniqueness provider */
@Suppress("MagicNumber") // database column length
@ThreadSafe
class JPAUniquenessProvider(
        val clock: Clock,
        val database: CordaPersistence,
        val config: JPANotaryConfiguration = JPANotaryConfiguration(),
        val notaryWorkerName: CordaX500Name,
        val signBatch: BatchSigningFunction
) : UniquenessProvider, SingletonSerializeAsToken() {

    // This is the prefix of the ID in the request log table, to allow running multiple instances that access the
    // same table.
    private val instanceId = UUID.randomUUID()

    @Entity
    @javax.persistence.Table(name = "${NODE_DATABASE_PREFIX}notary_request_log")
    @CordaSerializable
    class Request(
            @Id
            @Column(nullable = true, length = 76)
            var id: String? = null,

            @Column(name = "consuming_transaction_id", nullable = true, length = 144)
            val consumingTxHash: String?,

            @Column(name = "requesting_party_name", nullable = true, length = 255)
            var partyName: String?,

            @Lob
            @Column(name = "request_signature", nullable = false)
            val requestSignature: ByteArray,

            @Column(name = "request_timestamp", nullable = false)
            var requestDate: Instant,

            @Column(name = "worker_node_x500_name", nullable = true, length = 255)
            val workerNodeX500Name: String?
    )

    private data class CommitRequest(
            val states: List<StateRef>,
            val txId: SecureHash,
            val callerIdentity: Party,
            val requestSignature: NotarisationRequestSignature,
            val timeWindow: TimeWindow?,
            val references: List<StateRef>,
            val future: OpenFuture<UniquenessProvider.Result>,
            val requestEntity: Request,
            val committedStatesEntities: List<CommittedState>)

    @Entity
    @javax.persistence.Table(name = "${NODE_DATABASE_PREFIX}notary_committed_states")
    @NamedQuery(name = "CommittedState.select", query = "SELECT c from JPAUniquenessProvider\$CommittedState c WHERE c.id in :ids")
    class CommittedState(
            @EmbeddedId
            val id: PersistentStateRef,
            @Column(name = "consuming_transaction_id", nullable = false, length = 144)
            val consumingTxHash: String)

    @Entity
    @javax.persistence.Table(name = "${NODE_DATABASE_PREFIX}notary_committed_txs")
    class CommittedTransaction(
            @Id
            @Column(name = "transaction_id", nullable = false, length = 144)
            val transactionId: String
    )

    private val requestQueue = LinkedBlockingQueue<CommitRequest>(requestQueueSize)

    /** A requestEntity processor thread. */
    private val processorThread = thread(name = "Notary request queue processor", isDaemon = true) {
        try {
            val buffer = LinkedList<CommitRequest>()
            while (!Thread.interrupted()) {
                val drainedSize = Queues.drain(requestQueue, buffer, config.batchSize, config.batchTimeoutMs, TimeUnit.MILLISECONDS)
                if (drainedSize == 0) continue
                processRequests(buffer)
                buffer.clear()
            }
        } catch (_: InterruptedException) {
            log.debug { "Process interrupted."}
        }
        log.debug { "Shutting down with ${requestQueue.size} in-flight requests unprocessed." }
    }

    fun stop() {
        processorThread.interrupt()
    }

    companion object {
        private const val requestQueueSize = 100_000
        private const val jdbcBatchSize = 100_000
        private val log = contextLogger()

        fun encodeStateRef(s: StateRef): PersistentStateRef {
            return PersistentStateRef(s.txhash.toString(), s.index)
        }
    }

    /**
     * Generates and adds a [CommitRequest] to the requestEntity queue. If the requestEntity queue is full, this method will block
     * until space is available.
     *
     * Returns a future that will complete once the requestEntity is processed, containing the commit [Result].
     */
    override fun commit(
            states: List<StateRef>,
            txId: SecureHash,
            callerIdentity: Party,
            requestSignature: NotarisationRequestSignature,
            timeWindow: TimeWindow?,
            references: List<StateRef>
    ): CordaFuture<UniquenessProvider.Result> {
        val future = openFuture<UniquenessProvider.Result>()
        val requestEntities = Request(consumingTxHash = txId.toString(),
                partyName = callerIdentity.name.toString(),
                requestSignature = requestSignature.serialize(context = SerializationDefaults.STORAGE_CONTEXT.withEncoding(CordaSerializationEncoding.SNAPPY)).bytes,
                requestDate = clock.instant(),
                workerNodeX500Name = notaryWorkerName.toString())
        val stateEntities = states.map {
            CommittedState(
                    encodeStateRef(it),
                    txId.toString()
            )
        }
        val request = CommitRequest(states, txId, callerIdentity, requestSignature, timeWindow, references, future, requestEntities, stateEntities)

        requestQueue.put(request)

        return future
    }

    // Safe up to 100k requests per second.
    private var nextRequestId = System.currentTimeMillis() * 100

    private fun logRequests(requests: List<CommitRequest>) {
        database.transaction {
            for (request in requests) {
                request.requestEntity.id = "$instanceId:${(nextRequestId++).toString(16)}"
                session.persist(request.requestEntity)
            }
        }
    }

    private fun commitRequests(session: Session, requests: List<CommitRequest>) {
        for (request in requests) {
            for (cs in request.committedStatesEntities) {
                session.persist(cs)
            }
            session.persist(CommittedTransaction(request.txId.toString()))
        }
    }

    private fun findAlreadyCommitted(session: Session, states: List<StateRef>, references: List<StateRef>): Map<StateRef, StateConsumptionDetails> {
        val persistentStateRefs = (states + references).mapToSet(::encodeStateRef)
        val committedStates = mutableListOf<CommittedState>()

        for (idsBatch in persistentStateRefs.chunked(config.maxInputStates)) {
            @Suppress("UNCHECKED_CAST")
            val existing = session
                    .createNamedQuery("CommittedState.select")
                    .setParameter("ids", idsBatch)
                    .resultList as List<CommittedState>
            committedStates.addAll(existing)
        }

        return committedStates.associate {
            val stateRef = it.id.toStateRef()
            val consumingTxId = SecureHash.create(it.consumingTxHash)
            if (stateRef in references) {
                stateRef to StateConsumptionDetails(consumingTxId.reHash(), type = StateConsumptionDetails.ConsumedStateType.REFERENCE_INPUT_STATE)
            } else {
                stateRef to StateConsumptionDetails(consumingTxId.reHash())
            }
        }
    }

    private fun<T> withRetry(block: () -> T): T {
        var retryCount = 0
        var backOff = config.backOffBaseMs
        var exceptionCaught: SQLException? = null
        while (retryCount <= config.maxDBTransactionRetryCount) {
            try {
                return block()
            } catch (e: SQLException) {
                retryCount++
                Thread.sleep(backOff)
                backOff *= 2
                exceptionCaught = e
            }
        }
        throw exceptionCaught!!
    }

    private fun findAllConflicts(session: Session, requests: List<CommitRequest>): MutableMap<StateRef, StateConsumptionDetails> {
        log.info("Processing notarization requests with ${requests.sumOf { it.states.size }} input states and ${requests.sumOf { it.references.size }} references")

        val allStates = requests.flatMap { it.states }
        val allReferences = requests.flatMap { it.references }
        return findAlreadyCommitted(session, allStates, allReferences).toMutableMap()
    }

    private fun processRequest(
            session: Session,
            request: CommitRequest,
            consumedStates: MutableMap<StateRef, StateConsumptionDetails>,
            processedTxIds: MutableMap<SecureHash, InternalResult>,
            toCommit: MutableList<CommitRequest>
    ): InternalResult {
        val conflicts = (request.states + request.references).mapNotNull {
            if (consumedStates.containsKey(it)) it to consumedStates[it]!!
            else null
        }.toMap()

        return if (conflicts.isNotEmpty()) {
            handleStateConflicts(request, conflicts, session)
        } else {
            handleNoStateConflicts(request, toCommit, consumedStates, processedTxIds, session)
        }
    }

    /**
     * Process the [request] given there are conflicting states already present in the DB or current batch.
     *
     * To ensure idempotency, if the request's transaction matches a previously consumed transaction then the
     * same result (success) can be returned without committing it to the DB. Failure is only returned in the
     * case where the request is not a duplicate of a previously processed request and hence it is a genuine
     * double spend attempt.
     */
    private fun handleStateConflicts(
            request: CommitRequest,
            stateConflicts: Map<StateRef, StateConsumptionDetails>,
            session: Session
    ): InternalResult {
        return when {
            isConsumedByTheSameTx(request.txId.reHash(), stateConflicts) -> {
                InternalResult.Success
            }
            request.states.isEmpty() && isPreviouslyNotarised(session, request.txId) -> {
                InternalResult.Success
            }
            else -> {
                InternalResult.Failure(NotaryError.Conflict(request.txId, stateConflicts))
            }
        }
    }

    /**
     * Process the [request] given there are no conflicting states already present in the DB or current batch.
     *
     * This method performs time window validation and adds the request to the commit list if applicable.
     * It also checks the [processedTxIds] map to ensure that any time-window only duplicates within the batch
     * are only committed once.
     */
    private fun handleNoStateConflicts(
            request: CommitRequest,
            toCommit: MutableList<CommitRequest>,
            consumedStates: MutableMap<StateRef, StateConsumptionDetails>,
            processedTxIds: MutableMap<SecureHash, InternalResult>,
            session: Session
    ): InternalResult {
        return when {
            request.states.isEmpty() && isPreviouslyNotarised(session, request.txId) -> {
                InternalResult.Success
            }
            processedTxIds.containsKey(request.txId) -> {
                processedTxIds[request.txId]!!
            }
            else -> {
                val outsideTimeWindowError = validateTimeWindow(clock.instant(), request.timeWindow)
                val internalResult = if (outsideTimeWindowError != null) {
                    InternalResult.Failure(outsideTimeWindowError)
                } else {
                    // Mark states as consumed to capture conflicting transactions in the same batch
                    request.states.forEach {
                        consumedStates[it] = StateConsumptionDetails(request.txId.reHash())
                    }
                    toCommit.add(request)
                    InternalResult.Success
                }
                // Store transaction result to capture conflicting time-window only transactions in the same batch
                processedTxIds[request.txId] = internalResult
                internalResult
            }
        }
    }

    private fun isPreviouslyNotarised(session: Session, txId: SecureHash): Boolean {
        return session.find(CommittedTransaction::class.java, txId.toString()) != null
    }

    private fun processRequests(requests: List<CommitRequest>) {
        try {
            // Note that there is an additional retry mechanism within the transaction itself.
            val res = withRetry {
                database.transaction {
                    val em = session.entityManagerFactory.createEntityManager()
                    em.unwrap(Session::class.java).jdbcBatchSize = jdbcBatchSize

                    val toCommit = mutableListOf<CommitRequest>()
                    val consumedStates = findAllConflicts(session, requests)
                    val processedTxIds = mutableMapOf<SecureHash, InternalResult>()

                    val results = requests.map { request ->
                        processRequest(session, request, consumedStates, processedTxIds, toCommit)
                    }

                    logRequests(requests)
                    commitRequests(session, toCommit)

                    results
                }
            }
            completeResponses(requests, res)
        } catch (e: Exception) {
            log.warn("Error processing commit requests", e)
            for (request in requests) {
                respondWithError(request, e)
            }
        }
    }

    private fun completeResponses(requests: List<CommitRequest>, results: List<InternalResult>): Int {
        val zippedResults = requests.zip(results)
        val successfulRequests = zippedResults
                .filter { it.second is InternalResult.Success }
                .map { it.first.txId }
                .distinct()
        val signature = if (successfulRequests.isNotEmpty())
            signBatch(successfulRequests)
        else null

        var inputStateCount = 0
        for ((request, result) in zippedResults) {
            val resultToSet = when {
                result is InternalResult.Failure -> UniquenessProvider.Result.Failure(result.error)
                signature != null -> UniquenessProvider.Result.Success(signature.forParticipant(request.txId))
                else -> throw IllegalStateException("Signature is required but not found")
            }

            request.future.set(resultToSet)
            inputStateCount += request.states.size
        }
        return inputStateCount
    }

    private fun respondWithError(request: CommitRequest, exception: Exception) {
        if (exception is NotaryInternalException) {
            request.future.set(UniquenessProvider.Result.Failure(exception.error))
        } else {
            request.future.setException(NotaryInternalException(NotaryError.General(Exception("Internal service error."))))
        }
    }
}
