/*
 * Copyright (c) 2014-2015 Spotify AB
 *
 * 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 com.spotify.folsom.ketama;

import com.spotify.futures.CompletableFutures;
import java.util.ArrayList;
import java.util.function.Function;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import java.util.concurrent.CompletionStage;
import com.spotify.folsom.GetResult;
import com.spotify.folsom.RawMemcacheClient;
import com.spotify.folsom.client.AbstractMultiMemcacheClient;
import com.spotify.folsom.client.Request;
import com.spotify.folsom.client.MultiRequest;

import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

public class KetamaMemcacheClient extends AbstractMultiMemcacheClient {

  private static  Collection<RawMemcacheClient> clientsOnly(
    final Collection<AddressAndClient> addressAndClients) {

    int size = addressAndClients.size();
    final List<RawMemcacheClient> clients = Lists.newArrayListWithCapacity(size);
    for (final AddressAndClient client : addressAndClients) {
      clients.add(client.getClient());
    }
    return clients;
  }

  private final Continuum continuum;

  public KetamaMemcacheClient(final Collection<AddressAndClient> clients) {
    super(clientsOnly(clients));

    this.continuum = new Continuum(clients);
  }

  private RawMemcacheClient getClient(final byte[] key) {
    return continuum.findClient(key);
  }

  @Override
  public <T> CompletionStage<T> send(final Request<T> request) {
    if (request instanceof MultiRequest) {
      // T for Request should always be List<GetResult<byte[]>> here
      // which means that MultiRequest should have T = GetResult<byte[]>
      final MultiRequest<GetResult<byte[]>> multiRequest =
              (MultiRequest<GetResult<byte[]>>) request;
      if (multiRequest.getKeys().size() > 1) {
        return (CompletionStage<T>) sendSplitRequest(multiRequest);
      }
    }

    return getClient(request.getKey()).send(request);
  }

  private <T> CompletionStage<List<T>> sendSplitRequest(final MultiRequest<T> multiRequest) {
    final List<byte[]> keys = multiRequest.getKeys();

    final Map<RawMemcacheClient, List<byte[]>> routing = Maps.newIdentityHashMap();
    final List<RawMemcacheClient> routing2 = Lists.newArrayListWithCapacity(keys.size());
    for (final byte[] key : keys) {
      final RawMemcacheClient client = getClient(key);
      List<byte[]> subKeys = routing.computeIfAbsent(client, k -> Lists.newArrayList());
      subKeys.add(key);
      routing2.add(client);
    }

    final Map<RawMemcacheClient, CompletionStage<List<T>>> futures = Maps.newIdentityHashMap();

    for (final Map.Entry<RawMemcacheClient, List<byte[]>> entry : routing.entrySet()) {
      final List<byte[]> subKeys = entry.getValue();
      final Request<List<T>> subRequest = multiRequest.create(subKeys);
      final RawMemcacheClient client = entry.getKey();
      CompletionStage<List<T>> send = client.send(subRequest);
      futures.put(client, send);
    }
    final ArrayList<CompletionStage<List<T>>> list = new ArrayList<>(futures.values());
    final CompletionStage<List<List<T>>> allFutures = CompletableFutures.allAsList(list);
    return allFutures.thenApply(new Assembler<>(futures, routing2));
  }

  private static class Assembler<T, R> implements Function<List<List<T>>, List<T>> {
    private final Map<R, CompletionStage<List<T>>> futures;
    private final List<R> routing2;

    public Assembler(Map<R, CompletionStage<List<T>>> futures, List<R> routing2) {
      this.futures = futures;
      this.routing2 = routing2;
    }

    @Override
    public List<T> apply(final List<List<T>> ignored) {
      final Map<R, Iterator<T>> map = Maps.newIdentityHashMap();
      for (final Map.Entry<R, CompletionStage<List<T>>> entry : futures.entrySet()) {
        final R client = entry.getKey();
        map.put(client, CompletableFutures.getCompleted(entry.getValue()).iterator());
      }
      final List<T> result = Lists.newArrayList();
      for (final R memcacheClient : routing2) {
        final Iterator<T> iterator = map.get(memcacheClient);
        result.add(iterator.next());
      }
      return result;
    }
  }
}

@@ -693,25 +693,44 @@
 va.util.
-ArrayList
+concurrent.CompletableFuture
 ;%0Aimport
@@ -3574,25 +3574,26 @@
   final 
-ArrayList
+Collection
 %3CComplet
@@ -3615,31 +3615,17 @@
 %3E%3E%3E 
-list = new ArrayList%3C%3E(
+values = 
 futu
@@ -3640,28 +3640,28 @@
 es()
-)
 ;%0A    
-final
+return
  Complet
 ionS
@@ -3660,102 +3660,99 @@
 plet
-ionStage%3CList%3CList%3CT%3E%3E%3E allFutures = CompletableFutures.allAsList(list);%0A    return allFutures
+ableFuture%0A        .allOf(values.toArray(new CompletableFuture%3C?%3E%5Bvalues.size()%5D))%0A        
 .the
@@ -3859,29 +3859,20 @@
 unction%3C
-List%3CList%3CT%3E%3E
+Void
 , List%3CT
@@ -4174,21 +4174,12 @@
 nal 
-List%3CList%3CT%3E%3E
+Void
  ign

Use allOf instead of allAsList in ketama implementation