/*
 * Copyright 2021 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;

import org.jboss.logging.Logger;
import org.keycloak.Config;
import org.keycloak.common.Profile;
import org.keycloak.component.AmphibianProviderFactory;
import org.keycloak.models.ActionTokenValueModel;
import org.keycloak.models.AuthenticatedClientSessionModel;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.KeycloakSessionFactory;
import org.keycloak.models.UserSessionModel;
import org.keycloak.models.map.authorization.entity.MapPermissionTicketEntity;
import org.keycloak.models.map.authorization.entity.MapPolicyEntity;
import org.keycloak.models.map.authorization.entity.MapResourceEntity;
import org.keycloak.models.map.authorization.entity.MapResourceServerEntity;
import org.keycloak.models.map.authorization.entity.MapScopeEntity;
import org.keycloak.models.map.authSession.MapAuthenticationSessionEntity;
import org.keycloak.models.map.authSession.MapRootAuthenticationSessionEntity;
import org.keycloak.models.map.clientscope.MapClientScopeEntity;
import org.keycloak.models.map.common.AbstractEntity;
import org.keycloak.models.map.common.StringKeyConverter;
import org.keycloak.models.map.events.MapAdminEventEntity;
import org.keycloak.models.map.events.MapAuthEventEntity;
import org.keycloak.models.map.group.MapGroupEntity;
import org.keycloak.models.map.loginFailure.MapUserLoginFailureEntity;
import org.keycloak.models.map.realm.MapRealmEntity;
import org.keycloak.models.map.realm.entity.MapAuthenticationExecutionEntity;
import org.keycloak.models.map.realm.entity.MapAuthenticationFlowEntity;
import org.keycloak.models.map.realm.entity.MapAuthenticatorConfigEntity;
import org.keycloak.models.map.realm.entity.MapClientInitialAccessEntity;
import org.keycloak.models.map.realm.entity.MapComponentEntity;
import org.keycloak.models.map.realm.entity.MapIdentityProviderEntity;
import org.keycloak.models.map.realm.entity.MapIdentityProviderMapperEntity;
import org.keycloak.models.map.realm.entity.MapOTPPolicyEntity;
import org.keycloak.models.map.realm.entity.MapRequiredActionProviderEntity;
import org.keycloak.models.map.realm.entity.MapRequiredCredentialEntity;
import org.keycloak.models.map.realm.entity.MapWebAuthnPolicyEntity;
import org.keycloak.models.map.role.MapRoleEntity;
import org.keycloak.models.map.singleUseObject.MapSingleUseObjectEntity;
import org.keycloak.models.map.storage.MapKeycloakTransaction;
import org.keycloak.models.map.storage.chm.ConcurrentHashMapKeycloakTransaction;
import org.keycloak.models.map.storage.chm.MapFieldPredicates;
import org.keycloak.models.map.storage.chm.MapModelCriteriaBuilder;
import org.keycloak.models.map.storage.hotRod.authSession.HotRodAuthenticationSessionEntityDelegate;
import org.keycloak.models.map.storage.hotRod.authSession.HotRodRootAuthenticationSessionEntityDelegate;
import org.keycloak.models.map.storage.hotRod.authorization.HotRodPermissionTicketEntityDelegate;
import org.keycloak.models.map.storage.hotRod.authorization.HotRodPolicyEntityDelegate;
import org.keycloak.models.map.storage.hotRod.authorization.HotRodResourceEntityDelegate;
import org.keycloak.models.map.storage.hotRod.authorization.HotRodResourceServerEntityDelegate;
import org.keycloak.models.map.storage.hotRod.authorization.HotRodScopeEntityDelegate;
import org.keycloak.models.map.storage.hotRod.common.AbstractHotRodEntity;
import org.keycloak.models.map.storage.hotRod.common.AutogeneratedHotRodDescriptors;
import org.keycloak.models.map.storage.hotRod.common.HotRodEntityDelegate;
import org.keycloak.models.map.storage.hotRod.events.HotRodAdminEventEntityDelegate;
import org.keycloak.models.map.storage.hotRod.events.HotRodAuthEventEntityDelegate;
import org.keycloak.models.map.storage.hotRod.loginFailure.HotRodUserLoginFailureEntityDelegate;
import org.keycloak.models.map.storage.hotRod.role.HotRodRoleEntityDelegate;
import org.keycloak.models.map.storage.hotRod.client.HotRodClientEntityDelegate;
import org.keycloak.models.map.storage.hotRod.client.HotRodProtocolMapperEntityDelegate;
import org.keycloak.models.map.client.MapClientEntity;
import org.keycloak.models.map.client.MapProtocolMapperEntity;
import org.keycloak.models.map.common.DeepCloner;
import org.keycloak.models.map.storage.hotRod.clientscope.HotRodClientScopeEntityDelegate;
import org.keycloak.models.map.storage.hotRod.common.HotRodEntityDescriptor;
import org.keycloak.models.map.storage.hotRod.connections.HotRodConnectionProvider;
import org.keycloak.models.map.storage.MapStorageProvider;
import org.keycloak.models.map.storage.MapStorageProviderFactory;
import org.keycloak.models.map.storage.hotRod.group.HotRodGroupEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.HotRodRealmEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodAuthenticationExecutionEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodAuthenticationFlowEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodAuthenticatorConfigEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodClientInitialAccessEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodComponentEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodIdentityProviderEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodIdentityProviderMapperEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodOTPPolicyEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodRequiredActionProviderEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodRequiredCredentialEntityDelegate;
import org.keycloak.models.map.storage.hotRod.realm.entity.HotRodWebAuthnPolicyEntityDelegate;
import org.keycloak.models.map.storage.hotRod.singleUseObject.HotRodSingleUseObjectEntityDelegate;
import org.keycloak.models.map.storage.hotRod.user.HotRodUserConsentEntityDelegate;
import org.keycloak.models.map.storage.hotRod.user.HotRodUserCredentialEntityDelegate;
import org.keycloak.models.map.storage.hotRod.user.HotRodUserEntityDelegate;
import org.keycloak.models.map.storage.hotRod.user.HotRodUserFederatedIdentityEntityDelegate;
import org.keycloak.models.map.storage.hotRod.userSession.HotRodAuthenticatedClientSessionEntity;
import org.keycloak.models.map.storage.hotRod.userSession.HotRodAuthenticatedClientSessionEntityDelegate;
import org.keycloak.models.map.storage.hotRod.userSession.HotRodUserSessionEntityDelegate;
import org.keycloak.models.map.storage.hotRod.userSession.HotRodUserSessionTransaction;
import org.keycloak.models.map.user.MapUserConsentEntity;
import org.keycloak.models.map.user.MapUserCredentialEntity;
import org.keycloak.models.map.user.MapUserEntity;
import org.keycloak.models.map.user.MapUserFederatedIdentityEntity;
import org.keycloak.models.map.userSession.MapAuthenticatedClientSessionEntity;
import org.keycloak.models.map.userSession.MapUserSessionEntity;
import org.keycloak.provider.EnvironmentDependentProviderFactory;
import org.keycloak.storage.SearchableModelField;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class HotRodMapStorageProviderFactory implements AmphibianProviderFactory<MapStorageProvider>, MapStorageProviderFactory, EnvironmentDependentProviderFactory {

    public static final String PROVIDER_ID = "hotrod";
    private final Map<Class<?>, HotRodMapStorage> storages = new ConcurrentHashMap<>();

    private static final Map<SearchableModelField<AuthenticatedClientSessionModel>, MapModelCriteriaBuilder.UpdatePredicatesFunc<Object, AbstractEntity, AuthenticatedClientSessionModel>> clientSessionPredicates = MapFieldPredicates.basePredicates(HotRodAuthenticatedClientSessionEntity.ID);

    private final static DeepCloner CLONER = new DeepCloner.Builder()
            .constructor(MapRootAuthenticationSessionEntity.class,  HotRodRootAuthenticationSessionEntityDelegate::new)
            .constructor(MapAuthenticationSessionEntity.class,      HotRodAuthenticationSessionEntityDelegate::new)

            .constructor(MapClientEntity.class,                     HotRodClientEntityDelegate::new)
            .constructor(MapProtocolMapperEntity.class,             HotRodProtocolMapperEntityDelegate::new)

            .constructor(MapClientScopeEntity.class,                HotRodClientScopeEntityDelegate::new)

            .constructor(MapGroupEntity.class,                      HotRodGroupEntityDelegate::new)

            .constructor(MapRoleEntity.class,                       HotRodRoleEntityDelegate::new)

            .constructor(MapSingleUseObjectEntity.class,            HotRodSingleUseObjectEntityDelegate::new)

            .constructor(MapUserEntity.class,                       HotRodUserEntityDelegate::new)
            .constructor(MapUserCredentialEntity.class,             HotRodUserCredentialEntityDelegate::new)
            .constructor(MapUserFederatedIdentityEntity.class,      HotRodUserFederatedIdentityEntityDelegate::new)
            .constructor(MapUserConsentEntity.class,                HotRodUserConsentEntityDelegate::new)

            .constructor(MapUserLoginFailureEntity.class,           HotRodUserLoginFailureEntityDelegate::new)

            .constructor(MapRealmEntity.class,                      HotRodRealmEntityDelegate::new)
            .constructor(MapAuthenticationExecutionEntity.class,    HotRodAuthenticationExecutionEntityDelegate::new)
            .constructor(MapAuthenticationFlowEntity.class,         HotRodAuthenticationFlowEntityDelegate::new)
            .constructor(MapAuthenticatorConfigEntity.class,        HotRodAuthenticatorConfigEntityDelegate::new)
            .constructor(MapClientInitialAccessEntity.class,        HotRodClientInitialAccessEntityDelegate::new)
            .constructor(MapComponentEntity.class,                  HotRodComponentEntityDelegate::new)
            .constructor(MapIdentityProviderEntity.class,           HotRodIdentityProviderEntityDelegate::new)
            .constructor(MapIdentityProviderMapperEntity.class,     HotRodIdentityProviderMapperEntityDelegate::new)
            .constructor(MapOTPPolicyEntity.class,                  HotRodOTPPolicyEntityDelegate::new)
            .constructor(MapRequiredActionProviderEntity.class,     HotRodRequiredActionProviderEntityDelegate::new)
            .constructor(MapRequiredCredentialEntity.class,         HotRodRequiredCredentialEntityDelegate::new)
            .constructor(MapWebAuthnPolicyEntity.class,             HotRodWebAuthnPolicyEntityDelegate::new)

            .constructor(MapUserSessionEntity.class,                HotRodUserSessionEntityDelegate::new)
            .constructor(MapAuthenticatedClientSessionEntity.class, HotRodAuthenticatedClientSessionEntityDelegate::new)

            .constructor(MapResourceServerEntity.class,             HotRodResourceServerEntityDelegate::new)
            .constructor(MapResourceEntity.class,                   HotRodResourceEntityDelegate::new)
            .constructor(MapScopeEntity.class,                      HotRodScopeEntityDelegate::new)
            .constructor(MapPolicyEntity.class,                     HotRodPolicyEntityDelegate::new)
            .constructor(MapPermissionTicketEntity.class,           HotRodPermissionTicketEntityDelegate::new)

            .constructor(MapAuthEventEntity.class,                  HotRodAuthEventEntityDelegate::new)
            .constructor(MapAdminEventEntity.class,                 HotRodAdminEventEntityDelegate::new)

            .build();

    @Override
    public MapStorageProvider create(KeycloakSession session) {
        return new HotRodMapStorageProvider(session, this);
    }

    public HotRodEntityDescriptor<?, ?> getEntityDescriptor(Class<?> c) {
        return AutogeneratedHotRodDescriptors.ENTITY_DESCRIPTOR_MAP.get(c);
    }

    public <E extends AbstractHotRodEntity, V extends HotRodEntityDelegate<E> & AbstractEntity, M> HotRodMapStorage<String, E, V, M> getHotRodStorage(KeycloakSession session, Class<M> modelType, MapStorageProviderFactory.Flag... flags) {
        // We need to preload client session store before we load user session store to avoid recursive update of storages map
        if (modelType == UserSessionModel.class) getHotRodStorage(session, AuthenticatedClientSessionModel.class, flags);

        return storages.computeIfAbsent(modelType, c -> createHotRodStorage(session, modelType, flags));
    }

    private <E extends AbstractHotRodEntity, V extends HotRodEntityDelegate<E> & AbstractEntity, M> HotRodMapStorage<String, E, V, M> createHotRodStorage(KeycloakSession session, Class<M> modelType, MapStorageProviderFactory.Flag... flags) {
        HotRodConnectionProvider connectionProvider = session.getProvider(HotRodConnectionProvider.class);
        HotRodEntityDescriptor<E, V> entityDescriptor = (HotRodEntityDescriptor<E, V>) getEntityDescriptor(modelType);

        if (modelType == ActionTokenValueModel.class) {
            return new SingleUseObjectHotRodMapStorage(connectionProvider.getRemoteCache(entityDescriptor.getCacheName()), StringKeyConverter.StringKey.INSTANCE, entityDescriptor, CLONER);
        } if (modelType == AuthenticatedClientSessionModel.class) {
            return new HotRodMapStorage(connectionProvider.getRemoteCache(entityDescriptor.getCacheName()),
                    StringKeyConverter.StringKey.INSTANCE,
                    entityDescriptor,
                    CLONER) {
                @Override
                protected MapKeycloakTransaction createTransactionInternal(KeycloakSession session) {
                    return new ConcurrentHashMapKeycloakTransaction(this, keyConverter, cloner, clientSessionPredicates);
                }
            };
        } if (modelType == UserSessionModel.class) {
            HotRodMapStorage clientSessionStore = getHotRodStorage(session, AuthenticatedClientSessionModel.class);
            return new HotRodMapStorage(connectionProvider.getRemoteCache(entityDescriptor.getCacheName()),
                    StringKeyConverter.StringKey.INSTANCE,
                    entityDescriptor,
                    CLONER) {
                @Override
                protected MapKeycloakTransaction createTransactionInternal(KeycloakSession session) {
                    Map<SearchableModelField<? super UserSessionModel>, MapModelCriteriaBuilder.UpdatePredicatesFunc<String, MapUserSessionEntity, UserSessionModel>> fieldPredicates = MapFieldPredicates.getPredicates((Class<UserSessionModel>) storedEntityDescriptor.getModelTypeClass());
                    return new HotRodUserSessionTransaction(this, keyConverter, cloner, fieldPredicates, clientSessionStore.createTransaction(session));
                }
            };
        }
        return new HotRodMapStorage<>(connectionProvider.getRemoteCache(entityDescriptor.getCacheName()), StringKeyConverter.StringKey.INSTANCE, entityDescriptor, CLONER);
    }

    @Override
    public void init(Config.Scope config) {

    }

    @Override
    public void postInit(KeycloakSessionFactory factory) {

    }

    @Override
    public String getId() {
        return PROVIDER_ID;
    }

    @Override
    public boolean isSupported() {
        return Profile.isFeatureEnabled(Profile.Feature.MAP_STORAGE);
    }

    @Override
    public String getHelpText() {
        return "HotRod map storage";
    }
}
