package com.atlassian.bitbucket.mesh.execution;

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

import java.util.Optional;
import java.util.concurrent.Callable;
import java.util.function.Consumer;
import java.util.function.Function;

import static org.junit.Assert.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class GrpcExecutionManagerTest {

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

    private final GrpcExecutionManager executionManager = new GrpcExecutionManager();

    @Mock
    private Runnable cleanupCallback;

    @Test
    public void testCleanupCallbackExceptionsAreHandled() {
        doThrow(new RuntimeException("mock exception")).when(cleanupCallback).run();
        executionManager.withContext(ctx -> {
            ctx.addCleanupCallback(cleanupCallback);
            return null;
        });

        verify(cleanupCallback).run();
        //We should have gotten a WARN log message detailing the failure
        logListener.assertWarn("Error while executing cleanup callback");
    }

    @Test
    public void testCleanupCallbackRunsImmediatelyIfContextAlreadyClosed() {
        //Mock an exception to validate that exceptions are handled
        doThrow(new RuntimeException("mock exception")).when(cleanupCallback).run();
        ExecutionContext[] context = new ExecutionContext[1];
        executionManager.withContext(ctx -> context[0] = ctx);

        //At this point the ExecutionContext should already have been closed
        context[0].addCleanupCallback(cleanupCallback);
        verify(cleanupCallback).run();
        //We should have gotten a WARN log message detailing the failure
        logListener.assertWarn("Error while executing cleanup callback");
    }

    @Test
    public void testContextPropagation() {
        inGrpcContext(() -> {
            //Create and return the execution context. The context remains open while the outer
            //CancellableContext is open
            ExecutionContext context = executionManager.withContext(ctx -> ctx);
            assertNotNull(context);
            assertNotNull(context.getId());

            //Since we're not in withContext closure, there should be no current context
            assertEquals(Optional.empty(), executionManager.getContext());

            //Validate that ExecutionContext.run sets up the current context
            context.run(() -> assertEquals(context, executionManager.getContext().orElse(null)));
            //Validate that ExecutionContext.wrap sets up the current context
            context.wrap(() -> assertEquals(context, executionManager.getContext().orElse(null))).run();

            //Validate that ExecutionContext.call sets up the current context
            assertEquals("test-call", context.call(() -> {
                assertEquals(context, executionManager.getContext().orElse(null));
                return "test-call";
            }));
            assertEquals("test-wrap-call", context.wrap(() -> {
                assertEquals(context, executionManager.getContext().orElse(null));
                return "test-wrap-call";
            }).call());
        });
    }

    @Test
    public void testExecutionLocal() {
        ExecutionLocal<String> executionLocal = executionManager.newExecutionLocal();

        //Verify ExecutionLocal behavior when no ExecutionContext is active
        assertNotPresent(executionLocal);
        assertFalse(executionLocal.isActive());
        assertNull(executionLocal.computeIfAbsent(() -> "test"));
        assertNull(executionLocal.get());
        assertFalse(executionLocal.set("test"));
        assertNull(executionLocal.remove());

        //Create two separate contexts for testing, and set a value for each
        ExecutionContext context1 = executionManager.withContext(ctx -> {
            assertTrue(executionLocal.isActive());
            executionLocal.checkValid();
            executionLocal.set("context1");
            assertPresentAndEqualTo("context1", executionLocal);
            return ctx;
        });
        ExecutionContext context2 = executionManager.withContext(ctx -> {
            assertTrue(executionLocal.isActive());
            assertNotPresent(executionLocal);
            assertEquals("context2", executionLocal.computeIfAbsent(() -> "context2"));
            assertPresentAndEqualTo("context2", executionLocal);
            return ctx;
        });
        assertNotNull(context1);
        assertNotNull(context2);

        //Outside of the contexts, the executionLocal should still be empty
        assertFalse(executionLocal.isPresent());
        context1.run(() -> {
            assertPresentAndEqualTo("context1", executionLocal);
            executionLocal.set("context1-1");
            assertPresentAndEqualTo("context1-1", executionLocal);
            assertEquals("context1-1", executionLocal.remove());
            assertNotPresent(executionLocal);
        });

        context2.run(() -> assertPresentAndEqualTo("context2", executionLocal));
    }

    @SuppressWarnings("unchecked")
    @Test
    public void testExecutionLocalCleanupListener() {
        Consumer<String> cleanupListener = mock(Consumer.class);
        //Mock throwing an exception to verify that exceptions are handled
        doThrow(new IllegalStateException("mocking cleanup failure")).when(cleanupListener).accept(anyString());
        ExecutionLocal<String> executionLocal = executionManager.newExecutionLocal(cleanupListener);

        String value = executionManager.withContext(ctx -> {
            executionLocal.set("magic");
            assertPresentAndEqualTo("magic", executionLocal);

            verifyNoInteractions(cleanupListener);
            return executionLocal.get();
        });
        verify(cleanupListener).accept(eq(value));
    }

    @Test
    public void testExecutionLocalCompute() {
        ExecutionLocal<String> executionLocal = executionManager.newExecutionLocal();

        executionManager.withContext(ctx -> {
            //Compute when empty
            executionLocal.compute(curVal -> {
                assertNull(curVal);
                return "computed";
            });
            assertPresentAndEqualTo("computed", executionLocal);
            executionLocal.compute(curVal -> {
                assertEquals("computed", curVal);
                return "altered";
            });
            assertPresentAndEqualTo("altered", executionLocal);
            return null;
        });
        assertNotPresent(executionLocal);
    }

    @Test
    public void testExecutionLocalRemoveIf() {
        ExecutionLocal<String> executionLocal = executionManager.newExecutionLocal();

        executionManager.withContext(ctx -> {
            assertNull(executionLocal.removeIf(value -> {
                fail("Predicate should not be called if no value is present");
                return true;
            }));
            executionLocal.set("value");
            assertNull(executionLocal.removeIf(value -> {
                assertEquals("value", value);
                return false;
            }));
            assertPresentAndEqualTo("value", executionLocal);
            assertEquals("value", executionLocal.removeIf(value -> true));
            assertNotPresent(executionLocal);
            return null;
        });
        assertNotPresent(executionLocal);
    }

    @Test
    public void testGetContextReturnsEmptyWhenNoContextIsSet() {
        assertEquals(Optional.empty(), new GrpcExecutionManager().getContext());
    }

    @Test(expected = IllegalStateException.class)
    public void testGetContextOrThrowWhenNoContextIsSet() {
        new GrpcExecutionManager().getContextOrThrow();
    }

    @Test
    public void testMaybeWrapCallable() throws Exception {
        Callable<Boolean> callable = () -> executionManager.getContext().isPresent();

        //When no context is present, maybeWrap should just return the callable unchanged
        Callable<Boolean> wrapped = executionManager.maybeWrap(callable);
        assertSame(callable, wrapped);
        assertFalse(wrapped.call());

        //When a context _is_ present, the callable should be decorated
        wrapped = executionManager.withContext(ctx -> executionManager.maybeWrap(callable));
        assertNotNull(wrapped);
        assertNotSame(callable, wrapped);
        assertTrue(wrapped.call());
    }

    @Test
    public void testMaybeWrapRunnable() {
        Runnable assertNoContextPresent = () -> assertFalse(executionManager.getContext().isPresent());

        //When no context is present, maybeWrap should just return the runnable unchanged
        Runnable wrapped = executionManager.maybeWrap(assertNoContextPresent);
        assertSame(assertNoContextPresent, wrapped);
        assertNoContextPresent.run();

        //When a context _is_ present, the callable should be decorated
        Runnable assertContextPresent = () -> assertTrue(executionManager.getContext().isPresent());
        wrapped = executionManager.withContext(ctx -> {
            Runnable w = executionManager.maybeWrap(assertContextPresent);
            w.run();
            return w;
        });
        assertNotNull(wrapped);
        assertNotSame(assertContextPresent, wrapped);
        wrapped.run();
    }

    @Test
    public void testWithContextAutoClosesContextIfNoContextActive() {
        String returnValue = executionManager.withContext(ctx -> {
            ctx.addCleanupCallback(cleanupCallback);

            assertEquals(ctx, executionManager.getContextOrThrow());
            assertEquals(ctx, executionManager.getContext().orElse(null));
            return "test";
        });

        assertEquals("test", returnValue);
        verify(cleanupCallback).run();
        assertEquals(Optional.empty(), executionManager.getContext());
    }

    @Test
    public void testWithContextAutoClosesContextIfNonCancellableContextActive() {
        //Simulate a non-root gRPC context that is not cancellable
        Context.ROOT.withValue(Context.key("testing"), "testing").run(() -> {
            String returnValue = executionManager.withContext(ctx -> {
                ctx.addCleanupCallback(cleanupCallback);

                assertEquals(ctx, executionManager.getContextOrThrow());
                assertEquals(ctx, executionManager.getContext().orElse(null));
                return "test";
            });

            assertEquals("test", returnValue);
            verify(cleanupCallback).run();
            assertEquals(Optional.empty(), executionManager.getContext());
        });
    }

    @Test
    public void testWithContextDoesNotAutoCloseWhenGrpcContextActive() {
        inGrpcContext(() -> {
            //This simulates the gRPC context set up by gRPC when we're called as part of a gRPC request
            String returnValue = executionManager.withContext(ctx -> {
                ctx.addCleanupCallback(cleanupCallback);

                assertEquals(ctx, executionManager.getContextOrThrow());
                assertEquals(ctx, executionManager.getContext().orElse(null));
                return "test";
            });

            assertEquals("test", returnValue);

            //The cleanupCallback shouldn't be called until the outer grpcContext is closed
            assertEquals(Optional.empty(), executionManager.getContext());
            verifyNoInteractions(cleanupCallback);
        });

        verify(cleanupCallback).run();
    }

    @Test
    public void testWithContextDoesNotAutoCloseWhenNestedGrpcContextActive() {
        inGrpcContext(() -> {
            //This simulates the gRPC context set up by gRPC when we're called as part of a gRPC request
            Context.current().withValue(Context.key("testing"), "value").run(() -> {
                //This simulates a nested gRPC context that does not directly implement CancellableContext
                String returnValue = executionManager.withContext(ctx -> {
                    ctx.addCleanupCallback(cleanupCallback);

                    assertEquals(ctx, executionManager.getContextOrThrow());
                    assertEquals(ctx, executionManager.getContext().orElse(null));
                    return "test";
                });

                assertEquals("test", returnValue);
            });

            //The cleanupCallback shouldn't be called until the outer grpcContext is closed
            assertEquals(Optional.empty(), executionManager.getContext());
            verifyNoInteractions(cleanupCallback);
        });

        verify(cleanupCallback).run();
    }

    @Test
    public void testWrappingProxy() {
        inGrpcContext(() -> {
            ExecutionLocal<String> contextName = executionManager.newExecutionLocal();
            Function<String, String> contextNamePrepender = value -> contextName.get() + ":" + value;

            Function<String, String> proxiedPrepender = executionManager.withContext(ctx -> {
                contextName.set("magic");
                 return ctx.wrap(contextNamePrepender, Function.class);
            });

            assertEquals("null:isn't it?", contextNamePrepender.apply("isn't it?"));
            assertEquals("magic:isn't it?", proxiedPrepender.apply("isn't it?"));
        });
    }

    @SuppressWarnings("unchecked")
    private static void assertPresentAndEqualTo(String value, ExecutionLocal<String> executionLocal) {
        Consumer<String> mock = mock(Consumer.class);

        assertTrue(executionLocal.isActive());
        assertTrue(executionLocal.isPresent());
        assertEquals(value, executionLocal.get());
        assertTrue(executionLocal.ifPresent(mock));

        verify(mock).accept(eq(value));
    }

    @SuppressWarnings("unchecked")
    private static void assertNotPresent(ExecutionLocal<String> executionLocal) {
        Consumer<String> mock = mock(Consumer.class);

        assertFalse(executionLocal.isPresent());
        assertFalse(executionLocal.ifPresent(mock));

        verifyNoInteractions(mock);
    }

    private interface ThrowingRunnable {

        void run() throws Exception;
    }

    private static void inGrpcContext(ThrowingRunnable testAction) {
        try (Context.CancellableContext grpcContext = Context.current().withCancellation()) {
            //This simulates the gRPC context set up by gRPC when we're called as part of a gRPC request
            grpcContext.run(() -> {
                try {
                    testAction.run();
                } catch (Exception e) {
                    Throwables.throwIfUnchecked(e);
                    throw new RuntimeException(e);
                }
            });
        }
    }
}