package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.junit.LogListener;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.google.common.base.Throwables;
import io.grpc.Context;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.InOrder;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.util.concurrent.*;

import static com.atlassian.bitbucket.mesh.test.MockitoUtils.returnFirst;
import static java.util.concurrent.CompletableFuture.completedFuture;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.same;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class DefaultTransactionTest {

    private static final String TRANSACTION_ID = Transaction.formatId("leader", "tx-id");

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

    private ExecutorService executorService;
    @Mock
    private TransactionLeader leader;
    @Mock
    private PreparedAction preparedAction;
    @Mock
    private Repository repository;
    @Mock
    private TransactionListener listener;
    private DefaultTransaction transaction;

    @Before
    public void setUp() {
        transaction = new DefaultTransaction(TRANSACTION_ID, leader);
        transaction.addListener(listener);

        when(preparedAction.getRepository()).thenReturn(repository);
        when(preparedAction.populateTransaction(any())).thenAnswer(returnFirst());
        when(repository.getId()).thenReturn("repository-id");
        when(leader.onPrepared(same(repository), any())).thenReturn(completedFuture(PrepareResult.SUCCESS));

        executorService = Executors.newFixedThreadPool(1);
    }

    @After
    public void tearDown() throws Exception {
        executorService.shutdownNow();
        executorService.awaitTermination(1, TimeUnit.SECONDS);
    }

    @Test
    public void testAddListenerAfterCommit() {
        transaction.prepared(preparedAction);
        transaction.close();

        assertEquals(TransactionState.COMMITTED, transaction.getState());
        TransactionListener sync = mock(TransactionListener.class);

        transaction.addListener(sync);

        verify(sync, never()).beforeCommit(transaction);
        verify(sync).afterCommit(transaction);
        verify(sync).afterCompletion(transaction);
    }

    @Test
    public void testAddListenerAfterRollback() {
        when(leader.onPrepared(same(repository), any())).thenReturn(completedFuture(PrepareResult.FAILURE));
        transaction.prepared(preparedAction);
        transaction.close();

        assertEquals(TransactionState.ROLLED_BACK, transaction.getState());
        TransactionListener sync = mock(TransactionListener.class);

        transaction.addListener(sync);

        verify(sync, never()).beforeCommit(transaction);
        verify(sync, never()).afterCommit(transaction);
        verify(sync).afterCompletion(transaction);
    }

    @Test
    public void testCommitThatThrows() {
        RuntimeException exception = new RuntimeException("testing commit that throws");
        doThrow(exception).when(preparedAction).commit();

        transaction.prepared(preparedAction);
        assertThrows(RuntimeException.class, () -> transaction.close());

        verify(listener).beforeCommit(transaction);
        verify(leader).onCommitted(exception);
        verify(listener).afterCompletion(transaction);
        assertEquals(TransactionState.ROLLED_BACK, transaction.getState());
    }

    @Test
    public void testCloseWithoutPrepareRollsBack() {
        transaction.close();

        assertTrue(transaction.isRollbackOnly());
        assertEquals(TransactionState.ROLLED_BACK, transaction.getState());

        InOrder inOrder = inOrder(leader, listener);
        inOrder.verify(leader).onRolledBack(null);
        inOrder.verify(listener).afterCompletion(transaction);
        inOrder.verifyNoMoreInteractions();
    }

    @Test
    public void testInitialState() {
        assertEquals(TransactionState.INITIAL, transaction.getState());
        assertEquals(TRANSACTION_ID, transaction.getId());
    }

    @Test
    public void testPrepareAfterCloseReturnsFailure() {
        transaction.close();

        PrepareResult result = transaction.prepared(preparedAction);

        assertFalse(result.isSuccess());
        verify(leader).onRolledBack(null);
        verifyNoMoreInteractions(leader);
    }

    @Test
    public void testPrepareAfterRollbackReturnsFailure() {
        transaction.setRollbackOnly();
        assertEquals(TransactionState.ROLLING_BACK, transaction.getState());

        PrepareResult result = transaction.prepared(preparedAction);

        //Verify that prepare failed, but no onRolledBack callback has been sent to the leader yet.
        assertFalse(result.isSuccess());
        verifyNoInteractions(leader);
        assertEquals(TransactionState.ROLLING_BACK, transaction.getState());

        //Closing the transaction should trigger the onRolledBack callback to the leader.
        transaction.close();
        verify(leader).onRolledBack(null);
        verifyNoMoreInteractions(leader);
    }

    @Test
    public void testPrepareAlreadyPreparedTransactionThrows() {
        transaction.prepared(preparedAction);
        assertThrows(IllegalStateException.class, () -> transaction.prepared(preparedAction));
    }

    @Test
    public void testPrepareAndCommit() {
        TransactionListener listener1 = mock(TransactionListener.class);
        TransactionListener listener2 = mock(TransactionListener.class);
        when(listener1.getOrder()).thenReturn(10); // after listener2

        //Prepare the transaction
        transaction.addListener(listener1);
        PrepareResult result = transaction.prepared(preparedAction);
        transaction.addListener(listener2);

        assertTrue(result.isSuccess());
        assertEquals(TransactionState.COMMITTING, transaction.getState());
        InOrder inOrder = inOrder(preparedAction, leader, listener1, listener2);
        inOrder.verify(preparedAction).getRepository();
        inOrder.verify(preparedAction).populateTransaction(any());
        inOrder.verify(leader).onPrepared(same(repository), any());
        inOrder.verify(listener1).beforeCommit(transaction);
        inOrder.verify(listener2).beforeCommit(transaction);
        inOrder.verifyNoMoreInteractions();

        //Close the transaction to commit it
        transaction.close();
        assertEquals(TransactionState.COMMITTED, transaction.getState());

        inOrder.verify(preparedAction).commit();
        inOrder.verify(leader).onCommitted(null);
        inOrder.verify(listener2).afterCommit(transaction);
        inOrder.verify(listener2).afterCompletion(transaction);
        inOrder.verify(listener1).afterCommit(transaction);
        inOrder.verify(listener1).afterCompletion(transaction);
        inOrder.verifyNoMoreInteractions();
    }

    @Test
    public void testPrepareAndCommitWhenLeaderThrowsFromOnCommitted() {
        doThrow(new RuntimeException("simulating leader error")).when(leader).onCommitted(null);

        PrepareResult result = transaction.prepared(preparedAction);
        assertTrue(result.isSuccess());

        transaction.close();

        verify(leader).onCommitted(null);
        //Verify that the TransactionListeners are still invoked when the leader throws
        verify(listener).beforeCommit(transaction);
        verify(listener).afterCommit(transaction);
        verify(listener).afterCompletion(transaction);
        logListener.assertWarn("Failed to notify the leader that ");
    }

    @Test
    public void testPrepareReturnsFailureWhenStateIsNotRolledBackOnly() {
        transaction.close();

        Future<?> result = executorService.submit(() -> {
            PrepareResult prepareResult = transaction.prepared(preparedAction);

            assertEquals(TransactionState.ROLLED_BACK, transaction.getState());
            assertFalse(prepareResult.isSuccess());
        });

        try {
            result.get();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } catch (ExecutionException e) {
            Throwables.propagateIfPossible(e.getCause(), RuntimeException.class);
            throw new RuntimeException(e.getCause());
        }
    }

    @Test
    public void testPrepareRollsBackWhenLeaderRollsBack() {
        TransactionListener listener1 = mock(TransactionListener.class);
        TransactionListener listener2 = mock(TransactionListener.class);
        when(listener1.getOrder()).thenReturn(10); // after listener2

        //Mock the leader voting to roll back
        when(leader.onPrepared(same(repository), any())).thenReturn(completedFuture(PrepareResult.FAILURE));

        //Prepare the transaction
        transaction.addListener(listener1);
        PrepareResult result = transaction.prepared(preparedAction);
        transaction.addListener(listener2);

        assertFalse(result.isSuccess());
        assertEquals(TransactionState.ROLLING_BACK, transaction.getState());
        InOrder inOrder = inOrder(preparedAction, leader, listener1, listener2);
        inOrder.verify(preparedAction).getRepository();
        inOrder.verify(preparedAction).populateTransaction(any());
        inOrder.verify(leader).onPrepared(same(repository), any());
        inOrder.verify(leader).getNode(); //Debug logging
        inOrder.verifyNoMoreInteractions();

        //Close the transaction to mark it as fully rolled back
        transaction.close();
        assertEquals(TransactionState.ROLLED_BACK, transaction.getState());

        inOrder.verify(preparedAction).rollBack();
        inOrder.verify(leader).onRolledBack(null);
        inOrder.verify(listener2).afterCompletion(transaction);
        inOrder.verify(listener1).afterCompletion(transaction);
        inOrder.verifyNoMoreInteractions();
    }

    @Test(timeout = 2000L)
    public void testPrepareRollsBackWhenLeaderTimesOut() {
        ScheduledExecutorService executor = mock(ScheduledExecutorService.class);
        //Mock a leader that does not return
        when(leader.onPrepared(same(repository), any())).thenReturn(new CompletableFuture<>());
        //But call it from a gRPC context with an expired deadline to ensure that the transaction does
        //not wait for the leader to respond
        Context.current().withDeadlineAfter(0L, TimeUnit.MILLISECONDS, executor).run(() -> {
            PrepareResult result = transaction.prepared(preparedAction);

            assertFalse(result.isSuccess());
            assertTrue(transaction.isRollbackOnly());
            assertEquals(TransactionState.ROLLING_BACK, transaction.getState());

            transaction.close();
            assertEquals(TransactionState.ROLLED_BACK, transaction.getState());
        });
        logListener.assertWarn("Timed out waiting for outcome of transaction");
    }

    @Test
    public void testPrepareRollsBackWhenLeaderThrows() {
        //Mock the leader throwing an error
        CompletableFuture<PrepareResult> leaderFuture = new CompletableFuture<>();
        leaderFuture.completeExceptionally(new RuntimeException("simulating connection error"));
        when(leader.onPrepared(same(repository), any())).thenReturn(leaderFuture);
        doThrow(new RuntimeException("simulating connection error")).when(leader).onRolledBack(null);

        PrepareResult result = transaction.prepared(preparedAction);

        assertFalse(result.isSuccess());
        assertTrue(transaction.isRollbackOnly());
        assertEquals(TransactionState.ROLLING_BACK, transaction.getState());
        verify(leader).onPrepared(same(repository), any());
        logListener.assertWarn("Error while waiting for outcome of transaction");

        transaction.close();
        assertEquals(TransactionState.ROLLED_BACK, transaction.getState());
        verify(leader).onRolledBack(null);
        //Verify that the TransactionListeners are still invoked when the leader throws
        verify(listener).afterCompletion(transaction);
        verify(listener, never()).afterCommit(any());
        logListener.assertWarn("Failed to notify the leader");
    }

    @Test
    public void testRollBackThatThrows() {
        RuntimeException exception = new RuntimeException("testing rollBack that throws");
        doThrow(exception).when(preparedAction).rollBack();
        when(leader.onPrepared(same(repository), any())).thenReturn(completedFuture(PrepareResult.FAILURE));

        transaction.prepared(preparedAction);
        assertThrows(RuntimeException.class, () -> transaction.close());

        verify(leader).onRolledBack(exception);
        verify(listener).afterCompletion(transaction);
        assertEquals(TransactionState.ROLLED_BACK, transaction.getState());
    }

    @Test
    public void testListenersThatThrowAreHandled() {
        TransactionListener listener1 = mock(TransactionListener.class, invocation -> {
            String methodName = invocation.getMethod().getName();
            if (methodName.startsWith("after") || methodName.startsWith("before")) {
                throw new UnsupportedOperationException("listener 1 error");
            }
            return invocation.callRealMethod();
        });
        TransactionListener listener2 = mock(TransactionListener.class, invocation -> {
            String methodName = invocation.getMethod().getName();
            if (methodName.startsWith("after") || methodName.startsWith("before")) {
                throw new UnsupportedOperationException("listener 1 error");
            }
            return invocation.callRealMethod();
        });
        transaction.addListener(listener1);
        transaction.addListener(listener2);

        transaction.prepared(preparedAction);
        transaction.close();

        verify(listener).beforeCommit(transaction);
        verify(listener).afterCommit(transaction);
        verify(listener).afterCompletion(transaction);
        verify(listener1).beforeCommit(transaction);
        verify(listener1).afterCommit(transaction);
        verify(listener1).afterCompletion(transaction);
        verify(listener2).beforeCommit(transaction);
        verify(listener2).afterCommit(transaction);
        verify(listener2).afterCompletion(transaction);

        logListener.assertWarn("beforeCommit(..) failed for transaction");
        logListener.assertWarn("afterCommit(..) failed for transaction");
        logListener.assertWarn("afterCompletion(..) failed for transaction");
    }
}