package com.netflix.conductor.jedis;

import com.google.common.collect.Lists;

import com.netflix.dyno.connectionpool.Host;
import com.netflix.dyno.connectionpool.HostSupplier;
import com.netflix.dyno.connectionpool.TokenMapSupplier;
import com.netflix.dyno.connectionpool.impl.lb.HostToken;

import java.util.Arrays;
import java.util.List;
import java.util.Set;

import javax.inject.Inject;
import javax.inject.Provider;

public class TokenMapSupplierProvider implements Provider<TokenMapSupplier> {
    private final HostSupplier hostSupplier;

    @Inject
    public TokenMapSupplierProvider(HostSupplier hostSupplier) {
        this.hostSupplier = hostSupplier;
    }

    @Override
    public TokenMapSupplier get() {
        return new TokenMapSupplier() {

            // FIXME This isn't particularly safe, but it is equivalent to the existing code.
            // FIXME It seems like we should be supply tokens for more than one host?
            HostToken token = new HostToken(1L, Lists.newArrayList(hostSupplier.getHosts()).get(0));

            @Override
            public List<HostToken> getTokens(Set<Host> activeHosts) {
                return Arrays.asList(token);
            }

            @Override
            public HostToken getTokenForHost(Host host, Set<Host> activeHosts) {
                return token;
            }
        };
    }
}

@@ -1,79 +1,634 @@
-package com.netflix.conductor.jedis;%0A%0Aimport com.google.common.collect.List
+/*%0A * Copyright 2016 Netflix, Inc.%0A * %3Cp%3E%0A * Licensed under the Apache License, Version 2.0 (the %22License%22); you may not use this file except in compliance with%0A * the License. You may obtain a copy of the License at%0A * %3Cp%3E%0A * http://www.apache.org/licenses/LICENSE-2.0%0A * %3Cp%3E%0A * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on%0A * an %22AS IS%22 BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the%0A * specific language governing permissions and limitations under the License.%0A */%0Apackage com.netflix.conductor.jedi
 s;%0A%0A
@@ -711,20 +711,24 @@
 ionpool.
-Host
+TokenMap
 Supplier
@@ -768,32 +768,33 @@
 ionpool.
-TokenMapSupplier
+impl.lb.HostToken
 ;%0Aimport
@@ -835,69 +835,88 @@
 mpl.
-lb.HostToken;%0A%0Aimport java.util.Arrays;%0Aimport java.util.List
+utils.CollectionUtils;%0A%0Aimport javax.inject.Inject;%0Aimport javax.inject.Provider
 ;%0Aim
@@ -934,14 +934,19 @@
 til.
-Se
+ArrayLis
 t;%0A
-%0A
 impo
@@ -956,53 +956,40 @@
 java
-x.inject.Inject;%0Aimport javax.inject.Provider
+.util.List;%0Aimport java.util.Set
 ;%0A%0Ap
@@ -1087,33 +1087,34 @@
 nal 
-HostSupplier hostSupplier
+List%3CHostToken%3E hostTokens
 ;%0A%0A 
@@ -1164,77 +1164,55 @@
 der(
-HostSupplier hostSupplier) %7B%0A        this.hostSupplier = hostSupplier
+) %7B%0A        this.hostTokens = new ArrayList%3C%3E()
 ;%0A  
@@ -1306,17 +1306,16 @@
 ier() %7B%0A
-%0A
         
@@ -1322,176 +1322,189 @@
     
-// FIXME This isn't particularly safe, but it is equivalent to the existing code.%0A            // FIXME It seems like we should be supply tokens for more than one host?%0A
+@Override%0A            public List%3CHostToken%3E getTokens(Set%3CHost%3E activeHosts) %7B%0A                long i = activeHosts.size();%0A                for (Host host : activeHosts) %7B%0A        
     
@@ -1521,17 +1521,21 @@
 stToken 
-t
+hostT
 oken = n
@@ -1551,66 +1551,65 @@
 ken(
-1L, Lists.newArrayList(hostSupplier.getHosts()).get(0)
+i, host);%0A                    hostTokens.add(hostToken
 );%0A
-%0A
     
@@ -1612,33 +1612,36 @@
             
-@Override
+        i--;
 %0A           
@@ -1645,65 +1645,13 @@
     
-public List%3CHostToken%3E getTokens(Set%3CHost%3E activeHosts) %7B
+    %7D
 %0A   
@@ -1674,28 +1674,18 @@
 urn 
-Arrays.asList(t
+hostT
 oken
-)
+s
 ;%0A  
@@ -1827,13 +1827,87 @@
 urn 
-token
+CollectionUtils.find(hostTokens, token -%3E token.getHost().compareTo(host) == 0)
 ;%0A  

fix for token supplier provider with multiple dyno hosts