/*
 * 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.testsuite.util.cli;

import org.infinispan.AdvancedCache;
import org.infinispan.Cache;
import org.infinispan.context.Flag;
import org.keycloak.common.util.Time;
import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
import org.keycloak.models.ClientModel;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.RealmModel;
import org.keycloak.models.UserModel;
import org.keycloak.models.UserSessionModel;
import org.keycloak.models.sessions.infinispan.changes.SessionEntityWrapper;
import org.keycloak.models.sessions.infinispan.entities.SessionEntity;
import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity;
import org.keycloak.models.utils.KeycloakModelUtils;
import org.keycloak.services.managers.UserSessionManager;

import java.util.Arrays;
import java.util.Set;
import java.util.TreeSet;

/**
 * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
 */
public abstract class AbstractSessionCacheCommand extends AbstractCommand {

    private static final Set<String> SUPPORTED_CACHE_NAMES = new TreeSet<>(Arrays.asList(
      InfinispanConnectionProvider.USER_SESSION_CACHE_NAME,
      InfinispanConnectionProvider.OFFLINE_USER_SESSION_CACHE_NAME,
      InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME,
      InfinispanConnectionProvider.OFFLINE_CLIENT_SESSION_CACHE_NAME
    ));

    @Override
    protected void doRunCommand(KeycloakSession session) {
        InfinispanConnectionProvider provider = session.getProvider(InfinispanConnectionProvider.class);
        String cacheName = getArg(0);
        if (! SUPPORTED_CACHE_NAMES.contains(cacheName)) {
            log.errorf("Invalid cache name: '%s', Only cache names '%s' are supported", cacheName, SUPPORTED_CACHE_NAMES);
            throw new HandledException();
        }

        Cache<String, SessionEntityWrapper> ispnCache = provider.getCache(cacheName);
        doRunCacheCommand(session, ispnCache);

        ispnCache.entrySet().stream().skip(0).limit(10).collect(java.util.stream.Collectors.toMap(new java.util.function.Function() {

            public Object apply(Object entry) {
                return ((java.util.Map.Entry) entry).getKey();
            }
        }, new java.util.function.Function() {

            public Object apply(Object entry) {
                return ((java.util.Map.Entry) entry).getValue();
            }
        }));
    }

    protected void printSession(String id, UserSessionEntity userSession) {
        if (userSession == null) {
            log.info("Not found session with Id: " + id);
        } else {
            log.info("Found session. ID: " + toString(userSession));
        }
    }

    protected String toString(UserSessionEntity userSession) {
        int clientSessionsSize = userSession.getClientSessions().size();
        return "ID: " + userSession.getId() + ", realm: " + userSession.getRealmId()+ ", lastAccessTime: " + Time.toDate(userSession.getLastSessionRefresh()) +
                ", authenticatedClientSessions: " + clientSessionsSize;
    }

    @Override
    public String printUsage() {
        return getName() + " <cache-name>";
    }

