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.config.ReplicationConfig;
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.junit.LogListener;
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.ManagedRepository;
import com.atlassian.bitbucket.mesh.repository.RepositoryAttributeManager;
import com.atlassian.bitbucket.mesh.repository.RepositoryHashService;
import com.atlassian.bitbucket.mesh.repository.RepositoryManager;
import com.google.common.collect.ImmutableSet;
import org.apache.commons.lang3.mutable.MutableObject;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;

import static com.atlassian.bitbucket.mesh.test.MockitoUtils.returnFirst;
import static java.util.Optional.empty;
import static java.util.Optional.of;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.Silent.class)
public class DefaultTransactionManagerTest {

    @Rule
    public final LogListener logListener = new LogListener(DefaultTransaction.class);

    private DefaultTransaction activeTransaction;
    @Captor
    private ArgumentCaptor<Throwable> exceptionCaptor;
    private Consumer<DefaultTransaction> executionCallback;
    @Mock
    private ExecutionManager executionManager;
    @Mock
    private MeshNode localNode;
    @Mock
    private LogicalClock logicalClock;
    @Mock
    private MeshClient meshClient;
    @Mock
    private MeshManager meshManager;
    @Mock
    private PreparedAction preparedAction;
    private TransactionLeader remoteLeader;
    @Mock
    private ManagedRepository repository;
    @Mock
    private RepositoryAttributeManager repositoryAttributeManager;
    @Mock
    private RepositoryHashService repositoryHashService;
    @Mock
    private RepositoryManager repositoryManager;
    @Mock
    private ExecutionLocal<DefaultTransaction> transactionLocal;
    private DefaultTransactionManager transactionManager;
    @Mock
    private WriteTokenRegistry writeTokenRegistry;

    @Before
    public void setUp() {
        when(executionManager.newExecutionLocal(any())).thenAnswer(invocation -> {
            executionCallback = invocation.getArgument(0);
            return transactionLocal;
        });
        when(transactionLocal.compute(any())).thenAnswer(invocation -> {
            Function<DefaultTransaction, DefaultTransaction> callback = invocation.getArgument(0);
            if (callback != null) {
                activeTransaction = callback.apply(activeTransaction);
            }
            return activeTransaction;
        });
        when(transactionLocal.get()).thenAnswer(invocation -> activeTransaction);
        when(transactionLocal.ifPresent(any())).thenAnswer(invocation -> {
            Consumer<? super DefaultTransaction> callback = invocation.getArgument(0);
            if (activeTransaction != null) {
                callback.accept(activeTransaction);
                return true;
            }
            return false;
        });
        when(transactionLocal.removeIf(any())).thenAnswer(invocation -> {
            DefaultTransaction result = null;
            if (activeTransaction != null) {
                Predicate<DefaultTransaction> predicate = invocation.getArgument(0);
                if (predicate.test(activeTransaction)) {
                    result = activeTransaction;
                    activeTransaction = null;
                }
            }
            return result;
        });
        when(transactionLocal.set(any())).thenAnswer(invocation -> {
            activeTransaction = invocation.getArgument(0);
            return true;
        });

        when(localNode.isLocal()).thenReturn(true);
        when(localNode.getId()).thenReturn("local-node");
        when(meshManager.getLocalNode()).thenReturn(localNode);
        when(meshManager.getNodes()).thenReturn(ImmutableSet.of(localNode));

        when(preparedAction.getRepository()).thenReturn(repository);
        when(preparedAction.populateTransaction(any())).thenAnswer(returnFirst());
        when(repository.getId()).thenReturn("p/0001/deadbeef/repo");

        ReplicationConfig replicationConfig = mock(ReplicationConfig.class);
        when(replicationConfig.getQuorumMode()).thenReturn(QuorumMode.UNANIMOUS);
        ApplicationConfig config = mock(ApplicationConfig.class);
        when(config.getReplication()).thenReturn(replicationConfig);

        when(repositoryAttributeManager.getVersion(anyString())).thenReturn(3L);
        when(repositoryAttributeManager.incrementAndGetVersion(anyString())).thenReturn(4L);
        when(repositoryHashService.getContentHash(repository)).thenReturn("content-hash");

        remoteLeader = mockRemoteTransactionLeader();

        transactionManager = new DefaultTransactionManager(
                config, mock(EventPublisher.class), executionManager, logicalClock, meshClient, meshManager,
                repositoryAttributeManager, repositoryHashService, repositoryManager, mock(SidebandManager.class),
                writeTokenRegistry);
    }

