package org.realityforge.replicant.server.ee;

import javax.annotation.Nonnull;
import javax.annotation.Resource;
import javax.naming.InitialContext;
import javax.naming.NamingException;
import javax.persistence.PostPersist;
import javax.persistence.PostUpdate;
import javax.persistence.PreRemove;
import javax.transaction.TransactionSynchronizationRegistry;
import org.realityforge.replicant.server.EntityMessage;
import org.realityforge.replicant.server.EntityMessageGenerator;
import org.realityforge.replicant.server.EntityMessageSet;

/**
 * Abstract class to extend to collect changes to entities so that they can be replicated by the library.
 * Library users should implement the {@link #getEntityMessageGenerator()}} method to provide message
 * generator specific to the particular application.
 */
public abstract class ChangeRecorder
{
  /**
   * Standard JNDI key for TransactionSynchronizationRegistry.
   */
  private static final String REGISTRY_KEY = "java:comp/TransactionSynchronizationRegistry";

  /**
   * The registry is actually accessed via JNDI. The @Resource annotation is ignored in production as JPA 2.0 does not
   * support it. However our Guice based test infrastructure uses it to populate and avoid instantiation of JNDI
   * resources.
   */
  @Resource
  private TransactionSynchronizationRegistry _registry;

  @PostUpdate
  @PostPersist
  public void postUpdate( final Object object )
  {
    recordEntityMessagesForEntity( object, true );
  }

  /**
   * Collect messages before they are committed to the database with the
   * assumption that the remove will not fail. This allows us to traverse
   * the object graph before it is deleted. Note: This is a different strategy
   * from postUpdate() but PostUpdate may be changed in the future to match
   * remove hook.
   *
   * @param object the entity removed.
   */
  @PreRemove
  public void preRemove( final Object object )
  {
    recordEntityMessagesForEntity( object, false );
  }

  /**
   * Record the EntityMessage for the specified entity and dependent entities in the transactions EntityMessageSet.
   *
   * @param entity the entity to record.
   * @param isUpdate true if change is an update, false if it is a delete.
   */
  protected final void recordEntityMessagesForEntity( @Nonnull final Object entity, final boolean isUpdate )
  {
    if ( !getRegistry().getRollbackOnly() )
    {
      recordEntityMessageForEntity( entity, isUpdate );
    }
  }

  /**
   * Record the EntityMessage for specified entity in the transactions EntityMessageSet.
   *
   * @param entity the entity to record.
   * @param isUpdate true if change is an update, false if it is a delete.
   */
  protected final void recordEntityMessageForEntity( final Object entity, final boolean isUpdate )
  {
    final EntityMessage entityMessage = getEntityMessageGenerator().convertToEntityMessage( entity, isUpdate );
    if ( null != entityMessage )
    {
      getEntityMessageSet().merge( entityMessage );
    }
  }

  /**
   * Return the EntityMessageSet used to collect messages for the current transaction.
   *
   * @return the EntityMessageSet used to collect messages for the current transaction.
   */
  protected final EntityMessageSet getEntityMessageSet()
  {
    return EntityMessageCacheUtil.getEntityMessageSet( getRegistry() );
  }

  /**
   * @return the message generator for the specific application.
   */
  protected abstract EntityMessageGenerator getEntityMessageGenerator();

  private TransactionSynchronizationRegistry getRegistry()
  {
    if ( null == _registry )
    {
      _registry = lookupTransactionSynchronizationRegistry();
    }
    return _registry;
  }

  private static TransactionSynchronizationRegistry lookupTransactionSynchronizationRegistry()
  {
    try
    {
      return (TransactionSynchronizationRegistry) new InitialContext().lookup( REGISTRY_KEY );
    }
    catch ( final NamingException ne )
    {
      final String message =
        "Unable to locate TransactionSynchronizationRegistry at " + REGISTRY_KEY + " due to " + ne;
      throw new IllegalStateException( message, ne );
    }
  }
}


@@ -1339,16 +1339,469 @@
 istry;%0A%0A
+  private static TransactionSynchronizationRegistry lookupTransactionSynchronizationRegistry()%0A  %7B%0A    try%0A    %7B%0A      return (TransactionSynchronizationRegistry) new InitialContext().lookup( REGISTRY_KEY );%0A    %7D%0A    catch ( final NamingException ne )%0A    %7B%0A      final String message =%0A        %22Unable to locate TransactionSynchronizationRegistry at %22 + REGISTRY_KEY + %22 due to %22 + ne;%0A      throw new IllegalStateException( message, ne );%0A    %7D%0A  %7D%0A%0A
   @PostU
@@ -1865,32 +1865,65 @@
 ct object )%0A  %7B%0A
+    if ( !inRollback() )%0A    %7B%0A  
     recordEntity
@@ -1953,24 +1953,30 @@
 ct, true );%0A
+    %7D%0A
   %7D%0A%0A  /**%0A 
@@ -2404,24 +2404,57 @@
 bject )%0A  %7B%0A
+    if ( !inRollback() )%0A    %7B%0A  
     recordEn
@@ -2489,24 +2489,183 @@
 t, false );%0A
+    %7D%0A  %7D%0A%0A  /**%0A   * @return true if transaction is in rollback.%0A   */%0A  protected final boolean inRollback()%0A  %7B%0A    return getRegistry().getRollbackOnly();%0A
   %7D%0A%0A  /**%0A 
@@ -2795,32 +2795,34 @@
 * @param entity 
+  
 the entity to re
@@ -3025,60 +3025,8 @@
   %7B%0A
-    if ( !getRegistry().getRollbackOnly() )%0A    %7B%0A  
     
@@ -3071,30 +3071,24 @@
 isUpdate );%0A
-    %7D%0A
   %7D%0A%0A  /**%0A 
@@ -3199,16 +3199,18 @@
  entity 
+  
 the enti
@@ -4297,461 +4297,8 @@
 %0A  %7D
-%0A%0A  private static TransactionSynchronizationRegistry lookupTransactionSynchronizationRegistry()%0A  %7B%0A    try%0A    %7B%0A      return (TransactionSynchronizationRegistry) new InitialContext().lookup( REGISTRY_KEY );%0A    %7D%0A    catch ( final NamingException ne )%0A    %7B%0A      final String message =%0A        %22Unable to locate TransactionSynchronizationRegistry at %22 + REGISTRY_KEY + %22 due to %22 + ne;%0A      throw new IllegalStateException( message, ne );%0A    %7D%0A  %7D
 %0A%7D%0A%0A

Move rollback check out of recordEntityMessagesForEntity and into callers