package io.cattle.platform.allocator.port;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import io.cattle.platform.core.addon.PortInstance;
import io.cattle.platform.core.constants.InstanceConstants;
import io.cattle.platform.core.util.PortSpec;
import io.cattle.platform.metadata.MetadataManager;

import java.util.Collection;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;

import static java.util.stream.Collectors.*;

public class PortManagerImpl implements PortManager {

    Cache<Long, Set<PortInstance>> cache = CacheBuilder.newBuilder()
            .expireAfterAccess(5, TimeUnit.MINUTES)
            .build();
    MetadataManager metadataManager;

    public PortManagerImpl(MetadataManager metadataManager) {
        this.metadataManager = metadataManager;
    }

    @Override
    public boolean portsFree(long clusterId, long hostId, Collection<PortSpec> ports) {
        return portsFree(clusterId, hostId, null, ports);
    }

    private boolean portsFree(long clusterId, long hostId, Long instanceId, Collection<PortSpec> ports) {
        Set<PortInstance> portsUsedByHost = getPorts(clusterId, hostId);
        if (portsUsedByHost == null || portsUsedByHost.size() == 0) {
            return true;
        }

        for (PortInstance portUsed : portsUsedByHost) {
            for (PortSpec requestedPort : ports) {
                if (requestedPort.getPublicPort() != null &&
                        requestedPort.getPublicPort().equals(portUsed.getPublicPort()) &&
                        publicIpTheSame(requestedPort, portUsed) &&
                        requestedPort.getProtocol().equals(portUsed.getProtocol())) {
                    if (instanceId != null && Objects.equals(instanceId, portUsed.getInstanceId())) {
                        continue;
                    }

                    return false;
                }
            }
        }

        return true;
    }


    private boolean publicIpTheSame(PortSpec requestedPort, PortInstance portUsed) {
        if (portUsed.isBindAll()) {
            return true;
        }
        return Objects.equals(requestedPort.getIpAddress(), portUsed.getBindIpAddress());
    }

    @Override
    public void assignPorts(long clusterId, long hostId, Collection<PortInstance> ports) {
        Set<PortInstance> portMap = getPorts(clusterId, hostId);
        portMap.addAll(ports);
    }

    @Override
    public boolean optionallyAssignPorts(long clusterId, long hostId, long instanceId, Collection<PortInstance> ports) {
        List<PortSpec> specs = ports.stream()
                .map(PortSpec::new)
                .collect(toList());

        if (specs.size() == 0) {
            return true;
        }

        if (!portsFree(clusterId, hostId, instanceId, specs)) {
            return false;
        }

        assignPorts(clusterId, hostId, ports);
        return true;
    }

    @Override
    public void releasePorts(long clusterId, long hostId, Collection<PortInstance> ports) {
        Set<PortInstance> portMap = getPorts(clusterId, hostId);
        portMap.removeAll(ports);
    }

    private Set<PortInstance> getPorts(long clusterId, long hostId) {
        try {
            return cache.get(hostId, () ->
                metadataManager.getMetadataForCluster(clusterId)
                        .getInstances().stream()
                        .filter(instance -> Objects.equals(hostId, instance.getHostId()) &&
                                InstanceConstants.STATE_RUNNING.equals(instance.getState()))
                        .flatMap(instanceInfo -> instanceInfo.getPorts().stream())
                        .collect(toSet()));
        } catch (ExecutionException e) {
            throw new IllegalStateException(e);
        }
    }

}

@@ -399,23 +399,19 @@
 va.util.
-Objects
+Map
 ;%0Aimport
@@ -421,19 +421,23 @@
 va.util.
-Set
+Objects
 ;%0Aimport
@@ -515,16 +515,88 @@
 imeUnit;
+%0Aimport java.util.function.Function;%0Aimport java.util.stream.Collectors;
 %0A%0Aimport
@@ -706,20 +706,28 @@
 e%3CLong, 
-Set%3C
+Map%3CString, 
 PortInst
@@ -1269,36 +1269,44 @@
 orts) %7B%0A        
-Set%3C
+Map%3CString, 
 PortInstance%3E po
@@ -1508,16 +1508,25 @@
 edByHost
+.values()
 ) %7B%0A    
@@ -2500,36 +2500,44 @@
 orts) %7B%0A        
-Set%3C
+Map%3CString, 
 PortInstance%3E po
@@ -2585,30 +2585,95 @@
     
-portMap.addAll(ports);
+for (PortInstance port : ports) %7B%0A            portMap.put(toKey(port), port);%0A        %7D
 %0A   
@@ -3289,20 +3289,28 @@
         
-Set%3C
+Map%3CString, 
 PortInst
@@ -3346,32 +3346,78 @@
 terId, hostId);%0A
+        for (PortInstance port : ports) %7B%0A    
         portMap.
@@ -3426,20 +3426,33 @@
 move
-All
+(toKey
 (port
-s
+)
 );%0A
+        %7D%0A
     
@@ -3470,12 +3470,20 @@
 ate 
-Set%3C
+Map%3CString, 
 Port
@@ -4008,117 +4008,397 @@
 ect(
-toSet()));%0A        %7D catch (ExecutionException e) %7B%0A            throw new IllegalStateException(e);%0A        %7D
+Collectors.toConcurrentMap(%0A                                this::toKey,%0A                                Function.identity())));%0A        %7D catch (ExecutionException e) %7B%0A            throw new IllegalStateException(e);%0A        %7D%0A    %7D%0A%0A    private String toKey(PortInstance port) %7B%0A        return String.format(%22%25s:%25s/%25s%22, port.getBindIpAddress(), port.getPublicPort(), port.getProtocol());
 %0A   

Make portManager use keys and not object identity (which is mutable here)