    @Test
    public void testAddSynchronizationWhenNoTransactionIsActive() {
        TransactionListener synchronization = mock(TransactionListener.class);
        assertFalse(transactionManager.addListener(synchronization));
        verifyNoInteractions(synchronization);
    }

    @Test
    public void testAddSynchronizationWhenTransactionIsActive() {
        TransactionListener synchronization = mock(TransactionListener.class);
        Transaction transaction = transactionManager.execute(tx -> {
            assertTrue(transactionManager.addListener(synchronization));
            return tx;
        });

        //The transaction should have rolled back because no operation was prepared
        verify(synchronization).afterCompletion(transaction);
        verify(synchronization, never()).afterCommit(any());
    }

    @Test
    public void testExecute() {
        Transaction transaction = transactionManager.execute(tx -> {
            //Verify that the active transaction was set up in the ExecutionContext
            assertNotNull(activeTransaction);
            assertEquals(of(activeTransaction), transactionManager.getCurrentTransaction());
            assertEquals(TransactionState.INITIAL, tx.getState());
            assertFalse(tx.isRollbackOnly());
            assertNotNull(tx.getId());
            return tx;
        });

        //The current transaction should be cleared again
        assertNull(activeTransaction);
        assertEquals(empty(), transactionManager.getCurrentTransaction());

        //Verify that the transaction was rolled back because nothing was prepared
        assertNotNull(transaction);
        assertEquals(TransactionState.ROLLED_BACK, transaction.getState());
        //No interactions with replicas expected either
        verifyNoInteractions(meshClient);
    }

    @Test
    public void testExecuteDoesNotRollBackAfterExceptionIfTransactionAlreadyActive() {
        transactionManager.execute(outerTx -> {
            assertThrows(UnsupportedOperationException.class, () -> transactionManager.execute(innerTx -> {
                assertSame(outerTx, innerTx);
                throw new UnsupportedOperationException("testing exception handling");
            }));

            assertNotNull(outerTx);
            assertFalse(outerTx.isRollbackOnly());
            assertEquals(TransactionState.INITIAL, outerTx.getState());
            return null;
        });
    }

    @Test
    public void testExecuteLocalOnlyDoesNotRollbackSuspendedTransaction() {
        transactionManager.execute(outerTx -> {
            assertThrows(IllegalStateException.class, () -> transactionManager.executeLocalOnly(innerTx -> {
                throw new IllegalStateException("testing exception handling");
            }));

            assertEquals(outerTx, transactionManager.getCurrentTransaction().orElse(null));
            assertFalse(outerTx.isRollbackOnly());
            assertEquals(TransactionState.INITIAL, outerTx.getState());
            return null;
        });
    }

    @Test
    public void testExecuteLocalOnlyRollback() {
        TransactionListener listener = mock(TransactionListener.class);

        transactionManager.execute(outerTx -> {
            transactionManager.executeLocalOnly(innerTx -> {
                assertNotSame(outerTx, innerTx);
                assertEquals(innerTx, transactionManager.getCurrentTransaction().orElse(null));

                innerTx.addListener(listener);
                innerTx.setRollbackOnly();

                //Preparing a transactional operation should result in a rollback result
                assertFalse(innerTx.prepared(preparedAction).isSuccess());

                //But the rollback should not affect the suspended transaction
                assertEquals(TransactionState.INITIAL, outerTx.getState());
                return null;
            });

            assertEquals(outerTx, transactionManager.getCurrentTransaction().orElse(null));
            assertFalse(outerTx.isRollbackOnly());
            assertEquals(TransactionState.INITIAL, outerTx.getState());

            verify(preparedAction).rollBack();
            verify(listener).afterCompletion(any());
            verify(listener, never()).afterCommit(any());

            return null;
        });
    }

    @Test
    public void testExecuteLocalOnlySuspendsActiveTransaction() {
        transactionManager.execute(outerTx -> {
            transactionManager.executeLocalOnly(innerTx -> {
                assertNotSame(outerTx, innerTx);
                assertEquals(localNode, innerTx.getLeader());
                assertEquals(innerTx, transactionManager.getCurrentTransaction().orElse(null));

                assertEquals(TransactionState.INITIAL, outerTx.getState());
                return null;
            });

            assertEquals(outerTx, transactionManager.getCurrentTransaction().orElse(null));
            assertFalse(outerTx.isRollbackOnly());
            assertEquals(TransactionState.INITIAL, outerTx.getState());
            return null;
        });
    }

