/*
 * Copyright 2016 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.sessions.infinispan.changes;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import org.infinispan.Cache;
import org.infinispan.context.Flag;
import org.jboss.logging.Logger;
import org.keycloak.models.AbstractKeycloakTransaction;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.RealmModel;
import org.keycloak.models.sessions.infinispan.entities.SessionEntity;
import org.keycloak.models.sessions.infinispan.remotestore.RemoteCacheInvoker;

/**
 * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
 */
public class InfinispanChangelogBasedTransaction<S extends SessionEntity> extends AbstractKeycloakTransaction {

    public static final Logger logger = Logger.getLogger(InfinispanChangelogBasedTransaction.class);

    private final KeycloakSession kcSession;
    private final String cacheName;
    private final Cache<String, SessionEntityWrapper<S>> cache;
    private final RemoteCacheInvoker remoteCacheInvoker;

    private final Map<String, SessionUpdatesList<S>> updates = new HashMap<>();

    public InfinispanChangelogBasedTransaction(KeycloakSession kcSession, String cacheName, Cache<String, SessionEntityWrapper<S>> cache, RemoteCacheInvoker remoteCacheInvoker) {
        this.kcSession = kcSession;
        this.cacheName = cacheName;
        this.cache = cache;
        this.remoteCacheInvoker = remoteCacheInvoker;
    }


    public void addTask(String key, SessionUpdateTask<S> task) {
        SessionUpdatesList<S> myUpdates = updates.get(key);
        if (myUpdates == null) {
            // Lookup entity from cache
            SessionEntityWrapper<S> wrappedEntity = cache.get(key);
            if (wrappedEntity == null) {
                logger.warnf("Not present cache item for key %s", key);
                return;
            }

            RealmModel realm = kcSession.realms().getRealm(wrappedEntity.getEntity().getRealm());

            myUpdates = new SessionUpdatesList<>(realm, wrappedEntity);
            updates.put(key, myUpdates);
        }

        // Run the update now, so reader in same transaction can see it (TODO: Rollback may not work correctly. See if it's an issue..)
        task.runUpdate(myUpdates.getEntityWrapper().getEntity());
        myUpdates.add(task);
    }


    // Create entity and new version for it
    public void addTask(String key, SessionUpdateTask<S> task, S entity) {
        if (entity == null) {
            throw new IllegalArgumentException("Null entity not allowed");
        }

        RealmModel realm = kcSession.realms().getRealm(entity.getRealm());
        SessionEntityWrapper<S> wrappedEntity = new SessionEntityWrapper<>(entity);
        SessionUpdatesList<S> myUpdates = new SessionUpdatesList<>(realm, wrappedEntity);
        updates.put(key, myUpdates);

        // Run the update now, so reader in same transaction can see it
        task.runUpdate(entity);
        myUpdates.add(task);
    }


    public void reloadEntityInCurrentTransaction(RealmModel realm, String key, SessionEntityWrapper<S> entity) {
        if (entity == null) {
            throw new IllegalArgumentException("Null entity not allowed");
        }

        SessionEntityWrapper<S> latestEntity = cache.get(key);
        if (latestEntity == null) {
            return;
        }

        SessionUpdatesList<S> newUpdates = new SessionUpdatesList<>(realm, latestEntity);

        SessionUpdatesList<S> existingUpdates = updates.get(key);
        if (existingUpdates != null) {
            newUpdates.setUpdateTasks(existingUpdates.getUpdateTasks());
        }

        updates.put(key, newUpdates);
    }


    public SessionEntityWrapper<S> get(String key) {
        SessionUpdatesList<S> myUpdates = updates.get(key);
        if (myUpdates == null) {
            SessionEntityWrapper<S> wrappedEntity = cache.get(key);
            if (wrappedEntity == null) {
                return null;
            }

            RealmModel realm = kcSession.realms().getRealm(wrappedEntity.getEntity().getRealm());

            myUpdates = new SessionUpdatesList<>(realm, wrappedEntity);
            updates.put(key, myUpdates);

            return wrappedEntity;
        } else {
            return myUpdates.getEntityWrapper();
        }
    }


    @Override
    protected void commitImpl() {
        for (Map.Entry<String, SessionUpdatesList<S>> entry : updates.entrySet()) {
            SessionUpdatesList<S> sessionUpdates = entry.getValue();
            SessionEntityWrapper<S> sessionWrapper = sessionUpdates.getEntityWrapper();

            RealmModel realm = sessionUpdates.getRealm();

            MergedUpdate<S> merged = MergedUpdate.computeUpdate(sessionUpdates.getUpdateTasks(), sessionWrapper);

            if (merged != null) {
                // Now run the operation in our cluster
                runOperationInCluster(entry.getKey(), merged, sessionWrapper);

                // Check if we need to send message to second DC
                remoteCacheInvoker.runTask(kcSession, realm, cacheName, entry.getKey(), merged, sessionWrapper);
            }
        }
    }


    private void runOperationInCluster(String key, MergedUpdate<S> task,  SessionEntityWrapper<S> sessionWrapper) {
        S session = sessionWrapper.getEntity();
        SessionUpdateTask.CacheOperation operation = task.getOperation(session);

        // Don't need to run update of underlying entity. Local updates were already run
        //task.runUpdate(session);

        switch (operation) {
            case REMOVE:
                // Just remove it
                cache
                        .getAdvancedCache().withFlags(Flag.IGNORE_RETURN_VALUES)
                        .remove(key);
                break;
            case ADD:
                cache
                        .getAdvancedCache().withFlags(Flag.IGNORE_RETURN_VALUES)
                        .put(key, sessionWrapper, task.getLifespanMs(), TimeUnit.MILLISECONDS);
                break;
            case ADD_IF_ABSENT:
                SessionEntityWrapper existing = cache.putIfAbsent(key, sessionWrapper);
                if (existing != null) {
                    throw new IllegalStateException("There is already existing value in cache for key " + key);
                }
                break;
            case REPLACE:
                replace(key, task, sessionWrapper);
                break;
            default:
                throw new IllegalStateException("Unsupported state " +  operation);
        }

    }


    private void replace(String key, MergedUpdate<S> task, SessionEntityWrapper<S> oldVersionEntity) {
        boolean replaced = false;
        S session = oldVersionEntity.getEntity();

        while (!replaced) {
            SessionEntityWrapper<S> newVersionEntity = generateNewVersionAndWrapEntity(session, oldVersionEntity.getLocalMetadata());

            // Atomic cluster-aware replace
            replaced = cache.replace(key, oldVersionEntity, newVersionEntity);

            // Replace fail. Need to load latest entity from cache, apply updates again and try to replace in cache again
            if (!replaced) {
                logger.debugf("Replace failed for entity: %s . Will try again", key);

                oldVersionEntity = cache.get(key);

                if (oldVersionEntity == null) {
                    logger.debugf("Entity %s not found. Maybe removed in the meantime. Replace task will be ignored", key);
                    return;
                }

                session = oldVersionEntity.getEntity();

                task.runUpdate(session);
            } else {
                if (logger.isTraceEnabled()) {
                    logger.tracef("Replace SUCCESS for entity: %s . old version: %d, new version: %d", key, oldVersionEntity.getVersion(), newVersionEntity.getVersion());
                }
            }
        }

    }


    @Override
    protected void rollbackImpl() {
    }

    private SessionEntityWrapper<S> generateNewVersionAndWrapEntity(S entity, Map<String, String> localMetadata) {
        return new SessionEntityWrapper<>(localMetadata, entity);
    }

}
