package org.infinispan.util;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.infinispan.commands.CommandsFactory;
import org.infinispan.commands.functional.ReadWriteKeyCommand;
import org.infinispan.commands.functional.ReadWriteKeyValueCommand;
import org.infinispan.commands.functional.ReadWriteManyCommand;
import org.infinispan.commands.functional.ReadWriteManyEntriesCommand;
import org.infinispan.commands.functional.WriteOnlyKeyCommand;
import org.infinispan.commands.functional.WriteOnlyKeyValueCommand;
import org.infinispan.commands.functional.WriteOnlyManyCommand;
import org.infinispan.commands.functional.WriteOnlyManyEntriesCommand;
import org.infinispan.commands.triangle.BackupWriteCommand;
import org.infinispan.commands.triangle.MultiEntriesFunctionalBackupWriteCommand;
import org.infinispan.commands.triangle.MultiKeyFunctionalBackupWriteCommand;
import org.infinispan.commands.triangle.PutMapBackupWriteCommand;
import org.infinispan.commands.triangle.SingleKeyBackupWriteCommand;
import org.infinispan.commands.triangle.SingleKeyFunctionalBackupWriteCommand;
import org.infinispan.commands.write.PutKeyValueCommand;
import org.infinispan.commands.write.PutMapCommand;
import org.infinispan.commands.write.RemoveCommand;
import org.infinispan.commands.write.RemoveExpiredCommand;
import org.infinispan.commands.write.ReplaceCommand;
import org.infinispan.commons.util.InfinispanCollections;
import org.infinispan.distribution.LocalizedCacheTopology;
import org.infinispan.remoting.responses.Response;
import org.infinispan.remoting.responses.ValidResponse;
import org.infinispan.remoting.transport.Address;

/**
 * Some utility functions for {@link org.infinispan.interceptors.distribution.TriangleDistributionInterceptor}.
 *
 * @author Pedro Ruivo
 * @since 8.4
 */
public final class TriangleFunctionsUtil {

   private TriangleFunctionsUtil() {
   }

   public static PutMapCommand copy(PutMapCommand command, Collection<Object> keys) {
      PutMapCommand copy = new PutMapCommand(command, false);
      copy.setMap(filterEntries(command.getMap(), keys));
      return copy;
   }

   public static <K, V> WriteOnlyManyEntriesCommand<K, V> copy(WriteOnlyManyEntriesCommand<K, V> command,
         Collection<Object> keys) {
      return new WriteOnlyManyEntriesCommand<>(command).withEntries(filterEntries(command.getEntries(), keys));
   }

   public static <K, V> WriteOnlyManyCommand<K, V> copy(WriteOnlyManyCommand<K, V> command, Collection<Object> keys) {
      WriteOnlyManyCommand<K, V> copy = new WriteOnlyManyCommand<>(command);
      //noinspection unchecked
      copy.setKeys((Collection<? extends K>) keys);
      return copy;
   }

   public static <K, V, R> ReadWriteManyCommand<K, V, R> copy(ReadWriteManyCommand<K, V, R> command,
         Collection<Object> keys) {
      ReadWriteManyCommand<K, V, R> copy = new ReadWriteManyCommand<>(command);
      //noinspection unchecked
      copy.setKeys((Collection<? extends K>) keys);
      return copy;
   }

   public static <K, V, R> ReadWriteManyEntriesCommand<K, V, R> copy(ReadWriteManyEntriesCommand<K, V, R> command,
         Collection<Object> keys) {
      return new ReadWriteManyEntriesCommand<K, V, R>(command).withEntries(filterEntries(command.getEntries(), keys));
   }

   public static Map<Object, Object> mergeHashMap(Map<Address, Response> responses, Map<Object, Object> resultMap) {
      ValidResponse response = (ValidResponse) responses.values().iterator().next();
      //noinspection unchecked
      Map<Object, Object> remoteMap = (Map<Object, Object>) response.getResponseValue();
      return InfinispanCollections.mergeMaps(resultMap, remoteMap);
   }

   @SuppressWarnings("unused")
   public static Void voidMerge(Map<Address, Response> ignored1, Void ignored2) {
      return null;
   }

   public static List<Object> mergeList(Map<Address, Response> responses, List<Object> resultList) {
      ValidResponse response = (ValidResponse) responses.values().iterator().next();
      //noinspection unchecked
      List<Object> list = (List<Object>) response.getResponseValue();
      return InfinispanCollections.mergeLists(list, resultList);
   }

