package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.config.ReplicationConfig;
import com.atlassian.bitbucket.mesh.execution.ExecutionLocal;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.repository.ManagedRepository;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.atlassian.bitbucket.mesh.repository.RepositoryHashService;
import com.atlassian.bitbucket.mesh.repository.RepositoryVisitor;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcTransaction;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcTransactionCommitFailed;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcTransactionRollbackFailed;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcUpdateRefsRequest;
import com.atlassian.bitbucket.mesh.util.Person;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.concurrent.ThreadSafe;

import static com.atlassian.bitbucket.mesh.transaction.ReferenceTransactionState.*;
import static com.google.common.base.Preconditions.checkState;

/**
 * {@link ReferenceTransactionListener} implementation that coordinates {@code git} ref updates with
 * application transaction management. Subsequently, application transaction management will coordinate with
 * other replicas to apply the ref update(s) consistently.
 */
public class CoordinatingReferenceTransactionListener implements ReferenceTransactionListener {

    private static final Logger log = LoggerFactory.getLogger(CoordinatingReferenceTransactionListener.class);

    private final ReplicationConfig config;
    private final ExecutionLocal<PreparedRefUpdates> preparedUpdate;
    private final RepositoryHashService repositoryHashService;
    private final TransactionManager transactionManager;

    public CoordinatingReferenceTransactionListener(ReplicationConfig config,
                                                    ExecutionManager executionManager,
                                                    RepositoryHashService repositoryHashService,
                                                    TransactionManager transactionManager) {
        this.config = config;
        this.repositoryHashService = repositoryHashService;
        this.transactionManager = transactionManager;

        preparedUpdate = executionManager.newExecutionLocal();
    }

    @Override
    public void onReferenceTransaction(@Nonnull ReferenceTransaction refTransaction) {
        //We only want the coordination logic to run on managed repositories. All the other types, including
        //shared objects repositories, should bypass this logic because they either don't have replicas or
        //don't have transactional updates
        refTransaction.getRepository().accept(new RepositoryVisitor<Void>() {

            @Override
            public Void visit(@Nonnull ManagedRepository repository) {
                switch (refTransaction.getState()) {
                    case ABORTED:
                        maybeRollBack(refTransaction);
                        break;
                    case COMMITTED:
                        maybeCommit(refTransaction);
                        break;
                    case PREPARED:
                        maybePrepare(refTransaction);
                        break;
                }

                return null;
            }
        });
    }

    private void maybeCommit(ReferenceTransaction refTransaction) {
        if (preparedUpdate.ifPresent(prepared -> prepared.onHookCallback(refTransaction))) {
            transactionManager.getCurrentTransaction().ifPresent(Transaction::close);
        } else {
            Repository repository = refTransaction.getRepository();
            if (preparedUpdate.isActive()) {
                log.warn("[{}] {} callback received but no transaction has been PREPARED",
                        repository, refTransaction.getState());
            } else {
                //If there's no active context, there will never be a prepared update
                //from a previous reference-transaction callback
                log.debug("[{}] {} callback received without an execution context. It cannot be correlated to" +
                                " a PREPARED callback", repository, refTransaction.getState());
            }
            repositoryHashService.resetContentHash(repository);
        }
    }