    protected abstract void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache);


    // IMPLS

    public static class PutCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "put";
        }

        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            UserSessionEntity userSession = new UserSessionEntity(getArg(1));
            userSession.setRealmId(getArg(2));

            userSession.setLastSessionRefresh(Time.currentTime());
            cache.put(userSession.getId(), new SessionEntityWrapper(userSession));
        }

        @Override
        public String printUsage() {
            return getName() + " <cache-name> <user-session-id> <realm-name>";
        }
    }


    public static class GetCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "get";
        }

        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            String id = getArg(1);
            UserSessionEntity userSession = (UserSessionEntity) cache.get(id).getEntity();
            printSession(id, userSession);
        }

        @Override
        public String printUsage() {
            return getName() + " <cache-name> <user-session-id>";
        }
    }

    // Just to check performance of multiple get calls. And comparing what's the change between the case when item is available locally or not.
    public static class GetMultipleCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "getMulti";
        }

        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            String id = getArg(1);
            int count = getIntArg(2);

            long start = System.currentTimeMillis();
            for (int i=0 ; i<count ; i++) {
                UserSessionEntity userSession = (UserSessionEntity) cache.get(id).getEntity();
                //printSession(id, userSession);
            }
            long took = System.currentTimeMillis() - start;
            log.infof("Took %d milliseconds", took);
        }

        @Override
        public String printUsage() {
            return getName() + " <cache-name> <user-session-id> <count-of-gets>";
        }
    }


    public static class RemoveCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "remove";
        }

        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            String id = getArg(1);
            cache.remove(id);
        }

        @Override
        public String printUsage() {
            return getName() + " <cache-name> <user-session-id>";
        }
    }


    public static class ClearCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "clear";
        }

        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            cache.clear();
        }
    }


    public static class SizeCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "size";
        }

        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            log.info("Size: " + cache.size());
        }
    }


    public static class ListCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "list";
        }

        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            for (var entry : cache.entrySet()) {
                SessionEntity entity = entry.getValue().getEntity();
                if (!(entity instanceof UserSessionEntity)) {
                    continue;
                }
                UserSessionEntity userSession = (UserSessionEntity) entry.getValue().getEntity();
                log.info("list: key=" + entry.getKey() + ", value=" + toString(userSession));
            }
        }
    }


    public static class GetLocalCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "getLocal";
        }


        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            String id = getArg(1);
            cache = ((AdvancedCache) cache).withFlags(Flag.CACHE_MODE_LOCAL);
            UserSessionEntity userSession = (UserSessionEntity) cache.get(id).getEntity();
            printSession(id, userSession);
        }

        @Override
        public String printUsage() {
            return getName() + " <cache-name> <user-session-id>";
        }
    }


    public static class SizeLocalCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "sizeLocal";
        }

        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            log.info("Size local: " + cache.getAdvancedCache().withFlags(Flag.CACHE_MODE_LOCAL).size());
        }
    }


    public static class CreateManySessionsCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "createManySessions";
        }

        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            String realmName = getArg(1);
            int count = getIntArg(2);
            int batchCount = getIntArg(3);

            BatchTaskRunner.runInBatches(0, count, batchCount, session.getKeycloakSessionFactory(), (KeycloakSession batchSession, int firstInIteration, int countInIteration) -> {
                for (int i=0 ; i<countInIteration ; i++) {
                    UserSessionEntity userSession = new UserSessionEntity(KeycloakModelUtils.generateId());
                    userSession.setRealmId(realmName);

                    userSession.setLastSessionRefresh(Time.currentTime());
                    cache.put(userSession.getId(), new SessionEntityWrapper(userSession));
                }

                log.infof("Created '%d' sessions started from offset '%d'", countInIteration, firstInIteration);
            });

            log.infof("Created all '%d' sessions", count);
        }

        @Override
        public String printUsage() {
            return getName() + " <cache-name> <realm-name> <count> <count-in-batch>";
        }

    }


    // This will propagate creating sessions to remoteCache too
    public static class CreateManySessionsProviderCommand extends AbstractSessionCacheCommand {

        @Override
        public String getName() {
            return "createManySessionsProvider";
        }

        @Override
        protected void doRunCacheCommand(KeycloakSession session, Cache<String, SessionEntityWrapper> cache) {
            String realmName = getArg(1);
            String clientId = getArg(2);
            String username = getArg(3);
            int count = getIntArg(4);
            int batchCount = getIntArg(5);

            BatchTaskRunner.runInBatches(0, count, batchCount, session.getKeycloakSessionFactory(), (KeycloakSession batchSession, int firstInIteration, int countInIteration) -> {
                RealmModel realm = batchSession.realms().getRealmByName(realmName);
                ClientModel client = realm.getClientByClientId(clientId);
                UserModel user = batchSession.users().getUserByUsername(realm, username);

                UserSessionManager userSessionManager = new UserSessionManager(session);
                for (int i=0 ; i<countInIteration ; i++) {
                    UserSessionModel userSession = userSessionManager.createUserSession(realm, user, username, "127.0.0.1", "form", false, null, null);

                    session.sessions().createClientSession(userSession.getRealm(), client, userSession);
                }

                log.infof("Created '%d' sessions started from offset '%d'", countInIteration, firstInIteration);
            });

            log.infof("Created all '%d' sessions", count);
        }

        @Override
        public String printUsage() {
            return getName() + " <cache-name> <realm-name> <client-id> <user-name> <count> <count-in-batch>";
        }

    }

}