    @Test
    public void testExecuteLocalRollbackCommit() {
        TransactionListener listener = mock(TransactionListener.class);

        transactionManager.execute(outerTx -> {
            transactionManager.executeLocalOnly(innerTx -> {
                assertNotSame(outerTx, innerTx);
                assertEquals(innerTx, transactionManager.getCurrentTransaction().orElse(null));

                innerTx.addListener(listener);

                //Preparing a transactional operation should result in a success result
                assertTrue(innerTx.prepared(preparedAction).isSuccess());

                //But the prepared change should not affect the suspended transaction
                assertEquals(TransactionState.INITIAL, outerTx.getState());
                return null;
            });

            assertEquals(outerTx, transactionManager.getCurrentTransaction().orElse(null));
            assertEquals(TransactionState.INITIAL, outerTx.getState());

            verify(preparedAction).commit();
            verify(listener).afterCompletion(any());
            verify(listener).afterCommit(any());

            return null;
        });
    }

    @Test
    public void testExecuteLocalWhenNoTransactionIsActive() {
        transactionManager.executeLocalOnly(transaction -> {
            assertEquals(transaction, transactionManager.getCurrentTransaction().orElse(null));

            assertEquals(TransactionState.INITIAL, transaction.getState());
            return null;
        });

        assertNull(transactionManager.getCurrentTransaction().orElse(null));
    }

    @Test
    public void testExecuteRollsBackAfterSetRollbackOnlyIfTransactionAlreadyActive() {
        transactionManager.execute(outerTx -> {
            transactionManager.execute(innerTx -> {
                assertSame(outerTx, innerTx);
                innerTx.setRollbackOnly();

                assertTrue(outerTx.isRollbackOnly());
                assertEquals(TransactionState.ROLLING_BACK, outerTx.getState());
                return null;
            });

            assertNotNull(outerTx);
            assertTrue(outerTx.isRollbackOnly());
            assertEquals(TransactionState.ROLLING_BACK, outerTx.getState());
            return null;
        });
    }

    @Test
    public void testExecuteRollsBackWhenCallbackThrows() {
        MutableObject<Transaction> transaction = new MutableObject<>();
        assertThrows(UnsupportedOperationException.class, () -> transactionManager.execute(tx -> {
            //Verify that the active transaction was set up in the ExecutionContext
            assertNotNull(activeTransaction);
            transaction.setValue(tx);

            throw new UnsupportedOperationException("testing exception handling");
        }));

        assertNull(activeTransaction);
        assertEquals(empty(), transactionManager.getCurrentTransaction());

        //Verify the transaction was rolled back
        assertNotNull(transaction.getValue());
        assertEquals(TransactionState.ROLLED_BACK, transaction.getValue().getState());
        //No interactions with replicas expected either
        verifyNoInteractions(meshClient);
    }

    @Test
    public void testExecuteWithoutResult() {
        MutableObject<Transaction> transaction = new MutableObject<>();
        transactionManager.executeWithoutResult(tx -> {
            //Verify that the active transaction was set up in the ExecutionContext
            assertNotNull(activeTransaction);
            assertEquals(of(activeTransaction), transactionManager.getCurrentTransaction());
            assertEquals(TransactionState.INITIAL, tx.getState());
            assertFalse(tx.isRollbackOnly());
            assertNotNull(tx.getId());

            transaction.setValue(tx);
        });

        assertNull(activeTransaction);
        assertEquals(empty(), transactionManager.getCurrentTransaction());

        //Verify that the transaction was rolled back because no operation was prepared
        assertNotNull(transaction.getValue());
        assertEquals(TransactionState.ROLLED_BACK, transaction.getValue().getState());
        //No interactions with replicas expected either
        verifyNoInteractions(meshClient);
    }

    @Test
    public void testFollowerTransactionCommitsWhenLeaderVotesYes() {
        TransactionListener synchronization1 = mock(TransactionListener.class);
        TransactionListener synchronization2 = mock(TransactionListener.class);
        withFollowerTransaction(PrepareResult.SUCCESS, tx -> {
            //Verify that an active transaction has been set up
            assertNotNull(activeTransaction);

            tx.addListener(synchronization1);
            tx.prepared(preparedAction);
            transactionManager.addListener(synchronization2);
        });

        //Verify that the active transaction has been cleared
        assertNull(activeTransaction);
        verify(preparedAction).commit();

        //Verify that the synchronization callbacks have been invoked
        verify(synchronization1).afterCommit(any());
        verify(synchronization1).afterCompletion(any());
        verify(synchronization2).afterCommit(any());
        verify(synchronization2).afterCompletion(any());
    }