    private void maybePrepare(ReferenceTransaction refTransaction) {
        //Only coordinate with application transaction management if the reference transaction hasn't been
        //rejected yet (by other ReferenceTransactionListeners such as the listener that enforces read-only
        //mode).
        if (refTransaction.isRejected()) {
            //Another reference-transaction listener has already rejected the update. No need to coordinate
            //with the transaction manager, but we should still initialize preparedUpdate to allow the subsequent
            //ABORTED callback to match up with them (and not WARN log).
            PreparedRefUpdates refUpdates = preparedRefUpdatesFor(refTransaction);
            refUpdates.expect(PrepareResult.FAILURE);
            preparedUpdate.set(refUpdates);
        } else {
            try {
                PreparedRefUpdates refUpdates = preparedRefUpdatesFor(refTransaction);
                PrepareResult result = transactionManager.prepared(refUpdates);
                refUpdates.expect(result);
                if (!result.isSuccess()) {
                    refTransaction.reject("Ref update could not be replicated");
                }

                //Store the transactional in the request context so we can look it up again in
                //the ABORTED or COMMITTED reference-transaction hook callbacks
                preparedUpdate.set(refUpdates);
            } catch (RuntimeException e) {
                log.error("[{}] Failed to replicate the ref update", refTransaction.getRepository(), e);
                refTransaction.reject("Failed to replicate ref update");
            }
        }
    }

    private void maybeRollBack(ReferenceTransaction refTransaction) {
        //Only roll back the application transaction if one is active
        if (preparedUpdate.ifPresent(prepared -> prepared.onHookCallback(refTransaction))) {
            transactionManager.getCurrentTransaction().ifPresent(appTransaction -> {
                appTransaction.setRollbackOnly();
                appTransaction.close();
            });
        } else if (refTransaction.getState() != PREPARED) {
            if (preparedUpdate.isActive()) {
                log.warn("[{}] {} callback received but no transaction has been PREPARED",
                        refTransaction.getRepository(), refTransaction.getState());
            } else {
                //If there's no active context, there will never be a prepared update
                //from a previous reference-transaction callback
                log.debug("[{}] {} callback received without an execution context. It cannot be correlated to" +
                        " a PREPARED callback", refTransaction.getRepository(), refTransaction.getState());
            }
        }
    }

    private PreparedRefUpdates preparedRefUpdatesFor(ReferenceTransaction refTransaction) {
        return new PreparedRefUpdates(refTransaction, repositoryHashService);
    }

    /**
     * Integrates the git reference-transaction callbacks with application transaction management to verify
     * that the expected reference-transaction callbacks have been received for when the transaction commits or
     * rolls back.
     */
    @ThreadSafe
    private static class PreparedRefUpdates implements PreparedAction {

        private final Object lock;
        private final Repository repository;
        private final RepositoryHashService repositoryHashService;
        private final String transactionHash;
        private final Person user;

        private ReferenceTransactionState expectedState;
        private volatile String fromHash;
        private ReferenceTransactionState lastKnownState;
        private volatile ReferenceTransaction referenceTransaction;

        PreparedRefUpdates(ReferenceTransaction referenceTransaction,
                           RepositoryHashService repositoryHashService) {
            this.referenceTransaction = referenceTransaction;
            this.repositoryHashService = repositoryHashService;

            lastKnownState = referenceTransaction.getState();
            lock = new Object();
            repository = referenceTransaction.getRepository();
            transactionHash = hash(referenceTransaction);
            user = referenceTransaction.getUser();
        }

        @Override
        public void commit() {
            synchronized (lock) {
                if (expectedState != COMMITTED) {
                    throw new TransactionCommitFailedException("Transaction committed unexpectedly (expected " +
                            expectedState + ", latest callback was " + lastKnownState + ")",
                            RpcTransactionCommitFailed.Reason.UNKNOWN);
                }
                if (lastKnownState != COMMITTED) {
                    throw new TransactionCommitFailedException("COMMITTED reference-transaction hook has not been " +
                            "received. The ref updates may not have been applied",
                            RpcTransactionCommitFailed.Reason.PUBLISH_FAILED, null);
                }
                repositoryHashService.compareAndUpdateContentHash(repository, fromHash, transactionHash);
            }
        }

        @Nonnull
        @Override
        public Repository getRepository() {
            return repository;
        }

