package io.quarkus.narayana.jta.runtime.context;

import java.lang.annotation.Annotation;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.stream.Collectors;

import javax.enterprise.context.ContextNotActiveException;
import javax.enterprise.context.spi.Contextual;
import javax.enterprise.context.spi.CreationalContext;
import javax.transaction.Status;
import javax.transaction.SystemException;
import javax.transaction.Transaction;
import javax.transaction.TransactionManager;
import javax.transaction.TransactionScoped;
import javax.transaction.TransactionSynchronizationRegistry;

import com.arjuna.ats.internal.jta.transaction.arjunacore.TransactionSynchronizationRegistryImple;

import io.quarkus.arc.ContextInstanceHandle;
import io.quarkus.arc.InjectableBean;
import io.quarkus.arc.InjectableContext;
import io.quarkus.arc.impl.ContextInstanceHandleImpl;

/**
 * {@link javax.enterprise.context.spi.Context} class which defines the {@link TransactionScoped} context.
 */
public class TransactionContext implements InjectableContext {
    // marker object to be put as a key for SynchronizationRegistry to gather all beans created in the scope
    private static final Object TRANSACTION_CONTEXT_MARKER = new Object();

    private final TransactionSynchronizationRegistry transactionSynchronizationRegistry = new TransactionSynchronizationRegistryImple();
    private final TransactionManager transactionManager = com.arjuna.ats.jta.TransactionManager.transactionManager();

    @Override
    public void destroy() {
        if (!isActive()) {
            return;
        }

        TransactionContextState contextState = (TransactionContextState) transactionSynchronizationRegistry
                .getResource(TRANSACTION_CONTEXT_MARKER);
        if (contextState == null) {
            return;
        }
        contextState.destroy();
    }

    @Override
    public void destroy(Contextual<?> contextual) {
        if (!isActive()) {
            return;
        }
        TransactionContextState contextState = (TransactionContextState) transactionSynchronizationRegistry
                .getResource(TRANSACTION_CONTEXT_MARKER);
        if (contextState == null) {
            return;
        }
        contextState.remove(contextual);
    }

    @Override
    public ContextState getState() {
        if (!isActive()) {
            throw new ContextNotActiveException("No active transaction on the current thread");
        }

        ContextState result;
        TransactionContextState contextState = (TransactionContextState) transactionSynchronizationRegistry
                .getResource(TRANSACTION_CONTEXT_MARKER);
        if (contextState == null) {
            result = new TransactionContextState<>();
        } else {
            result = contextState;
        }
        return result;
    }

    @Override
    public Class<? extends Annotation> getScope() {
        return TransactionScoped.class;
    }

    @Override
    @SuppressWarnings("unchecked")
    public <T> T get(Contextual<T> contextual, CreationalContext<T> creationalContext) {
        if (!isActive()) {
            throw new ContextNotActiveException();
        }
        if (contextual == null) {
            throw new IllegalArgumentException("Contextual parameter must not be null");
        }

        TransactionContextState<T> contextState;
        contextState = (TransactionContextState<T>) transactionSynchronizationRegistry
                .getResource(TRANSACTION_CONTEXT_MARKER);

        if (contextState == null) {
            contextState = new TransactionContextState<>();
            transactionSynchronizationRegistry.putResource(TRANSACTION_CONTEXT_MARKER, contextState);
        }

        ContextInstanceHandle<T> instanceHandle = contextState.get(contextual);
        if (instanceHandle != null) {
            return instanceHandle.get();
        } else if (creationalContext != null) {
            T createdInstance = contextual.create(creationalContext);
            instanceHandle = new ContextInstanceHandleImpl<>((InjectableBean<T>) contextual, createdInstance,
                    creationalContext);

            contextState.put(contextual, instanceHandle);

            return createdInstance;
        } else {
            return null;
        }
    }

    @Override
    public <T> T get(Contextual<T> contextual) {
        return get(contextual, null);
    }

