package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.client.MeshClient;
import com.atlassian.bitbucket.mesh.config.ApplicationConfig;
import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenRegistry;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.execution.ExecutionLocal;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.management.SidebandManager;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.repository.RepositoryAttributeManager;
import com.atlassian.bitbucket.mesh.repository.RepositoryHashService;
import com.atlassian.bitbucket.mesh.repository.RepositoryManager;
import io.grpc.Context;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
import java.util.function.Function;

import static com.google.common.base.Preconditions.checkState;
import static java.util.Objects.requireNonNull;
import static java.util.Optional.ofNullable;

public class DefaultTransactionManager implements TransactionManager {

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

    private final ExecutionLocal<DefaultTransaction> activeTransaction;
    private final ApplicationConfig config;
    private final String localIdPrefix;
    private final EventPublisher eventPublisher;
    private final LogicalClock logicalClock;
    private final MeshClient meshClient;
    private final MeshManager meshManager;
    private final AtomicLong nextLocalId;
    private final RepositoryAttributeManager repositoryAttributeManager;
    private final RepositoryHashService repositoryHashService;
    private final RepositoryManager repositoryManager;
    private final SidebandManager sidebandManager;
    private final UnregisterActiveTransactionOnCompletion unregisterActiveTransactionOnCompletion;
    private final WriteTokenRegistry writeTokenRegistry;

    public DefaultTransactionManager(ApplicationConfig config, EventPublisher eventPublisher,
                                     ExecutionManager executionManager, LogicalClock logicalClock,
                                     MeshClient meshClient, MeshManager meshManager,
                                     RepositoryAttributeManager repositoryAttributeManager,
                                     RepositoryHashService repositoryHashService, RepositoryManager repositoryManager,
                                     SidebandManager sidebandManager, WriteTokenRegistry writeTokenRegistry) {
        this.config = config;
        this.eventPublisher = eventPublisher;
        this.logicalClock = logicalClock;
        this.meshClient = meshClient;
        this.meshManager = meshManager;
        this.repositoryAttributeManager = repositoryAttributeManager;
        this.repositoryHashService = repositoryHashService;
        this.repositoryManager = repositoryManager;
        this.sidebandManager = sidebandManager;
        this.writeTokenRegistry = writeTokenRegistry;

        activeTransaction = executionManager.newExecutionLocal(DefaultTransactionManager::completeIfNotFinished);
        localIdPrefix = randomId();
        nextLocalId = new AtomicLong(1);
        unregisterActiveTransactionOnCompletion = new UnregisterActiveTransactionOnCompletion();
    }

    @Override
    public boolean addListener(@Nonnull TransactionListener listener) {
        return activeTransaction.ifPresent(tx -> tx.addListener(listener));
    }

    @Override
    public <T> T execute(@Nonnull Function<? super Transaction, T> callback) {
        requireNonNull(callback, "callback");

        boolean transactionStarted = false;
        DefaultTransaction transaction = activeTransaction.get();
        if (transaction == null) {
            //No transaction was currently active; start a new one
            transaction = createLeaderTransaction(true);
            if (activeTransaction.set(transaction)) {
                log.debug("Registering transaction {} to automatically roll back on cleanup of the current RPC",
                        transaction.getId());
            } else {
                log.debug("Automatic rollback on cleanup is not available for transaction {} because no RPC is active",
                        transaction.getId());
            }
            transactionStarted = true;
        }

        return invokeWithTransaction(transaction, callback, transactionStarted);
    }

    @Override
    public <T> T executeLocalOnly(@Nonnull Function<? super Transaction, T> callback) {
        requireNonNull(callback, "callback");

        DefaultTransaction suspendedTx = activeTransaction.get();
        DefaultTransaction localOnlyTx = createLeaderTransaction(false);
        activeTransaction.set(localOnlyTx);
        if (log.isDebugEnabled()) {
            log.debug("Started local-only transaction {}{}", localOnlyTx.getId(),
                    (suspendedTx == null ? "" : " and suspended " + suspendedTx.getId()));
        }
        try {
            return invokeWithTransaction(localOnlyTx, callback, true);
        } finally {
            if (suspendedTx != null) {
                activeTransaction.set(suspendedTx);
                log.debug("Resumed transaction {}", suspendedTx.getId());
            }
        }
    }

    @Override
    public void executeLocalOnlyWithoutResult(@Nonnull Consumer<? super Transaction> callback) {
        requireNonNull(callback, "callback");

        executeLocalOnly(tx -> {
            callback.accept(tx);
            return null;
        });
    }

    @Override
    public void executeWithoutResult(@Nonnull Consumer<? super Transaction> callback) {
        requireNonNull(callback, "callback");

        execute(tx -> {
            callback.accept(tx);
            return null;
        });
    }

    @Nonnull
    @Override
    public Optional<Transaction> getCurrentTransaction() {
        return ofNullable(activeTransaction.get());
    }

