/*
 * 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.mongo.keycloak.adapters;

import com.mongodb.BasicDBObject;
import com.mongodb.DBObject;
import com.mongodb.QueryBuilder;
import org.keycloak.connections.mongo.api.MongoStore;
import org.keycloak.connections.mongo.api.context.MongoStoreInvocationContext;
import org.keycloak.models.ClientModel;
import org.keycloak.models.ClientSessionModel;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.ModelException;
import org.keycloak.models.RealmModel;
import org.keycloak.models.UserModel;
import org.keycloak.models.UserSessionModel;
import org.keycloak.models.mongo.keycloak.entities.MongoOfflineUserSessionEntity;
import org.keycloak.models.mongo.keycloak.entities.MongoOnlineUserSessionEntity;
import org.keycloak.models.mongo.keycloak.entities.MongoUserSessionEntity;
import org.keycloak.models.mongo.keycloak.entities.PersistentClientSessionEntity;
import org.keycloak.models.mongo.keycloak.entities.PersistentUserSessionEntity;
import org.keycloak.models.session.PersistentClientSessionAdapter;
import org.keycloak.models.session.PersistentClientSessionModel;
import org.keycloak.models.session.PersistentUserSessionAdapter;
import org.keycloak.models.session.PersistentUserSessionModel;
import org.keycloak.models.session.UserSessionPersisterProvider;

import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;

/**
 * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
 */
public class MongoUserSessionPersisterProvider implements UserSessionPersisterProvider {

    private final MongoStoreInvocationContext invocationContext;
    private final KeycloakSession session;

    public MongoUserSessionPersisterProvider(KeycloakSession session, MongoStoreInvocationContext invocationContext) {
        this.session = session;
        this.invocationContext = invocationContext;
    }

    protected MongoStore getMongoStore() {
        return invocationContext.getMongoStore();
    }

    private MongoUserSessionEntity loadUserSession(String userSessionId, boolean offline) {
        Class<? extends MongoUserSessionEntity> clazz = offline ? MongoOfflineUserSessionEntity.class : MongoOnlineUserSessionEntity.class;
        return getMongoStore().loadEntity(clazz, userSessionId, invocationContext);
    }

    @Override
    public void createUserSession(UserSessionModel userSession, boolean offline) {
        PersistentUserSessionAdapter adapter = new PersistentUserSessionAdapter(userSession);
        PersistentUserSessionModel model = adapter.getUpdatedModel();

        MongoUserSessionEntity entity = offline ? new MongoOfflineUserSessionEntity() : new MongoOnlineUserSessionEntity();
        entity.setId(model.getUserSessionId());
        entity.setRealmId(adapter.getRealm().getId());
        entity.setUserId(adapter.getUser().getId());
        entity.setLastSessionRefresh(model.getLastSessionRefresh());
        entity.setData(model.getData());
        entity.setClientSessions(new ArrayList<PersistentClientSessionEntity>());
        getMongoStore().insertEntity(entity, invocationContext);
    }

    @Override
    public void createClientSession(ClientSessionModel clientSession, boolean offline) {
        PersistentClientSessionAdapter adapter = new PersistentClientSessionAdapter(clientSession);
        PersistentClientSessionModel model = adapter.getUpdatedModel();

        MongoUserSessionEntity userSession = loadUserSession(model.getUserSessionId(), offline);
        if (userSession == null) {
            throw new ModelException("Not userSession found with ID " + clientSession.getUserSession().getId() + ". Requested by clientSession: " + clientSession.getId());
        } else {
            PersistentClientSessionEntity entity = new PersistentClientSessionEntity();
            entity.setClientSessionId(clientSession.getId());
            entity.setClientId(clientSession.getClient().getId());
            entity.setData(model.getData());
            userSession.getClientSessions().add(entity);
            getMongoStore().updateEntity(userSession, invocationContext);
        }
    }

    @Override
    public void updateUserSession(UserSessionModel userSession, boolean offline) {
        PersistentUserSessionAdapter adapter;
        if (userSession instanceof PersistentUserSessionAdapter) {
            adapter = (PersistentUserSessionAdapter) userSession;
        } else {
            adapter = new PersistentUserSessionAdapter(userSession);
        }

        PersistentUserSessionModel model = adapter.getUpdatedModel();

        MongoUserSessionEntity entity = loadUserSession(model.getUserSessionId(), offline);
        if (entity == null) {
            throw new ModelException("UserSession with ID " + userSession.getId() + ", offline: " + offline + " not found");
        }
        entity.setLastSessionRefresh(model.getLastSessionRefresh());
        entity.setData(model.getData());

        getMongoStore().updateEntity(entity, invocationContext);
    }

