/*
 * 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.hotRod.userSession;

import org.keycloak.models.AuthenticatedClientSessionModel;
import org.keycloak.models.UserSessionModel;
import org.keycloak.models.map.common.AbstractEntity;
import org.keycloak.models.map.common.DeepCloner;
import org.keycloak.models.map.common.StringKeyConverter;
import org.keycloak.models.map.common.delegate.SimpleDelegateProvider;
import org.keycloak.models.map.storage.QueryParameters;
import org.keycloak.models.map.storage.CrudOperations;
import org.keycloak.models.map.storage.chm.ConcurrentHashMapCrudOperations;
import org.keycloak.models.map.storage.chm.ConcurrentHashMapStorage;
import org.keycloak.models.map.storage.chm.MapModelCriteriaBuilder;
import org.keycloak.models.map.storage.criteria.DefaultModelCriteria;
import org.keycloak.models.map.userSession.MapAuthenticatedClientSessionEntity;
import org.keycloak.models.map.userSession.MapAuthenticatedClientSessionEntityDelegate;
import org.keycloak.models.map.userSession.MapUserSessionEntity;
import org.keycloak.models.map.userSession.MapUserSessionEntityDelegate;
import org.keycloak.storage.SearchableModelField;

import java.util.Collections;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static org.keycloak.models.map.storage.ModelCriteriaBuilder.Operator.IN;

public class HotRodUserSessionMapStorage<K> extends ConcurrentHashMapStorage<K, MapUserSessionEntity, UserSessionModel, CrudOperations<MapUserSessionEntity, UserSessionModel>> {

    private final ConcurrentHashMapStorage<String, MapAuthenticatedClientSessionEntity, AuthenticatedClientSessionModel, CrudOperations<MapAuthenticatedClientSessionEntity, AuthenticatedClientSessionModel>> clientSessionStore;

    public HotRodUserSessionMapStorage(CrudOperations<MapUserSessionEntity, UserSessionModel> map,
                                       StringKeyConverter<K> keyConverter,
                                       DeepCloner cloner,
                                       Map<SearchableModelField<? super UserSessionModel>, MapModelCriteriaBuilder.UpdatePredicatesFunc<K, MapUserSessionEntity, UserSessionModel>> fieldPredicates,
                                       ConcurrentHashMapStorage<String, MapAuthenticatedClientSessionEntity, AuthenticatedClientSessionModel, CrudOperations<MapAuthenticatedClientSessionEntity, AuthenticatedClientSessionModel>> clientSessionStore
    ) {
        super(map, keyConverter, cloner, fieldPredicates);
        this.clientSessionStore = clientSessionStore;
    }

    private MapAuthenticatedClientSessionEntity wrapClientSessionEntityToClientSessionAwareDelegate(MapAuthenticatedClientSessionEntity d) {
        return new MapAuthenticatedClientSessionEntityDelegate(new HotRodAuthenticatedClientSessionEntityDelegateProvider(d) {
            @Override
            public MapAuthenticatedClientSessionEntity loadClientSessionFromDatabase() {
                return clientSessionStore.read(d.getId());
            }
        });
    }

    private MapUserSessionEntity wrapUserSessionEntityToClientSessionAwareDelegate(MapUserSessionEntity entity) {
        if (entity == null) return null;

        return new MapUserSessionEntityDelegate(new SimpleDelegateProvider<>(entity)) {

            private boolean filterAndRemoveNotExpired(MapAuthenticatedClientSessionEntity clientSession) {
                if (!clientSessionStore.exists(clientSession.getId())) {
                    // If client session does not exist, remove the reference to it from userSessionEntity loaded in this transaction
                    entity.removeAuthenticatedClientSession(clientSession.getClientId());
                    return false;
                }

                return true;
            }

            @Override
            public Set<MapAuthenticatedClientSessionEntity> getAuthenticatedClientSessions() {
                Set<MapAuthenticatedClientSessionEntity> clientSessions = super.getAuthenticatedClientSessions();
                return clientSessions == null ? null : clientSessions.stream()
                        // Find whether client session still exists in Infinispan and if not, remove the reference from user session
                        .filter(this::filterAndRemoveNotExpired)
                        .map(HotRodUserSessionMapStorage.this::wrapClientSessionEntityToClientSessionAwareDelegate)
                        .collect(Collectors.toSet());
            }

            @Override
            public Optional<MapAuthenticatedClientSessionEntity> getAuthenticatedClientSession(String clientUUID) {
                return super.getAuthenticatedClientSession(clientUUID)
                        // Find whether client session still exists in Infinispan and if not, remove the reference from user sessionZ
                        .filter(this::filterAndRemoveNotExpired)
                        .map(HotRodUserSessionMapStorage.this::wrapClientSessionEntityToClientSessionAwareDelegate);
            }

            @Override
            public void addAuthenticatedClientSession(MapAuthenticatedClientSessionEntity clientSession) {
                super.addAuthenticatedClientSession(clientSession);
                clientSessionStore.create(clientSession);
            }

            @Override
            public Boolean removeAuthenticatedClientSession(String clientUUID) {
                Optional<MapAuthenticatedClientSessionEntity> clientSession = getAuthenticatedClientSession(clientUUID);
                if (!clientSession.isPresent()) {
                    return false;
                }
                return super.removeAuthenticatedClientSession(clientUUID) && clientSessionStore.delete(clientSession.get().getId());
            }

            @Override
            public void clearAuthenticatedClientSessions() {
                Set<MapAuthenticatedClientSessionEntity> clientSessions = super.getAuthenticatedClientSessions();
                if (clientSessions != null) {
                    clientSessionStore.delete(QueryParameters.withCriteria(
                            DefaultModelCriteria.<AuthenticatedClientSessionModel>criteria()
                                    .compare(HotRodAuthenticatedClientSessionEntity.ID, IN, clientSessions.stream()
                                            .map(MapAuthenticatedClientSessionEntity::getId))
                    ));
                }
                super.clearAuthenticatedClientSessions();
            }
        };
    }


    @Override
    public MapUserSessionEntity read(String sKey) {
        return wrapUserSessionEntityToClientSessionAwareDelegate(super.read(sKey));
    }

    @Override
    public Stream<MapUserSessionEntity> read(QueryParameters<UserSessionModel> queryParameters) {
        return super.read(queryParameters).map(this::wrapUserSessionEntityToClientSessionAwareDelegate);
    }

    @Override
    public MapUserSessionEntity create(MapUserSessionEntity value) {
        return wrapUserSessionEntityToClientSessionAwareDelegate(super.create(value));
    }

    @Override
    public boolean delete(String key) {
        MapUserSessionEntity uSession = read(key);
        Set<MapAuthenticatedClientSessionEntity> clientSessions = uSession.getAuthenticatedClientSessions();
        if (clientSessions != null) {
            clientSessionStore.delete(QueryParameters.withCriteria(
                    DefaultModelCriteria.<AuthenticatedClientSessionModel>criteria()
                            .compare(HotRodAuthenticatedClientSessionEntity.ID, IN, clientSessions.stream()
                                    .map(MapAuthenticatedClientSessionEntity::getId))
            ));
        }

        return super.delete(key);
    }

    @Override
    public long delete(QueryParameters<UserSessionModel> queryParameters) {
        clientSessionStore.delete(QueryParameters.withCriteria(
                DefaultModelCriteria.<AuthenticatedClientSessionModel>criteria()
                        .compare(HotRodAuthenticatedClientSessionEntity.ID, IN, read(queryParameters)
                                .flatMap(userSession -> Optional.ofNullable(userSession.getAuthenticatedClientSessions()).orElse(Collections.emptySet()).stream().map(AbstractEntity::getId)))
        ));

        return super.delete(queryParameters);
    }
}
