package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamChunk;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcTransaction;
import com.atlassian.bitbucket.mesh.util.SatiableConsumer;
import io.grpc.Context;
import io.grpc.Deadline;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.Supplier;

import static com.atlassian.bitbucket.mesh.transaction.TransactionState.*;
import static com.google.common.base.Preconditions.checkState;
import static java.util.Objects.requireNonNull;

class DefaultTransaction implements Transaction {

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

    private static final long DEFAULT_DEADLINE_MILLIS = TimeUnit.HOURS.toMillis(16);

    private final String id;
    private final TransactionLeader leader;
    private final List<TransactionListener> listeners;
    private final Object lock;

    private volatile PreparedAction action;
    private volatile boolean closed;
    private volatile RuntimeException exceptionWhileCommitting;
    private volatile TransactionState state;

    DefaultTransaction(@Nonnull String id, @Nonnull TransactionLeader leader) {
        this.id = requireNonNull(id, "id");
        this.leader = requireNonNull(leader, "leader");

        listeners = new ArrayList<>();
        lock = new Object();
        state = INITIAL;
    }

    @Override
    public void addListener(@Nonnull TransactionListener listener) {
        requireNonNull(listener, "listener");

        synchronized (lock) {
            if (shouldInvokeListeners()) {
                invokeListener(listener);
            }
            if (!state.isFinished()) {
                listeners.add(listener);
            }
        }
    }

    @Override
    public void close() {
        synchronized (lock) {
            if (!closed) {
                closed = true;
                if (!state.isFinished()) {
                    if (state == COMMITTING) {
                        internalCommit();
                    } else {
                        internalRollBack();
                    }
                }
            }
        }
    }

    @Nonnull
    @Override
    public String getId() {
        return id;
    }

    @Nonnull
    @Override
    public Optional<Path> getIncomingDir() {
        return leader.getIncomingDir();
    }

    @Nonnull
    @Override
    public MeshNode getLeader() {
        return leader.getNode();
    }

    @Nonnull
    @Override
    public TransactionState getState() {
        return state;
    }

    @Override
    public boolean isRollbackOnly() {
        synchronized (lock) {
            return state.isRollbackOnly();
        }
    }

    @Nonnull
    @Override
    public Optional<SatiableConsumer<RpcStreamChunk>> newPackWriter(@Nonnull Repository repository) {
        synchronized (lock) {
            checkState(state == INITIAL, "Cannot open a PackWriter; state = %s", state.name());
        }
        return leader.newPackWriter(repository);
    }

    @Override
    @Nonnull
    public PrepareResult prepared(@Nonnull PreparedAction action) {
        checkState(this.action == null, "Another action has already been prepared");

        this.action = requireNonNull(action, "action");

        synchronized (lock) {
            if (state.isRollbackOnly()) {
                return PrepareResult.FAILURE;
            }
            transitionTo(PREPARED);
        }

        //Notify the transaction leader that this replica successfully prepared.
        Repository repository = action.getRepository();
        Future<PrepareResult> future = leader.onPrepared(repository,
                action.populateTransaction(RpcTransaction.newBuilder()
                        .setId(id)
                        .setRepository(repository.getId())
                        .setTimestamp(System.currentTimeMillis()))
                        //TODO: inject RpcPerson as well?
                        .build());

        //Then wait for the leader's decision
        PrepareResult result = PrepareResult.FAILURE;
        try {
            result = future.get(getDeadlineMillis(), TimeUnit.MILLISECONDS);
            if (log.isDebugEnabled()) {
                log.debug("[{}] Leader {} voted to {} transaction {}", action.getRepository(),
                        leader.getNode(), result.isSuccess() ? "commit" : "roll back", id);
            }
        } catch (InterruptedException e) {
            log.warn("Interrupted while waiting for outcome of transaction {}. Rolling back...", id);
            Thread.currentThread().interrupt();
        } catch (TimeoutException e) {
            log.warn("Timed out waiting for outcome of transaction {}. Rolling back...", id);
            future.cancel(true);
        } catch (ExecutionException e) {
            log.warn("Error while waiting for outcome of transaction {}. Rolling back...", id, e.getCause());
        }

        synchronized (lock) {
            //This further check is to ensure that if we are in the process of rolling back or have already rolled-back
            //the transaction due to other factors that impact the transaction (incomplete requests, etc.) we do not
            //even try to transition to COMMITTING as the result may have been calculated before our state changed.
            if (state.isRollbackOnly()) {
                log.debug("[{}] Transaction '{}' has been rolled back or is rolling back by another transaction;" +
                          " we cannot commit this transaction", action.getRepository(), id);
                return PrepareResult.FAILURE;
            } else {
                transitionTo(result.isSuccess() ? COMMITTING : ROLLING_BACK);
            }
        }
        return result;
    }

    @Override
    public void rollBackIfThrows(@Nonnull Runnable action) {
        rollBackIfThrows(() -> {
            requireNonNull(action, "action").run();
            return null;
        });
    }