    @Test
    public void testFollowerTransactionRollsBackWhenLeaderVotesNo() {
        withFollowerTransaction(PrepareResult.FAILURE, tx -> tx.prepared(preparedAction));

        verify(preparedAction).rollBack();
        verify(remoteLeader).onPrepared(same(repository), any());
        verify(remoteLeader).onRolledBack(null);
    }

    @Test
    public void testFollowerTransactionRollsBackWhenLocalOperationThrows() {
        assertThrows(IllegalStateException.class, () -> withFollowerTransaction(PrepareResult.SUCCESS, tx -> {
            throw new IllegalStateException("testing");
        }));

        //No exception is expected on the onRolledBack callback here because an exception should only provided when the
        //rollback _itself_ fails, not if an exception triggers a rollback of the local transaction.
        verify(remoteLeader).onRolledBack(null);
    }

    @Test
    public void testGetActiveTransactionWhenNoneActive() {
        assertEquals(empty(), transactionManager.getCurrentTransaction());
    }

    @Test
    public void testGetActiveTransactionWhenOneIsActive() {
        transactionManager.executeWithoutResult(tx ->
                assertEquals(of(tx), transactionManager.getCurrentTransaction()));
    }

    @Test
    public void testPrepareWithAnActiveTransaction() {
        transactionManager.executeWithoutResult(tx -> {
            assertTrue(transactionManager.prepared(preparedAction).isSuccess());

            assertEquals(TransactionState.COMMITTING, tx.getState());
            //Operation should not have committed or rolled back yet
            verify(preparedAction, never()).rollBack();
            verify(preparedAction, never()).commit();
        });

        //Once the transaction commits, PreparedAction.commit should be called
        verify(preparedAction).commit();
    }

    @Test
    public void testPrepareWithNoActiveTransactionAndNoExecutionContextImmediatelyCompletes() {
        when(transactionLocal.compute(any())).thenReturn(null); //Simulate no active execution context

        assertTrue(transactionManager.prepared(preparedAction).isSuccess());

        verify(preparedAction).commit();
    }

    @Test
    public void testPrepareWithNoActiveTransactionButAnActiveExecutionContextStartsOne() {
        assertTrue(transactionManager.prepared(preparedAction).isSuccess());

        Transaction transaction = transactionManager.getCurrentTransaction().orElse(null);

        assertNotNull(transaction);
        assertEquals(TransactionState.COMMITTING, transaction.getState());
        //Operation should not have committed or rolled back yet
        verify(preparedAction, never()).rollBack();
        verify(preparedAction, never()).commit();

        //Once the execution context closes, PreparedAction.commit should be called
        executionCallback.accept((DefaultTransaction) transaction); //Simulate execution context cleanup
        verify(preparedAction).commit();
    }

    @Test
    public void testStartDoesNotThrowWhenSameTransactionIsActive() {
        try (Transaction outer = transactionManager.start(new StartTransactionParameters.Builder(remoteLeader)
                .id("tx-id")
                .build())) {
            assertSame(outer, transactionManager.start(new StartTransactionParameters.Builder(remoteLeader)
                    .id("tx-id")
                    .build()));
        }
    }

    @Test
    public void testStartThrowsWhenAnotherTransactionIsActive() {
        try (Transaction ignored = transactionManager.start(new StartTransactionParameters.Builder(remoteLeader)
                .id("outer")
                .build())) {
            assertThrows(IllegalStateException.class,
                    () -> transactionManager.start(new StartTransactionParameters.Builder(remoteLeader)
                            .id("inner-tx")
                            .build()));
        }
    }

    @Test
    public void testStartWithoutParameters() {
        Transaction tx;
        try (Transaction transaction = (tx = transactionManager.start())) {
            assertEquals(transaction, transactionManager.getCurrentTransaction().orElse(null));
            assertEquals(TransactionState.INITIAL, transaction.getState());
            assertEquals(localNode, transaction.getLeader());
        }
        //Verify that the transaction is rolled back and that there is no transaction active
        assertFalse(transactionManager.getCurrentTransaction().isPresent());
        assertEquals(TransactionState.ROLLED_BACK, tx.getState());
    }

