package io.quarkus.arc.runtime.context;

import java.util.Map;

import org.eclipse.microprofile.context.ThreadContext;
import org.eclipse.microprofile.context.spi.ThreadContextController;
import org.eclipse.microprofile.context.spi.ThreadContextProvider;
import org.eclipse.microprofile.context.spi.ThreadContextSnapshot;

import io.quarkus.arc.Arc;
import io.quarkus.arc.ArcContainer;
import io.quarkus.arc.InjectableContext;
import io.quarkus.arc.ManagedContext;

/**
 * Context propagation for Arc
 * Only handles Request context as that's currently the only one in Arc that needs propagation.
 */
public class ArcContextProvider implements ThreadContextProvider {

    private static ThreadContextSnapshot NOOP_SNAPSHOT = () -> () -> {
    };

    @Override
    public ThreadContextSnapshot currentContext(Map<String, String> map) {
        ArcContainer arc = Arc.container();
        if (arc == null || !arc.isRunning()) {
            //return null as per docs to state that propagation of this context is not supported
            return null;
        }

        if (!isContextActiveOnThisThread(arc)) {
            // request context not active, nothing to propagate, return no-op
            return NOOP_SNAPSHOT;
        }

        // capture the state
        InjectableContext.ContextState state = arc.requestContext().getState();
        return () -> {
            // can be called later on, we should retrieve the container again
            ArcContainer arcContainer = Arc.container();
            if (arcContainer == null || !arcContainer.isRunning()) {
                throw new IllegalStateException("Arc context propagation was attempted but the container is not running.");
            }
            ThreadContextController controller;
            ManagedContext requestContext = arcContainer.requestContext();
            // this is executed on another thread, context can but doesn't need to be active here
            if (isContextActiveOnThisThread(arcContainer)) {
                // context active, store current state, feed it new one and restore state afterwards
                InjectableContext.ContextState stateToRestore = requestContext.getState();
                requestContext.deactivate();
                requestContext.activate(state);
                controller = () -> {
                    // clean up, reactivate context with previous values
                    requestContext.deactivate();
                    requestContext.activate(stateToRestore);
                };
            } else {
                // context not active, activate and pass it new instance, deactivate afterwards
                requestContext.activate(state);
                controller = () -> {
                    requestContext.deactivate();
                };
            }
            return controller;
        };
    }

    @Override
    public ThreadContextSnapshot clearedContext(Map<String, String> map) {
        // note that by cleared we mean that we still activate context if need be, just leave the contents blank
        ArcContainer arc = Arc.container();
        if (arc == null || !arc.isRunning()) {
            //return null as per docs to state that propagation of this context is not supported
            return null;
        }

        if (!isContextActiveOnThisThread(arc)) {
            // request context not active, nothing to propagate, return no-op
            return NOOP_SNAPSHOT;
        }

        return () -> {
            // can be called later on, we should retrieve the container again
            ArcContainer arcContainer = Arc.container();
            if (arcContainer == null || !arcContainer.isRunning()) {
                throw new IllegalStateException("Arc context propagation was attempted but the container is not running.");
            }
            ThreadContextController controller;
            ManagedContext requestContext = arcContainer.requestContext();
            // this is executed on another thread, context can but doesn't need to be active here
            if (isContextActiveOnThisThread(arcContainer)) {
                // context active, store current state, start blank context anew and restore state afterwards
                InjectableContext.ContextState stateToRestore = requestContext.getState();
                requestContext.deactivate();
                requestContext.activate();
                controller = () -> {
                    // clean up, reactivate context with previous values
                    requestContext.deactivate();
                    requestContext.activate(stateToRestore);
                };
            } else {
                // context not active, activate blank one, deactivate afterwards
                requestContext.activate();
                controller = () -> {
                    requestContext.deactivate();
                };
            }
            return controller;
        };
    }

    @Override
    public String getThreadContextType() {
        return ThreadContext.CDI;
    }

    private boolean isContextActiveOnThisThread(ArcContainer arc) {
        return arc.requestContext().isActive();
    }
}

@@ -666,87 +666,8 @@
  %7B%0A%0A
-    private static ThreadContextSnapshot NOOP_SNAPSHOT = () -%3E () -%3E %7B%0A    %7D;%0A%0A
     
@@ -987,200 +987,87 @@
     
-if (!isContextActiveOnThisThread(arc)) %7B%0A            // request context
+// capture the state, null indicates
  no
-t
  active
-, nothing to propagate, return no-op%0A            return NOOP_SNAPSHOT;%0A        %7D%0A%0A        // capture the state
+ context while capturing snapshot
 %0A   
@@ -1109,16 +1109,51 @@
  state =
+ isContextActiveOnThisThread(arc) ?
  arc.req
@@ -1168,32 +1168,39 @@
 ext().getState()
+ : null
 ;%0A        return
@@ -2061,32 +2061,143 @@
 t.deactivate();%0A
+                if (state != null) %7B%0A                    // only activate if previous thread had it active%0A    
                 
@@ -2220,32 +2220,50 @@
 ctivate(state);%0A
+                %7D%0A
                 
@@ -2348,32 +2348,153 @@
 previous values%0A
+                    if (state != null) %7B%0A                        // only deactivate if previous thread had it active%0A    
                 
@@ -2518,32 +2518,54 @@
 t.deactivate();%0A
+                    %7D%0A
                 
@@ -2737,32 +2737,143 @@
 vate afterwards%0A
+                if (state != null) %7B%0A                    // only activate if previous thread had it active%0A    
                 
@@ -2895,32 +2895,51 @@
 activate(state);
+%0A%0A                %7D
 %0A               
@@ -2952,32 +2952,153 @@
 oller = () -%3E %7B%0A
+                    if (state != null) %7B%0A                        // only deactivate if previous thread had it active%0A    
                 
@@ -3122,32 +3122,54 @@
 t.deactivate();%0A
+                    %7D%0A
                 
@@ -3664,180 +3664,8 @@
  %7D%0A%0A
-        if (!isContextActiveOnThisThread(arc)) %7B%0A            // request context not active, nothing to propagate, return no-op%0A            return NOOP_SNAPSHOT;%0A        %7D%0A%0A
     

Correct how ArcContextProvider handles susped/resume of context if propagating no context.