    @Override
    public <T> T rollBackIfThrows(@Nonnull Supplier<T> action) {
        try {
            return requireNonNull(action, "action").get();
        } catch (RuntimeException e) {
            synchronized (lock) {
                switch (state) {
                    case INITIAL:
                    case PREPARING:
                    case PREPARED:
                        //An exception was thrown from the transactional scope while the transaction can still cleanly
                        //roll back. Mark the transaction for roll back, but don't record the exception; we only report
                        //exceptions to the transaction leader when we've failed while in COMMITTING or ABORTING state.
                        setRollbackOnly();
                        break;
                    case COMMITTING:
                        //Exception was thrown after we've gone from PREPARED -> COMMITTING but before the transaction
                        //was closed. In this case, we can't commit as we promised the transaction leader that we would.
                        //Record the exception so we can inform the leader of our failure.
                        if (exceptionWhileCommitting == null) {
                            exceptionWhileCommitting = e;
                        } else {
                            exceptionWhileCommitting.addSuppressed(e);
                        }
                        break;
                }
                //We don't do anything for ROLLING_BACK, COMMITTED or ROLLED_BACK. The latter two are final states,
                //so we've already informed the transaction leader of our outcome. If we're currently ROLLING_BACK
                //(i.e. because the leader instructed us to roll back), it's somewhat expected for the transactional
                //scope to end with an exception. We do not need to inform the transaction leader of the exception
                //because we haven't (yet) failed to do what we promised the leader (commit / roll back).
            }
            throw e;
        }
    }

    @Override
    public void setRollbackOnly() {
        synchronized (lock) {
            transitionTo(ROLLING_BACK);
        }
    }

    private static long getDeadlineMillis() {
        Deadline deadline = Context.current().getDeadline();
        return deadline == null ? DEFAULT_DEADLINE_MILLIS : deadline.timeRemaining(TimeUnit.MILLISECONDS);
    }

    private void internalCommit() {
        checkState(action != null, "No operation was prepared before committing!");
        Repository repository = action.getRepository();
        String actionType = action.getClass().getSimpleName();
        RuntimeException commitException = null;
        if (exceptionWhileCommitting == null) {
            try {
                log.debug("[{}] Committing transaction {} (type = {})", repository, id, actionType);
                action.commit();
            } catch (RuntimeException e) {
                exceptionWhileCommitting = commitException = e;
            }
        } else {
            //The transactional operation failed from Transaction.run(operation) after transitioning to COMMITTING
            //but before closing the transaction. Roll back locally and inform the leader that we failed.
            try {
                log.error("[{}] Rolling back {} transaction {} (type = {}) which failed before it could be committed",
                        repository, state, id, actionType, exceptionWhileCommitting);
                action.rollBack();
            } catch (RuntimeException e) {
                log.warn("[{}] Transaction {} did not roll back cleanly", repository, id, e);
                exceptionWhileCommitting.addSuppressed(e);
            }
        }

        //Inform the transaction leader of the outcome of our attempt to commit
        try {
            leader.onCommitted(exceptionWhileCommitting);
        } catch (RuntimeException e) {
            log.warn("[{}] Failed to notify the leader that transaction {} {}",
                    repository, id, (exceptionWhileCommitting == null ? "was committed" : "failed to commit"), e);
        }
        if (exceptionWhileCommitting == null) {
            transitionTo(COMMITTED);
        } else {
            transitionTo(ROLLED_BACK);
        }

        if (commitException != null) {
            //PreparedAction.commit() failed; propagate the exception
            throw commitException;
        }
    }

    private void internalRollBack() {
        RuntimeException rollbackException = null;
        if (action != null) {
            log.debug("[{}] Rolling back transaction {} (state was {}))", action.getRepository(), id, state);
            try {
                action.rollBack();
            } catch (RuntimeException e) {
                rollbackException = e;
            }
        }

        //Notify the leader that the local transactional operation completed rolling back the transaction
        //(successfully if rollbackException == null).
        try {
            leader.onRolledBack(rollbackException);
        } catch (RuntimeException e) {
            log.warn("[{}] Failed to notify the leader that transaction {} was rolled back",
                    action == null ? "unknown" : action.getRepository(), id, e);
        }
        transitionTo(ROLLED_BACK);
        if (rollbackException != null) {
            //PreparedAction.rollBack failed; propagate the exception
            throw rollbackException;
        }
    }

    private void invokeListener(TransactionListener listener) {
        if (state == COMMITTING) {
            try {
                listener.beforeCommit(this);
            } catch (RuntimeException e) {
                log.warn("{}.beforeCommit(..) failed for transaction {}",
                        listener.getClass().getName(), id, e);
            }
            return;
        }
        if (state == COMMITTED) {
            try {
                listener.afterCommit(this);
            } catch (RuntimeException e) {
                log.warn("{}.afterCommit(..) failed for transaction {}",
                        listener.getClass().getName(), id, e);
            }
        }
        try {
            listener.afterCompletion(this);
        } catch (RuntimeException e) {
            log.warn("{}.afterCompletion(..) failed for transaction {}",
                    listener.getClass().getName(), id, e);
        }
    }

    private void invokeListeners() {
        listeners.stream()
                .sorted(Comparator.comparing(TransactionListener::getOrder))
                .forEach(this::invokeListener);
    }

    private boolean shouldInvokeListeners() {
        return state.isFinished() || state == COMMITTING;
    }

    private void transitionTo(TransactionState nextState) {
        if (state == nextState) {
            return;
        }
        if (state.canTransitionTo(nextState)) {
            log.trace("Transaction {} transitioned from {} to {}", id, state, nextState);
            state = nextState;
            if (shouldInvokeListeners()) {
                invokeListeners();
            }
        } else if (nextState == PREPARED && (state == COMMITTING || state == ROLLING_BACK)) {
            log.debug("Ignoring transition to PREPARED for transaction {}; the leader has already sent outcome {}",
                    id, state);
        } else if (nextState == ROLLING_BACK && state == ROLLED_BACK) {
            log.debug("Ignoring transition to ROLLING_BACK for already ROLLED_BACK transaction {}", id);
        } else {
            throw new IllegalTransactionStateException(id, state, nextState);
        }
    }
}