    /**
     * The transaction scoped context is active when a transaction is active.
     */
    @Override
    public boolean isActive() {
        Transaction transaction = getCurrentTransaction();
        if (transaction == null) {
            return false;
        }

        try {
            int currentStatus = transaction.getStatus();
            return currentStatus == Status.STATUS_ACTIVE ||
                    currentStatus == Status.STATUS_MARKED_ROLLBACK ||
                    currentStatus == Status.STATUS_PREPARED ||
                    currentStatus == Status.STATUS_UNKNOWN ||
                    currentStatus == Status.STATUS_PREPARING ||
                    currentStatus == Status.STATUS_COMMITTING ||
                    currentStatus == Status.STATUS_ROLLING_BACK;
        } catch (SystemException e) {
            throw new RuntimeException("Error getting the status of the current transaction", e);
        }
    }

    private Transaction getCurrentTransaction() {
        try {
            return transactionManager.getTransaction();
        } catch (SystemException e) {
            throw new RuntimeException("Error getting the current transaction", e);
        }
    }

    /**
     * Representing of the context state. It's a container for all available beans in the context.
     * It's filled during bean usage and cleared on destroy.
     */
    private static class TransactionContextState<T> implements ContextState {

        private final ConcurrentMap<Contextual<T>, ContextInstanceHandle<T>> mapBeanToInstanceHandle = new ConcurrentHashMap<>();

        /**
         * Put the contextual bean and its handle to the container.
         *
         * @param bean bean to be added
         * @param handle handle for the bean which incorporates the bean, contextual instance and the context
         */
        void put(Contextual<T> bean, ContextInstanceHandle<T> handle) {
            mapBeanToInstanceHandle.put(bean, handle);
        }

        /**
         * Remove the bean from the container.
         *
         * @param bean contextual bean instance
         */
        void remove(Contextual<T> bean) {
            mapBeanToInstanceHandle.remove(bean);
        }

        /**
         * Retrieve the bean saved in the container.
         *
         * @param bean retrieving the bean from the container, otherwise {@code null} is returned
         */
        ContextInstanceHandle<T> get(Contextual<T> bean) {
            return mapBeanToInstanceHandle.get(bean);
        }

        /**
         * Destroying all the beans in the container and clearing the container.
         */
        void destroy() {
            for (ContextInstanceHandle<T> handle : mapBeanToInstanceHandle.values()) {
                handle.destroy();
            }
            mapBeanToInstanceHandle.clear();
        }

        /**
         * Method required by the {@link io.quarkus.arc.InjectableContext.ContextState} interface
         * which is then used to get state of the scope in method {@link InjectableContext#getState()}
         *
         * @return list of context bean and the bean instances which are available in the container
         */
        @Override
        public Map<InjectableBean<?>, Object> getContextualInstances() {
            return mapBeanToInstanceHandle.values().stream()
                    .collect(Collectors.toMap(ContextInstanceHandle::getBean, ContextInstanceHandle::get));
        }

    }
}

@@ -2786,34 +2786,32 @@
 tionContextState
-%3C%3E
 ();%0A        %7D el
@@ -3394,19 +3394,16 @@
 extState
-%3CT%3E
  context
@@ -3456,19 +3456,16 @@
 extState
-%3CT%3E
 ) transa
@@ -3642,18 +3642,16 @@
 extState
-%3C%3E
 ();%0A    
@@ -5876,19 +5876,16 @@
 extState
-%3CT%3E
  impleme
@@ -5951,17 +5951,17 @@
 textual%3C
-T
+?
 %3E, Conte
@@ -5977,17 +5977,17 @@
 eHandle%3C
-T
+?
 %3E%3E mapBe
@@ -6291,24 +6291,28 @@
   */%0A       
+ %3CT%3E
  void put(Co
@@ -6563,24 +6563,28 @@
   */%0A       
+ %3CT%3E
  void remove
@@ -6852,32 +6852,36 @@
       */%0A       
+ %3CT%3E
  ContextInstance
@@ -6926,32 +6926,59 @@
           return
+ (ContextInstanceHandle%3CT%3E)
  mapBeanToInstan
@@ -7169,33 +7169,33 @@
 tInstanceHandle%3C
-T
+?
 %3E handle : mapBe

Remove unnecessary type from TransactionalContextState.