    @Nonnull
    @Override
    public PrepareResult prepared(@Nonnull PreparedAction action) {
        requireNonNull(action, "action");

        //If there is an active execution context, start a transaction that closes when the execution context closes.
        DefaultTransaction transaction = activeTransaction.compute(
                existing -> existing == null ? createLeaderTransaction(true) : existing);
        if (transaction == null) {
            //There is no active execution context, create a leader transaction that completes immediately after
            //the outcome is determined
            try (DefaultTransaction tx = createLeaderTransaction(true)) {
                return tx.prepared(action);
            }
        }
        return transaction.prepared(action);
    }

    @Nonnull
    @Override
    public Transaction start() {
        DefaultTransaction transaction = activeTransaction.get();
        if (transaction == null) {
            //No transaction was currently active; start a new one
            transaction = createLeaderTransaction(true);
            if (activeTransaction.set(transaction)) {
                log.debug("Registering transaction {} to automatically roll back on cleanup of the current RPC",
                        transaction.getId());
            } else {
                log.debug("Automatic rollback on cleanup is not available for transaction {} because no RPC is active",
                        transaction.getId());
            }
            return transaction;
        }
        throw new IllegalStateException("Another transaction " + transaction.getId() + " is already active");
    }

    @Nonnull
    @Override
    public DefaultTransaction start(@Nonnull StartTransactionParameters parameters) {
        requireNonNull(parameters, "parameters");
        activeTransaction.checkValid();

        //noinspection ConstantConditions
        return activeTransaction.compute(existing -> {
            TransactionLeader leader = parameters.getLeader();
            boolean local = leader.getNode().isLocal();
            String id = local ? nextLocalId() : requireNonNull(parameters.getId(), "transactionId");

            if (existing == null) {
                //No transaction is currently active
                log.debug("Starting transaction {} coordinated by {}", id, leader.getNode());

                DefaultTransaction transaction = new DefaultTransaction(id, leader);
                transaction.addListener(unregisterActiveTransactionOnCompletion);
                return transaction;
            }

            //Verify that the currently active transaction is compatible
            checkState(id.equals(existing.getId()), "Another transaction (%s) is already active " +
                    "when starting follower transaction %s", existing.getId(), id);
            return existing;
        });
    }

    @Override
    public RepositoryConsistencySummary verifyConsistency(@Nonnull String repositoryId) {
        requireNonNull(repositoryId, "repositoryId");

        String txId = Transaction.formatId(meshManager.getLocalNode().getId(), nextLocalId());
        return new RepositoryConsistencyVerifier(repositoryId, txId, logicalClock, eventPublisher, meshClient,
                meshManager, repositoryAttributeManager, repositoryHashService, repositoryManager, sidebandManager,
                writeTokenRegistry).call();
    }

    private static void completeIfNotFinished(Transaction transaction) {
        if (!transaction.getState().isFinished()) {
            log.debug("Completing dangling transaction {} at the end of RPC execution", transaction.getId(),
                    Context.current().cancellationCause());
            transaction.close();
        }
    }

    private static <T> T invokeWithTransaction(DefaultTransaction transaction,
                                               Function<? super Transaction, T> callback,
                                               boolean transactionStarted) {
        //Only handle the exception if the transaction was started here. Otherwise, just propagate any exception
        //thrown by the callback to allow the caller to recover from it.
        if (transactionStarted) {
            try (DefaultTransaction tx = transaction) {
                return tx.rollBackIfThrows(() -> callback.apply(tx));
            }
        }

        return callback.apply(transaction);
    }

    private static String randomId() {
        return Long.toHexString(Math.abs(UUID.randomUUID().getMostSignificantBits()))
                .substring(0, 8);
    }

    private DefaultTransaction createLeaderTransaction(boolean replicate) {
        MeshNode node = meshManager.getLocalNode();
        String txId = Transaction.formatId(node.getId(), nextLocalId());

        //Create a LocalTransactionLeader because the current node is the transaction leader
        TransactionLeader leader = replicate ?
                new LocalTransactionLeader(txId, config, eventPublisher, logicalClock, meshClient, meshManager,
                        repositoryAttributeManager, repositoryHashService,
                        sidebandManager, writeTokenRegistry) :
                new LocalOnlyTransactionLeader(txId, meshManager);

        DefaultTransaction transaction = new DefaultTransaction(txId, leader);
        transaction.addListener(unregisterActiveTransactionOnCompletion);
        return transaction;
    }

    private String nextLocalId() {
        return localIdPrefix + StringUtils.leftPad(Long.toHexString(nextLocalId.incrementAndGet()), 16, '0');
    }

    private class UnregisterActiveTransactionOnCompletion implements TransactionListener {

        @Override
        public void afterCompletion(@Nonnull Transaction transaction) {
            activeTransaction.removeIf(activeTx -> activeTx.equals(transaction));
        }

        @Override
        public int getOrder() {
            //Unregistering the active transaction should run last
            return Integer.MAX_VALUE;
        }
    }
}