   public static Map<Integer, Collection<Object>> filterBySegment(LocalizedCacheTopology cacheTopology,
         Collection<Object> keys) {
      Map<Integer, Collection<Object>> filteredKeys = new HashMap<>(
            cacheTopology.getReadConsistentHash().getNumSegments());
      for (Object key : keys) {
         filteredKeys.computeIfAbsent(cacheTopology.getSegment(key), integer -> new ArrayList<>()).add(key);
      }
      return filteredKeys;
   }


   public static <K, V> Map<K, V> filterEntries(Map<K, V> map, Collection<Object> keys) {
      //note: can't use Collector.toMap() since the implementation doesn't support null values.
      return map.entrySet().stream()
            .filter(entry -> keys.contains(entry.getKey()))
            .collect(HashMap::new, (rMap, entry) -> rMap.put(entry.getKey(), entry.getValue()), HashMap::putAll);
   }

   public static BackupWriteCommand backupFrom(CommandsFactory factory, PutKeyValueCommand command) {
      SingleKeyBackupWriteCommand cmd = factory.buildSingleKeyBackupWriteCommand();
      cmd.setPutKeyValueCommand(command);
      return cmd;
   }

   public static BackupWriteCommand backupFrom(CommandsFactory factory, RemoveCommand command) {
      SingleKeyBackupWriteCommand cmd = factory.buildSingleKeyBackupWriteCommand();
      cmd.setRemoveCommand(command, command.getCommandId() == RemoveExpiredCommand.COMMAND_ID);
      return cmd;
   }

   public static BackupWriteCommand backupFrom(CommandsFactory factory, ReplaceCommand command) {
      SingleKeyBackupWriteCommand cmd = factory.buildSingleKeyBackupWriteCommand();
      cmd.setReplaceCommand(command);
      return cmd;
   }

   public static BackupWriteCommand backupFrom(CommandsFactory factory, ReadWriteKeyValueCommand command) {
      SingleKeyFunctionalBackupWriteCommand cmd = factory.buildSingleKeyFunctionalBackupWriteCommand();
      cmd.setReadWriteKeyValueCommand(command);
      return cmd;
   }

   public static BackupWriteCommand backupFrom(CommandsFactory factory, ReadWriteKeyCommand command) {
      SingleKeyFunctionalBackupWriteCommand cmd = factory.buildSingleKeyFunctionalBackupWriteCommand();
      cmd.setReadWriteKeyCommand(command);
      return cmd;
   }

   public static BackupWriteCommand backupFrom(CommandsFactory factory, WriteOnlyKeyValueCommand command) {
      SingleKeyFunctionalBackupWriteCommand cmd = factory.buildSingleKeyFunctionalBackupWriteCommand();
      cmd.setWriteOnlyKeyValueCommand(command);
      return cmd;
   }

   public static BackupWriteCommand backupFrom(CommandsFactory factory, WriteOnlyKeyCommand command) {
      SingleKeyFunctionalBackupWriteCommand cmd = factory.buildSingleKeyFunctionalBackupWriteCommand();
      cmd.setWriteOnlyKeyCommand(command);
      return cmd;
   }

   public static BackupWriteCommand backupFrom(CommandsFactory factory, PutMapCommand command,
         Collection<Object> keys) {
      PutMapBackupWriteCommand cmd = factory.buildPutMapBackupWriteCommand();
      cmd.setPutMapCommand(command, keys);
      return cmd;
   }

   public static <K, V> BackupWriteCommand backupFrom(CommandsFactory factory,
         WriteOnlyManyEntriesCommand<K, V> command, Collection<Object> keys) {
      MultiEntriesFunctionalBackupWriteCommand cmd = factory.buildMultiEntriesFunctionalBackupWriteCommand();
      cmd.setWriteOnly(command, keys);
      return cmd;
   }

   public static <K, V, R> BackupWriteCommand backupFrom(CommandsFactory factory,
         ReadWriteManyEntriesCommand<K, V, R> command, Collection<Object> keys) {
      MultiEntriesFunctionalBackupWriteCommand cmd = factory.buildMultiEntriesFunctionalBackupWriteCommand();
      cmd.setReadWrite(command, keys);
      return cmd;
   }

   public static <K, V> BackupWriteCommand backupFrom(CommandsFactory factory,
         WriteOnlyManyCommand<K, V> command, Collection<Object> keys) {
      MultiKeyFunctionalBackupWriteCommand cmd = factory.buildMultiKeyFunctionalBackupWriteCommand();
      cmd.setWriteOnly(command, keys);
      return cmd;
   }

   public static <K, V, R> BackupWriteCommand backupFrom(CommandsFactory factory,
         ReadWriteManyCommand<K, V, R> command, Collection<Object> keys) {
      MultiKeyFunctionalBackupWriteCommand cmd = factory.buildMultiKeyFunctionalBackupWriteCommand();
      cmd.setReadWrite(command, keys);
      return cmd;
   }
}
