/*
 * JBoss, the OpenSource EJB server
 *
 * Distributable under LGPL license.
 * See terms of license at gnu.org.
 */
package org.jboss.ejb3.entity;

import java.util.Map;
import javax.persistence.EntityManager;
import javax.persistence.EntityManagerFactory;
import javax.persistence.FlushModeType;
import javax.persistence.PersistenceContextType;
import javax.persistence.TransactionRequiredException;
import javax.transaction.RollbackException;
import javax.transaction.Status;
import javax.transaction.Synchronization;
import javax.transaction.SystemException;
import javax.transaction.Transaction;
import org.jboss.ejb3.EJBContainer;
import org.jboss.logging.Logger;
import org.jboss.tm.TransactionLocal;
import org.jboss.tm.TxUtils;
import org.jboss.tm.TxManager;

/**
 * @author <a href="mailto:gavine@hibernate.org">Gavin King</a>
 * @version $Revision: 1.3.2.4 $
 */
public class ManagedEntityManagerFactory
{
   private static final Logger log = Logger.getLogger(ManagedEntityManagerFactory.class);

   public static final String LONG_LIVED_SESSION = "LONG_LIVED_SESSION";

   protected EntityManagerFactory entityManagerFactory;
   protected TransactionLocal session = new TransactionLocal(TxManager.getInstance());
   protected String jndiName;

   public ManagedEntityManagerFactory(EntityManagerFactory sf)
   {
      this(sf, EJBContainer.ENC_CTX_NAME + "/" + EJBContainer.MANAGED_ENTITY_MANAGER_FACTORY);
   }

   public ManagedEntityManagerFactory(EntityManagerFactory sf, String jndiName)
   {
      this.entityManagerFactory = sf;
      this.jndiName = jndiName;
   }

   public EntityManagerFactory getEntityManagerFactory()
   {
      return entityManagerFactory;
   }

   public String getJndiName()
   {
      return jndiName;
   }

   public void destroy()
   {
      entityManagerFactory.close();
   }

   private static class SessionSynchronization implements Synchronization
   {
      private EntityManager manager;
      private Transaction tx;
      private boolean closeAtTxCompletion;

      public SessionSynchronization(EntityManager session, Transaction tx, boolean close)
      {
         this.manager = session;
         this.tx = tx;
         closeAtTxCompletion = close;
      }

      public void beforeCompletion()
      {
         try
         {
            int status = tx.getStatus();
            if (status != Status.STATUS_ROLLEDBACK && status != Status.STATUS_ROLLING_BACK && status != Status.STATUS_MARKED_ROLLBACK)
            {
               if (FlushModeInterceptor.getTxFlushMode() != FlushModeType.NEVER)
               {
                  log.debug("************** flushing.....");
                  manager.flush();
               }
            }
         }
         catch (SystemException e)
         {
            throw new RuntimeException(e);
         }
      }

      public void afterCompletion(int status)
      {
         if (closeAtTxCompletion)
         {
            log.debug("************** closing entity managersession **************");
            manager.close();
         }
      }
   }

   public static ThreadLocal longLivedSession = new ThreadLocal();

   public TransactionLocal getTransactionSession()
   {
      return session;
   }

   public void registerPersistenceContext()
   {
       getSession();
   }

   public EntityManager getSession()
   {
      Transaction tx = session.getTransaction();
      if (tx == null) throw new TransactionRequiredException("EntityManager must be access within a transaction");
      if (!TxUtils.isActive(tx)) throw new TransactionRequiredException("Transaction must be active to access EntityManager");
      EntityManager rtnSession = (EntityManager) session.get();
      if (rtnSession == null)
      {
         log.debug("GETTING NEW EntityManager");

         Map propagated = (Map) longLivedSession.get();
         if (propagated != null)
         {
            log.debug("logLivedSession.get() was NOT NULL: " + jndiName);
            ExtendedPersistenceContext pc = (ExtendedPersistenceContext)propagated.get(jndiName);
            rtnSession = pc.getPersistenceContext();
         }
         else
         {
            log.debug("logLivedSession.get() was NULL");
         }

         boolean closeAtTxCompletion = false;
         if (rtnSession == null)
         {
            closeAtTxCompletion = true;
            rtnSession = createEntityManager(PersistenceContextType.TRANSACTION);
            log.debug("LONG LIVED NOT FOUND");
         }
         else
         {
            log.debug("using LONG LIVED!!!");
         }
         try
         {
            tx.registerSynchronization(new SessionSynchronization(rtnSession, tx, closeAtTxCompletion));
         }
         catch (RollbackException e)
         {
            throw new RuntimeException(e);  //To change body of catch statement use Options | File Templates.
         }
         catch (SystemException e)
         {
            throw new RuntimeException(e);  //To change body of catch statement use Options | File Templates.
         }
         session.set(rtnSession);
      }
      return rtnSession;
   }

   public EntityManager createEntityManager(PersistenceContextType type)
   {
      return entityManagerFactory.createEntityManager(type);
   }


}
