package io.quarkus.arc.impl;

import io.quarkus.arc.Arc;
import io.quarkus.arc.ManagedContext;
import io.smallrye.mutiny.Multi;
import io.smallrye.mutiny.Uni;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import javax.annotation.Priority;
import javax.enterprise.context.control.ActivateRequestContext;
import javax.interceptor.AroundInvoke;
import javax.interceptor.Interceptor;
import javax.interceptor.InvocationContext;

@Interceptor
@ActivateRequestContext
@Priority(Interceptor.Priority.PLATFORM_BEFORE + 100)
public class ActivateRequestContextInterceptor {

    @AroundInvoke
    Object aroundInvoke(InvocationContext ctx) throws Exception {
        if (ctx.getMethod().getReturnType().equals(Uni.class)) {
            return invokeUni(ctx);
        }

        if (ctx.getMethod().getReturnType().equals(CompletionStage.class)) {
            return invokeStage(ctx);
        }

        if (ctx.getMethod().getReturnType().equals(Multi.class)) {
            return invokeMulti(ctx);
        }

        return invoke(ctx);
    }

    private CompletionStage<?> invokeStage(InvocationContext ctx) {
        ManagedContext requestContext = Arc.container().requestContext();
        if (requestContext.isActive()) {
            return proceedWithStage(ctx);
        }

        return activate(requestContext)
                .thenCompose(v -> proceedWithStage(ctx))
                .whenComplete((r, t) -> requestContext.terminate());
    }

    private static CompletionStage<ManagedContext> activate(ManagedContext requestContext) {
        try {
            requestContext.activate();
            return CompletableFuture.completedStage(requestContext);
        } catch (Throwable t) {
            return CompletableFuture.failedStage(t);
        }
    }

    private CompletionStage<?> proceedWithStage(InvocationContext ctx) {
        try {
            return (CompletionStage<?>) ctx.proceed();
        } catch (Throwable t) {
            return CompletableFuture.failedStage(t);
        }
    }

    private Multi<?> invokeMulti(InvocationContext ctx) {
        return Multi.createFrom().deferred(() -> {
            ManagedContext requestContext = Arc.container().requestContext();
            if (requestContext.isActive()) {
                return proceedWithMulti(ctx);
            }

            return Multi.createFrom().deferred(() -> {
                requestContext.activate();
                return proceedWithMulti(ctx);
            }).onTermination().invoke(requestContext::terminate);
        });
    }

    private Multi<?> proceedWithMulti(InvocationContext ctx) {
        try {
            return (Multi<?>) ctx.proceed();
        } catch (Throwable t) {
            return Multi.createFrom().failure(t);
        }
    }

    private Uni<?> invokeUni(InvocationContext ctx) {
        return Uni.createFrom().deferred(() -> {
            ManagedContext requestContext = Arc.container().requestContext();
            if (requestContext.isActive()) {
                return proceedWithUni(ctx);
            }

            return Uni.createFrom().deferred(() -> {
                requestContext.activate();
                return proceedWithUni(ctx);
            }).eventually(requestContext::terminate);
        });
    }

    private Uni<?> proceedWithUni(InvocationContext ctx) {
        try {
            return (Uni<?>) ctx.proceed();
        } catch (Throwable t) {
            return Uni.createFrom().failure(t);
        }
    }

    private Object invoke(InvocationContext ctx) throws Exception {
        ManagedContext requestContext = Arc.container().requestContext();
        if (requestContext.isActive()) {
            return ctx.proceed();
        }

        try {
            requestContext.activate();
            return ctx.proceed();
        } finally {
            requestContext.terminate();
        }
    }
}

@@ -701,187 +701,141 @@
     
-if (ctx.getMethod().getReturnType().equals(Uni.class)) %7B%0A            return invokeUni(ctx);%0A        %7D%0A%0A        if (ctx.getMethod().getReturnType().equals(CompletionStage.class)) %7B
+switch (ReactiveType.valueOf(ctx.getMethod())) %7B%0A            case UNI:%0A                return invokeUni(ctx);%0A            case MULTI:
 %0A   
@@ -835,32 +835,36 @@
 TI:%0A            
+    
 return invokeSta
@@ -860,21 +860,21 @@
 n invoke
-Stage
+Multi
 (ctx);%0A 
@@ -881,81 +881,31 @@
     
-   %7D%0A%0A
        
- if (ctx.getMethod().getReturnType().equals(Multi.class)) %7B%0A
+case STAGE:%0A    
     
@@ -917,37 +917,37 @@
    return invoke
-Multi
+Stage
 (ctx);%0A        %7D
@@ -937,35 +937,53 @@
 e(ctx);%0A        
-%7D%0A%0A
+    default:%0A        
         return i
@@ -990,24 +990,34 @@
 nvoke(ctx);%0A
+        %7D%0A
     %7D%0A%0A    p

Use the new utility class