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 org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;

import static com.google.common.base.Preconditions.checkState;
import static java.util.Objects.requireNonNull;
import static java.util.Optional.empty;
import static java.util.Optional.ofNullable;
import static org.junit.Assert.assertEquals;

public class StubTransactionManager implements TransactionManager, TestRule {

    private static final MeshNode LOCAL_NODE = new MeshNode.Builder("local")
            .name("local")
            .local(true)
            .rpcUrl("http://should-never-be-called.nosuchdomain:7777")
            .build();

    private final MeshNode leader;

    private int rolledBackCount;
    private int committedCount;
    private boolean coordinatorVote;
    private SimpleTransaction lastTransaction;
    private SimpleTransaction transaction;

    public StubTransactionManager() {
        this(LOCAL_NODE);
    }

    public StubTransactionManager(MeshNode leader) {
        this.leader = requireNonNull(leader, "leader");
        reset();
    }

    @Override
    public boolean addListener(@Nonnull TransactionListener listener) {
        return false;
    }

    @Override
    public Statement apply(Statement base, Description description) {
        return new Statement() {

            @Override
            public void evaluate() throws Throwable {
                try {
                    base.evaluate();
                } finally {
                    reset();
                }
            }
        };
    }

    public void assertTransactions(int committed, int rolledBack) {
        String details = " (committed = " + committedCount + ", rolledBack = " + rolledBack + ")";
        assertEquals("Committed transactions do not match" + details, committed, committedCount);
        assertEquals("Rolled back transactions do not match" + details, rolledBack, rolledBackCount);
    }

    @Override
    public <T> T execute(@Nonnull Function<? super Transaction, T> callback) {
        transaction = transaction == null ? new SimpleTransaction() : new NestedTransaction(transaction);
        return invokeCallback(callback);
    }

    @Override
    public <T> T executeLocalOnly(@Nonnull Function<? super Transaction, T> callback) {
        SimpleTransaction suspendedTx = transaction;
        transaction = new SimpleTransaction();
        try {
            return invokeCallback(callback);
        } finally {
            transaction = suspendedTx;
        }
    }

    @Override
    public void executeLocalOnlyWithoutResult(@Nonnull Consumer<? super Transaction> callback) {
        executeLocalOnly(tx -> {
            callback.accept(tx);
            return null;
        });
    }

    @Override
    public void executeWithoutResult(@Nonnull Consumer<? super Transaction> callback) {
        execute(tx -> {
            callback.accept(tx);
            return null;
        });
    }

    public int getCommittedCount() {
        return committedCount;
    }

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

    public Transaction getLastTransaction() {
        return lastTransaction;
    }

    @Nullable
    public RpcTransaction getPreparedRpcTransaction() {
        return transaction == null ? null : transaction.getPreparedRpcTransaction();
    }

    public int getRolledBackCount() {
        return rolledBackCount;
    }

    @Nonnull
    @Override
    public PrepareResult prepared(@Nonnull PreparedAction operation) {
        if (transaction == null) {
            transaction = new SimpleTransaction();
        }
        return transaction.prepared(operation);
    }

    public void reset() {
        coordinatorVote = true;
        lastTransaction = null;
        transaction = null;
        committedCount = 0;
        rolledBackCount = 0;
    }

    @Nonnull
    @Override
    public Transaction start() {
        checkState(transaction == null, "Another transaction is already active!");
        transaction = new SimpleTransaction();
        return transaction;
    }

    @Nonnull
    @Override
    public Transaction start(@Nonnull StartTransactionParameters parameters) {
        throw new UnsupportedOperationException("not implemented");
    }

    public void stubCoordinatorVote(boolean success) {
        coordinatorVote = success;
    }

    @Override
    public RepositoryConsistencySummary verifyConsistency(@Nonnull String repositoryId) {
        return null;
    }

    private <T> T invokeCallback(Function<? super Transaction, T> callback) {
        RuntimeException exception = null;
        T result = null;
        SimpleTransaction tx = transaction;
        try {
            result = callback.apply(tx);
        } catch (RuntimeException e) {
            tx.setRollbackOnly();
            exception = e;
        } finally {
            try {
                tx.close();
            } catch (RuntimeException e) {
                if (exception == null) {
                    exception = e;
                } else {
                    exception.addSuppressed(e);
                }
            }
            if (tx.getState() == TransactionState.COMMITTED) {
                ++committedCount;
            } else {
                ++rolledBackCount;
            }
        }

        if (exception != null) {
            throw exception;
        }
        return result;
    }

    private class NestedTransaction extends SimpleTransaction {

        private final SimpleTransaction parent;

        NestedTransaction(SimpleTransaction parent) {
            this.parent = parent;
        }

        @Override
        public void close() {
            transaction = parent;
        }

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

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

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

        @Override
        public boolean isRollbackOnly() {
            return parent.isRollbackOnly();
        }

        @Nonnull
        @Override
        public Optional<SatiableConsumer<RpcStreamChunk>> newPackWriter(@Nonnull Repository repository) {
            return parent.newPackWriter(repository);
        }

        @Override
        public void setRollbackOnly() {
            parent.setRollbackOnly();
        }

        @Override
        RpcTransaction getPreparedRpcTransaction() {
            return parent.getPreparedRpcTransaction();
        }
    }

    private class SimpleTransaction implements Transaction {

        private final List<TransactionListener> synchronizations = new ArrayList<>();

        private PreparedAction prepared;
        private RpcTransaction preparedRpcTransaction;
        private boolean rollbackOnly;
        private TransactionState state = TransactionState.INITIAL;

        @Override
        public void addListener(@Nonnull TransactionListener listener) {
            synchronizations.add(listener);
        }

        @Override
        public void close() {
            if (state.isFinished()) {
                return;
            }

            if (rollbackOnly) {
                if (prepared != null) {
                    prepared.rollBack();
                }
                state = TransactionState.ROLLED_BACK;
                synchronizations.forEach(sync -> sync.afterCompletion(this));
            } else {
                if (prepared != null) {
                    prepared.commit();
                }
                state = TransactionState.COMMITTED;
                synchronizations.forEach(sync -> {
                    sync.afterCommit(this);
                    sync.afterCompletion(this);
                });
            }
            prepared = null;
            lastTransaction = this;
            transaction = null;
        }

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

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

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

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

        @Override
        public boolean isRollbackOnly() {
            return rollbackOnly;
        }

        @Override
        public void setRollbackOnly() {
            rollbackOnly = true;
        }

        @Nonnull
        @Override
        public PrepareResult prepared(@Nonnull PreparedAction action) {
            prepared = action;
            if (!coordinatorVote) {
                rollbackOnly = true;
            }
            if (rollbackOnly) {
                state = TransactionState.ROLLING_BACK;
                return PrepareResult.FAILURE;
            } else {
                preparedRpcTransaction = prepared.populateTransaction(RpcTransaction.newBuilder()
                        .setTimestamp(System.currentTimeMillis())
                        .setId(getId())
                        .setRepository(prepared.getRepository().getId()))
                        .build();
                state = TransactionState.COMMITTING;
                return PrepareResult.SUCCESS;
            }
        }

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

        @Override
        public <T> T rollBackIfThrows(@Nonnull Supplier<T> action) {
            requireNonNull(action, "action");
            try {
                return action.get();
            } catch (RuntimeException e) {
                setRollbackOnly();
                throw e;
            }
        }

        @Nullable
        RpcTransaction getPreparedRpcTransaction() {
            return preparedRpcTransaction;
        }
    }
}