    @Override
    public void removeUserSession(String userSessionId, boolean offline) {
        MongoUserSessionEntity entity = loadUserSession(userSessionId, offline);
        if (entity != null) {
            getMongoStore().removeEntity(entity, invocationContext);
        }
    }

    @Override
    public void removeClientSession(String clientSessionId, boolean offline) {
        DBObject query = new QueryBuilder()
                .and("clientSessions.clientSessionId").is(clientSessionId)
                .get();
        Class<? extends MongoUserSessionEntity> clazz = offline ? MongoOfflineUserSessionEntity.class : MongoOnlineUserSessionEntity.class;
        MongoUserSessionEntity userSession = getMongoStore().loadSingleEntity(clazz, query, invocationContext);
        if (userSession != null) {

            PersistentClientSessionEntity found = null;
            for (PersistentClientSessionEntity clientSession : userSession.getClientSessions()) {
                if (clientSession.getClientSessionId().equals(clientSessionId)) {
                    found = clientSession;
                    break;
                }
            }

            if (found != null) {
                userSession.getClientSessions().remove(found);

                // Remove userSession if it was last clientSession attached
                if (userSession.getClientSessions().size() == 0) {
                    getMongoStore().removeEntity(userSession, invocationContext);
                } else {
                    getMongoStore().updateEntity(userSession, invocationContext);
                }
            }
        }
    }

    @Override
    public void onRealmRemoved(RealmModel realm) {
        DBObject query = new QueryBuilder()
                .and("realmId").is(realm.getId())
                .get();
        getMongoStore().removeEntities(MongoOnlineUserSessionEntity.class, query, false, invocationContext);
        getMongoStore().removeEntities(MongoOfflineUserSessionEntity.class, query, false, invocationContext);
    }

    @Override
    public void onClientRemoved(RealmModel realm, ClientModel client) {
        DBObject query = new QueryBuilder()
                .and("clientSessions.clientId").is(client.getId())
                .get();

        List<MongoOnlineUserSessionEntity> userSessions = getMongoStore().loadEntities(MongoOnlineUserSessionEntity.class, query, invocationContext);
        for (MongoOnlineUserSessionEntity userSession : userSessions) {
            removeClientSessionOfClient(userSession, client.getId());
        }

        List<MongoOfflineUserSessionEntity> userSessions2 = getMongoStore().loadEntities(MongoOfflineUserSessionEntity.class, query, invocationContext);
        for (MongoOfflineUserSessionEntity userSession : userSessions2) {
            removeClientSessionOfClient(userSession, client.getId());
        }
    }

    private void removeClientSessionOfClient(MongoUserSessionEntity userSession, String clientId) {
        PersistentClientSessionEntity found = null;
        for (PersistentClientSessionEntity clientSession : userSession.getClientSessions()) {
            if (clientSession.getClientId().equals(clientId)) {
                found = clientSession;
                break;
            }
        }

        if (found != null) {
            userSession.getClientSessions().remove(found);

            // Remove userSession if it was last clientSession attached
            if (userSession.getClientSessions().size() == 0) {
                getMongoStore().removeEntity(userSession, invocationContext);
            } else {
                getMongoStore().updateEntity(userSession, invocationContext);
            }
        }
    }

    @Override
    public void onUserRemoved(RealmModel realm, UserModel user) {
        onUserRemoved(realm, user.getId());
    }

    private void onUserRemoved(RealmModel realm, String userId) {
        DBObject query = new QueryBuilder()
                .and("userId").is(userId)
                .get();
        getMongoStore().removeEntities(MongoOnlineUserSessionEntity.class, query, false, invocationContext);
        getMongoStore().removeEntities(MongoOfflineUserSessionEntity.class, query, false, invocationContext);
    }

    @Override
    public void clearDetachedUserSessions() {
        DBObject query = new QueryBuilder()
                .and("clientSessions").is(Collections.emptyList())
                .get();
        getMongoStore().removeEntities(MongoOnlineUserSessionEntity.class, query, false, invocationContext);
        getMongoStore().removeEntities(MongoOfflineUserSessionEntity.class, query, false, invocationContext);
    }

