/*
 * Copyright 2022. Red Hat, Inc. and/or its affiliates
 * and other contributors as indicated by the @author tags.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.keycloak.models.map.storage.jpa.hibernate.listeners;

import org.hibernate.FlushMode;
import org.hibernate.HibernateException;
import org.hibernate.engine.spi.ActionQueue;
import org.hibernate.engine.spi.PersistenceContext;
import org.hibernate.engine.spi.SessionEventListenerManager;
import org.hibernate.event.internal.DefaultAutoFlushEventListener;
import org.hibernate.event.spi.AutoFlushEvent;
import org.hibernate.event.spi.EventSource;
import org.hibernate.internal.CoreMessageLogger;
import org.hibernate.stat.spi.StatisticsImplementor;
import org.jboss.logging.Logger;
import org.keycloak.models.map.storage.jpa.JpaMapStorage;

/**
 * Extends Hibernate's {@link DefaultAutoFlushEventListener} to always flush queued inserts to allow correct handling
 * of orphans of that entities in the same transactions, and also to clear a session-level query cache.
 * <p />
 * If they wouldn't be flushed, they won't be orphaned (at least not in Hibernate 5.3.24.Final).
 * This class copies over all functionality of the base class that can't be overwritten via inheritance.
 * This is being tracked as part of <a href="https://github.com/keycloak/keycloak/issues/11666">keycloak/keycloak#11666</a>.
 * <p />
 * This also clears the JPA map store query level cache for the {@link JpaMapStorage} whenever there is some data written to the database.
 */
public class JpaAutoFlushListener extends DefaultAutoFlushEventListener {

    public static final JpaAutoFlushListener INSTANCE = new JpaAutoFlushListener();

    private static final CoreMessageLogger LOG = Logger.getMessageLogger(CoreMessageLogger.class, DefaultAutoFlushEventListener.class.getName());

    public void onAutoFlush(AutoFlushEvent event) throws HibernateException {
        final EventSource source = event.getSession();
        final SessionEventListenerManager eventListenerManager = source.getEventListenerManager();
        try {
            eventListenerManager.partialFlushStart();

            if (flushMightBeNeeded(source)) {
                // Need to get the number of collection removals before flushing to executions
                // (because flushing to executions can add collection removal actions to the action queue).
                final ActionQueue actionQueue = source.getActionQueue();
                final int oldSize = actionQueue.numberOfCollectionRemovals();
                flushEverythingToExecutions(event);
                if (flushIsReallyNeeded(event, source)) {
                    LOG.trace("Need to execute flush");
                    event.setFlushRequired(true);

                    // note: performExecutions() clears all collectionXxxxtion
                    // collections (the collection actions) in the session
                    performExecutions(source);
                    postFlush(source);

                    postPostFlush(source);

                    final StatisticsImplementor statistics = source.getFactory().getStatistics();
                    if (statistics.isStatisticsEnabled()) {
                        statistics.flush();
                    }
                } else {
                    LOG.trace("No need to execute flush");
                    event.setFlushRequired(false);
                    actionQueue.clearFromFlushNeededCheck(oldSize);
                }
            }
        } finally {
            eventListenerManager.partialFlushEnd(
                    event.getNumberOfEntitiesProcessed(),
                    event.getNumberOfEntitiesProcessed()
            );
        }
    }

    private boolean flushIsReallyNeeded(AutoFlushEvent event, final EventSource source) {
        boolean flushIsReallyNeeded = source.getHibernateFlushMode() == FlushMode.ALWAYS
                // START OF FIX for auto-flush-mode on inserts that might later be deleted in the same transaction
                || source.getActionQueue().numberOfInsertions() > 0
                // END OF FIX
                || source.getActionQueue().areTablesToBeUpdated(event.getQuerySpaces());
        // START OF HOOK: clear query cache when data is flushed to the database
        if (flushIsReallyNeeded) {
            // clear the per-session query cache, as changing an entity might change any of the cached query results
            JpaMapStorage.clearQueryCache(source.getSession());
        }
        // END OF HOOK
        return flushIsReallyNeeded;
    }

    private boolean flushMightBeNeeded(final EventSource source) {
        final PersistenceContext persistenceContext = source.getPersistenceContextInternal();
        return !source.getHibernateFlushMode().lessThan(FlushMode.AUTO)
                && (persistenceContext.getNumberOfManagedEntities() > 0
                || persistenceContext.getCollectionEntriesSize() > 0);
    }

}