package com.atlassian.bitbucket.mesh.test;

import com.atlassian.bitbucket.mesh.execution.ExecutionContext;
import com.atlassian.bitbucket.mesh.execution.ExecutionLocal;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import io.grpc.Context;
import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;

import javax.annotation.Nonnull;
import java.util.HashMap;
import java.util.Map;

import static java.util.Objects.requireNonNull;

public class ExecutionContextRule implements TestRule {

    private final ExecutionManager executionManager;
    private final Map<ExecutionContext, Context.CancellableContext> openExecutionContexts;

    public ExecutionContextRule(@Nonnull ExecutionManager executionManager) {
        this.executionManager = requireNonNull(executionManager, "executionManager");

        openExecutionContexts = new HashMap<>();
    }

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

            @Override
            public void evaluate() throws Throwable {
                Throwable throwable = executionManager.withContext(context -> {
                    try {
                        base.evaluate();
                        return null;
                    } catch (Throwable t) {
                        //Return the exception to be rethrown outside of the ExecutionCallback
                        return t;
                    } finally {
                        openExecutionContexts.values().forEach(ctx -> ctx.cancel(null));
                        openExecutionContexts.clear();
                    }
                });
                //Rethrow the inner exception if one was caught
                if (throwable != null) {
                    throw throwable;
                }
            }
        };
    }

    /**
     * Closes an {@link ExecutionContext} that was created with {@link #newContext()}, triggering any cleanup listeners
     * that have been registered for the context.
     *
     * @param executionContext the context to close
     * @return {@code true} if the context was closed, otherwise false
     */
    public boolean close(@Nonnull ExecutionContext executionContext) {
        Context.CancellableContext ctx = openExecutionContexts.remove(
                requireNonNull(executionContext, "executionContext"));
        if (ctx == null) {
            return false;
        }
        ctx.cancel(null);
        return true;
    }

    @Nonnull
    public ExecutionManager getExecutionManager() {
        return executionManager;
    }

    /**
     * Creates a new {@link ExecutionContext} that is separate from the currently active context, which is useful for
     * testing classes that use {@link ExecutionLocal} variables.
     * <p>
     * The new context can be closed by calling {@link #close(ExecutionContext)}, at which time any registered cleanup
     * callbacks will be invoked.
     *
     * @return a new {@link ExecutionContext}
     */
    @Nonnull
    public ExecutionContext newContext() {
        Context.CancellableContext grpcContext = Context.ROOT.withCancellation();
        Context prev = grpcContext.attach();
        try {
            return executionManager.withContext(ctx -> {
                openExecutionContexts.put(ctx, grpcContext);
                return ctx;
            });
        } finally {
            grpcContext.detach(prev);
        }
    }

    /**
     * Runs the provided {@link Runnable} in a new {@link ExecutionContext}, which will be torn down when the
     * runnable returns. This can be useful for tests that need to verify behavior triggered on close.
     *
     * @param runnable the operation to perform in a new context
     */
    public void runInNewContext(@Nonnull Runnable runnable) {
        Context.ROOT.run(() -> executionManager.withContext(context -> {
            requireNonNull(runnable, "runnable").run();
            return null;
        }));
    }

    /**
     * Runs the provided {@link Runnable} <i>without</i> an {@link ExecutionContext}. By default, when this rule is
     * used, all test methods are run in an {@link ExecutionContext}. This can be used to escape that context, for
     * tests that are verifying behavior when no context is available.
     *
     * @param runnable the operation to perform without a context
     */
    public void runWithoutContext(@Nonnull Runnable runnable) {
        Context.ROOT.run(runnable);
    }
}