    @Test
    public void testStartWithoutParametersThrowsWhenAnotherTransactionIsActive() {
        transactionManager.executeWithoutResult(
                tx -> assertThrows(IllegalStateException.class, transactionManager::start));
    }

    @Test
    public void testTransactionIsClosedOnExecutionContextCleanup() {
        DefaultTransaction transaction = transactionManager.start(new StartTransactionParameters.Builder(remoteLeader)
                .id("tx-id")
                .build());

        //Verify that the transaction has been started and is registered with the ExecutionLocal
        assertNotNull(activeTransaction);
        assertNotNull(executionCallback);
        assertEquals(TransactionState.INITIAL, transaction.getState());

        //Simulate the ExecutionContext being cleaned up
        executionCallback.accept(transaction);
        assertNull(activeTransaction);
        assertEquals(TransactionState.ROLLED_BACK, transaction.getState());
        verify(remoteLeader).onRolledBack(null);
    }

    @Test
    public void testTransactionThatThrowsFromPreparedActionCommit() {
        doThrow(new UnsupportedOperationException("testing exception handling"))
                .when(preparedAction).commit();

        assertThrows(UnsupportedOperationException.class, () -> {
            withFollowerTransaction(PrepareResult.SUCCESS, tx -> {
                assertTrue(tx.prepared(preparedAction).isSuccess());

                //The transaction should now be COMMITTING
                assertEquals(TransactionState.COMMITTING, tx.getState());
            });
        });

        verify(remoteLeader).onCommitted(exceptionCaptor.capture());
        assertThat(exceptionCaptor.getValue(), instanceOf(UnsupportedOperationException.class));
    }

    @Test
    public void testTransactionThatThrowsFromPreparedActionRollback() {
        doThrow(new UnsupportedOperationException("testing exception handling"))
                .when(preparedAction).rollBack();

        assertThrows(UnsupportedOperationException.class, () -> {
            withFollowerTransaction(PrepareResult.FAILURE, tx -> {
                assertFalse(tx.prepared(preparedAction).isSuccess());

                //The transaction should now be ROLLING_BACK
                assertEquals(TransactionState.ROLLING_BACK, tx.getState());
            });
        });

        verify(remoteLeader).onRolledBack(exceptionCaptor.capture());
        assertThat(exceptionCaptor.getValue(), instanceOf(UnsupportedOperationException.class));
    }

    @Test
    public void testTransactionThatThrowsWhenCommittingButBeforeClose() {
        assertThrows(UnsupportedOperationException.class, () -> {
            withFollowerTransaction(PrepareResult.SUCCESS, tx -> {
                assertTrue(tx.prepared(preparedAction).isSuccess());

                //The transaction should now be COMMITTING
                assertEquals(TransactionState.COMMITTING, tx.getState());

                throw new UnsupportedOperationException("Testing exception handling");
            });
        });

        verify(remoteLeader).onCommitted(exceptionCaptor.capture());
        assertThat(exceptionCaptor.getValue(), instanceOf(UnsupportedOperationException.class));

        logListener.assertError("Rolling back COMMITTING transaction");
    }

    @Test
    public void testTransactionThatThrowsWhenRollingBackButBeforeClose() {
        assertThrows(UnsupportedOperationException.class, () -> {
            withFollowerTransaction(PrepareResult.FAILURE, tx -> {
                assertFalse(tx.prepared(preparedAction).isSuccess());

                //The transaction should now be ROLLING_BACK
                assertEquals(TransactionState.ROLLING_BACK, tx.getState());

                throw new UnsupportedOperationException("Testing exception handling");
            });
        });

        verify(remoteLeader).onRolledBack(null);
    }

    private static TransactionLeader mockRemoteTransactionLeader() {
        MeshNode node = mock(MeshNode.class);
        when(node.getId()).thenReturn("remote");
        TransactionLeader leader = mock(TransactionLeader.class);
        when(leader.getNode()).thenReturn(node);
        return leader;
    }

    private void withFollowerTransaction(PrepareResult result, Consumer<Transaction> callback) {
        when(remoteLeader.onPrepared(any(), any())).thenReturn(CompletableFuture.completedFuture(result));

        try (Transaction tx = transactionManager.start(new StartTransactionParameters.Builder(remoteLeader)
                .id(Transaction.formatId("remote", "12345"))
                .build())) {
            tx.rollBackIfThrows(() -> callback.accept(tx));
        }
    }
}