/*
 * Copyright (c) 2010-2016. Axon Framework
 *
 * 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.axonframework.eventhandling.saga.repository.inmemory;

import org.axonframework.eventhandling.saga.AssociationValue;
import org.axonframework.eventhandling.saga.AssociationValues;
import org.axonframework.eventhandling.saga.repository.SagaStore;
import org.axonframework.eventsourcing.eventstore.TrackingToken;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.stream.Collectors;

/**
 * SagaRepository implementation that stores all Saga instances in memory.
 *
 * @author Allard Buijze
 * @since 0.7
 */
public class InMemorySagaStore implements SagaStore<Object> {

    private final ConcurrentMap<String, ManagedSaga> managedSagas = new ConcurrentHashMap<>();

    @Override
    public Set<String> findSagas(Class<?> sagaType, AssociationValue associationValue) {
        return managedSagas.entrySet()
                .stream()
                .filter(avEntry -> sagaType.isAssignableFrom(avEntry.getValue().saga().getClass()))
                .filter(avEntry -> avEntry.getValue().associationValues().contains(associationValue))
                .map(Map.Entry::getKey)
                .collect(Collectors.toSet());
    }

    @Override
    public void deleteSaga(Class<?> sagaType, String sagaIdentifier, Set<AssociationValue> associationValues) {
        managedSagas.remove(sagaIdentifier);
    }

    @SuppressWarnings("unchecked")
    @Override
    public <S> Entry<S> loadSaga(Class<S> sagaType, String sagaIdentifier) {
        return (Entry<S>) managedSagas.get(sagaIdentifier);
    }

    @Override
    public void insertSaga(Class<?> sagaType, String sagaIdentifier, Object saga, TrackingToken token, Set<AssociationValue> associationValues) {
        managedSagas.put(sagaIdentifier, new ManagedSaga(saga, associationValues));
    }

    @Override
    public void updateSaga(Class<?> sagaType, String sagaIdentifier, Object saga, TrackingToken token, AssociationValues associationValues) {
        managedSagas.put(sagaIdentifier, new ManagedSaga(saga, associationValues.asSet()));
    }

    /**
     * Returns the number of Sagas currently contained in this repository.
     *
     * @return the number of Sagas currently contained in this repository
     */
    public int size() {
        return managedSagas.size();
    }

    private static class ManagedSaga implements Entry<Object> {

        private final Object saga;
        private final Set<AssociationValue> associationValues;

        public ManagedSaga(Object saga, Set<AssociationValue> associationValues) {

            this.saga = saga;
            this.associationValues = associationValues;
        }

        @Override
        public TrackingToken trackingToken() {
            return null;
        }

        @Override
        public Set<AssociationValue> associationValues() {
            return associationValues;
        }

        @Override
        public Object saga() {
            return saga;
        }
    }
}

@@ -1599,22 +1599,16 @@
 e.is
-AssignableFrom
+Instance
 (avE
@@ -1633,19 +1633,8 @@
 ga()
-.getClass()
 ))%0A 

use isInstance(obj) instead of isAssignableFrom(obj.class)