    @Override
    public int getUserSessionsCount(boolean offline) {
        DBObject query = new QueryBuilder()
                .get();

        Class<? extends MongoUserSessionEntity> clazz = offline ? MongoOfflineUserSessionEntity.class : MongoOnlineUserSessionEntity.class;
        return getMongoStore().countEntities(clazz, query, invocationContext);
    }

    @Override
    public void updateAllTimestamps(int time) {
        // 1) Update timestamp of clientSessions

        DBObject timestampSubquery = new QueryBuilder()
                .and("timestamp").notEquals(time).get();

        DBObject query = new QueryBuilder()
                .and("clientSessions").elemMatch(timestampSubquery).get();


        DBObject update = new QueryBuilder()
                .and("$set").is(new BasicDBObject("clientSessions.$.timestamp", time)).get();

        // Not sure how to do in single query :/
        int countModified = 1;
        while (countModified > 0) {
            countModified = getMongoStore().updateEntities(MongoOfflineUserSessionEntity.class, query, update, invocationContext);
        }

        countModified = 1;
        while (countModified > 0) {
            countModified = getMongoStore().updateEntities(MongoOnlineUserSessionEntity.class, query, update, invocationContext);
        }

        // 2) update lastSessionRefresh of userSessions
        query = new QueryBuilder().get();

        update = new QueryBuilder()
                .and("$set").is(new BasicDBObject("lastSessionRefresh", time)).get();

        getMongoStore().updateEntities(MongoOfflineUserSessionEntity.class, query, update, invocationContext);
        getMongoStore().updateEntities(MongoOnlineUserSessionEntity.class, query, update, invocationContext);
    }

    @Override
    public List<UserSessionModel> loadUserSessions(int firstResult, int maxResults, boolean offline) {
        DBObject query = new QueryBuilder()
                .get();
        DBObject sort = new BasicDBObject("id", 1);

        Class<? extends MongoUserSessionEntity> clazz = offline ? MongoOfflineUserSessionEntity.class : MongoOnlineUserSessionEntity.class;

        List<? extends MongoUserSessionEntity> entities = getMongoStore().loadEntities(clazz, query, sort, firstResult, maxResults, invocationContext);

        List<UserSessionModel> results = new LinkedList<>();
        for (MongoUserSessionEntity entity : entities) {
            RealmModel realm = session.realms().getRealm(entity.getRealmId());
            UserModel user = session.users().getUserById(entity.getUserId(), realm);

            // Case when user was deleted in the meantime
            if (user == null) {
                onUserRemoved(realm, entity.getUserId());
                return loadUserSessions(firstResult, maxResults, offline);
            }

            PersistentUserSessionAdapter userSession = toAdapter(realm, user, entity);
            results.add(userSession);
        }
        return results;
    }

    private PersistentUserSessionAdapter toAdapter(RealmModel realm, UserModel user, PersistentUserSessionEntity entity) {
        PersistentUserSessionModel model = new PersistentUserSessionModel();
        model.setUserSessionId(entity.getId());
        model.setLastSessionRefresh(entity.getLastSessionRefresh());
        model.setData(entity.getData());

        List<ClientSessionModel> clientSessions = new LinkedList<>();
        PersistentUserSessionAdapter userSessionAdapter = new PersistentUserSessionAdapter(model, realm, user, clientSessions);
        for (PersistentClientSessionEntity clientSessEntity : entity.getClientSessions()) {
            PersistentClientSessionAdapter clientSessAdapter = toAdapter(realm, userSessionAdapter, clientSessEntity);
            clientSessions.add(clientSessAdapter);
        }

        return userSessionAdapter;
    }

    private PersistentClientSessionAdapter toAdapter(RealmModel realm, PersistentUserSessionAdapter userSession, PersistentClientSessionEntity entity) {
        ClientModel client = realm.getClientById(entity.getClientId());

        PersistentClientSessionModel model = new PersistentClientSessionModel();
        model.setClientSessionId(entity.getClientSessionId());
        model.setClientId(entity.getClientId());
        model.setUserSessionId(userSession.getId());
        model.setUserId(userSession.getUser().getId());
        model.setTimestamp(entity.getTimestamp());
        model.setData(entity.getData());
        return new PersistentClientSessionAdapter(model, realm, client, userSession);
    }

    @Override
    public void close() {

    }
}