        @Nonnull
        @Override
        public RpcTransaction.Builder populateTransaction(@Nonnull RpcTransaction.Builder builder) {
            checkState(referenceTransaction != null, "populateTransaction called after the transaction is %s",
                    (expectedState == COMMITTED ? "committing" : "rolling back"));

            RpcUpdateRefsRequest.Builder requestBuilder = RpcUpdateRefsRequest.newBuilder()
                    .setRepository(getRepository().getId());
            referenceTransaction.getRefUpdates().forEach(refUpdate ->
                    requestBuilder.addRefUpdates(refUpdate.toProposedRefUpdate()));
            if (user != null) {
                builder.setPerson(user.toPerson());
            }
            requestBuilder.setTransactional(true);
            return builder.setUpdateRefs(requestBuilder);
        }

        /**
         * Called by application transaction management when the _application_ transaction rolls back. This method
         * validates that the rollback was both <em>expected</em> because the application
         * {@link ReferenceTransaction#reject(String) rejected} the reference transaction, and that a
         * reference-transaction hook callback with status ABORTED has been received.
         */
        @Override
        public void rollBack() {
            synchronized (lock) {
                //expectedState can be one of three values:
                // - null: the transaction leader hasn't told us yet whether to commit or roll back
                // - COMMITTED: leader told us to commit
                // - ABORTED: leader told us to roll back
                if (expectedState == null) {
                    log.info("[{}] Transaction rolled back before the leader instructed us to commit or roll back",
                            repository);
                    return;
                }
                if (expectedState == COMMITTED) {
                    throw new TransactionCommitFailedException(
                            "Transaction rolled back unexpectedly for " + repository.getId() +
                            " (expected " + expectedState + ", latest callback was " + lastKnownState + ")",
                            RpcTransactionCommitFailed.Reason.PUBLISH_FAILED);
                }
                //expectedState must be ABORTED
                if (lastKnownState == COMMITTED) {
                    throw new TransactionRollbackFailedException(
                            "Cannot roll back already committed transaction for " + repository.getId(),
                            RpcTransactionRollbackFailed.Reason.CLEANUP_FAILED, null);
                } else if (lastKnownState != ABORTED) {
                    //As long as the last known state is not COMMITTED, and we know the PREPARED hook was
                    //rejected, we can _assume_ the ref update hasn't been applied (e.g. process timeout before
                    //the ABORTED timeout was processed)
                    log.info("[{}] Transaction rolled back as expected but no ABORTED reference-transaction hook" +
                            " has been received yet (last known state is {})", repository, lastKnownState);
                }
            }
        }

        void expect(PrepareResult prepareResult) {
            synchronized (lock) {
                if (prepareResult.isSuccess()) {
                    expectedState = COMMITTED;
                    //It's not safe to capture the fromHash until after the transaction has been successfully
                    //prepared and will be committed. At this point the write lock has been acquired, which
                    //means we shouldn't need to worry about racy updates
                    fromHash = repositoryHashService.getContentHash(repository);
                    log.debug("[{}] Recorded baseline content hash {}; waiting to commit", repository, fromHash);
                } else {
                    //If the transaction will roll back, we don't need to capture a fromHash; we won't update
                    //our content hash since the changes won't commit
                    expectedState = ABORTED;
                }

                //When expect is called, the transaction is fully prepared and either COMMITTING or ABORTING.
                //At this point it's safe to null out referenceTransaction to avoid holding on to potentially
                //a lot of memory between reference-transaction callbacks.
                referenceTransaction = null;
            }
        }

        void onHookCallback(ReferenceTransaction transaction) {
            synchronized (lock) {
                if (repository.equals(transaction.getRepository()) && transactionHash.equals(hash(transaction))) {
                    lastKnownState = transaction.getState();
                } else {
                    log.warn("[{}] Ignoring reference-transaction {} callback for {} while {}",
                            repository, transaction.getState(), transaction.getRefUpdates(), lastKnownState);
                }
            }
        }

        private String hash(ReferenceTransaction referenceTransaction) {
            return repositoryHashService.calculateTransactionHash(referenceTransaction.getRefUpdates());
        }
    }
}
