import java.lang.reflect.Field;
import java.util.Map;

import org.hibernate.cache.CacheException;
import org.hibernate.cache.infinispan.InfinispanRegionFactory;
import org.hibernate.cache.infinispan.timestamp.ClusteredTimestampsRegionImpl;
import org.hibernate.engine.spi.SharedSessionContractImplementor;
import org.infinispan.AdvancedCache;
import org.infinispan.notifications.Listener;

/**
 * See https://hibernate.atlassian.net/browse/HHH-10677
 * In a cluster environment, the timestamp cache is updated asynchronously, even for operations that happens in the same
 * node. This is a "hack" to make sure the timestamp local cache is always up-to-date in the node that originated the call.
 * 
 * This will ensure that a cached query will return the correct results just after the flush() was called.
 */
@Listener
public class CustomClusteredTimestampsRegionImpl extends ClusteredTimestampsRegionImpl {

    private final Map localCache = retrieveLocalCacheFieldValue(this);

    public CustomClusteredTimestampsRegionImpl(AdvancedCache cache, String name, InfinispanRegionFactory factory) {
        super(cache, name, factory);
    }

    private static Map retrieveLocalCacheFieldValue(ClusteredTimestampsRegionImpl target) {
        try {
            Field field = ClusteredTimestampsRegionImpl.class.getDeclaredField("localCache");
            field.setAccessible(true);
            return (Map) field.get(target);
        } catch (NoSuchFieldException | SecurityException | IllegalAccessException e) {
            throw new ExceptionInInitializerError(e);
        }
    }

    @Override
    public void put(SharedSessionContractImplementor session, final Object key, final Object value) throws CacheException {
        super.put(session, key, value);
        localCache.put(key, value);
    }

}