package com.hubspot.baragon.utils;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.concurrent.TimeUnit;

import com.google.common.base.Optional;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.net.HostAndPort;
import com.google.common.net.InetAddresses;
import com.google.inject.Singleton;

@Singleton
public class UpstreamResolver {
  private final Cache<String, String> resolveCache;

  public UpstreamResolver(long maxSize, long expireAfterDays) {
    this.resolveCache = CacheBuilder.newBuilder()
        .maximumSize(maxSize)
        .expireAfterAccess(expireAfterDays, TimeUnit.DAYS)
        .build();
  }

  public Optional<String> resolveUpstreamDNS(String address) {
    try {
      if (address.contains(":")) {
        HostAndPort hostAndPort = HostAndPort.fromString(address);

        try {
          InetAddresses.forString(hostAndPort.getHost());
          return Optional.of(address); // `address` is already an IP
        } catch (IllegalArgumentException e) {
          // `address` is not an IP, continue and try to resolve it
        }

        String cached = resolveCache.getIfPresent(hostAndPort.getHost());
        String ip;
        int port = hostAndPort.getPort();
        if (cached != null) {
          ip = cached;
        } else {
          ip = InetAddress.getByName(hostAndPort.getHost()).getHostAddress();
          resolveCache.put(hostAndPort.getHost(), ip);
        }
        return Optional.of(String.format("%s:%d", ip, port));
      } else {
        try {
          InetAddresses.forString(address);
          return Optional.of(address);
        } catch (IllegalArgumentException e) { }
        return Optional.of(InetAddress.getByName(address).getHostAddress());
      }
    } catch (UnknownHostException uhe) {
      // Don't let this exception block rendering of the template, the lb config check will still fail if the host is truly unknown
      return Optional.absent();
    }
  }
}

@@ -766,32 +766,86 @@
 ress) %7B%0A    
+S
 tr
-y %7B%0A  
+ing host;%0A    Optional%3CInteger%3E port = Optional.absent();%0A%0A
     if (addr
@@ -861,26 +861,24 @@
 ins(%22:%22)) %7B%0A
-  
       HostAn
@@ -930,35 +930,153 @@
 dress);%0A
-%0A        try %7B%0A    
+      host = hostAndPort.getHost();%0A      port = Optional.of(hostAndPort.getPort());%0A    %7D else %7B%0A      host = address;%0A    %7D%0A%0A    try %7B%0A
       In
@@ -1105,32 +1105,11 @@
 host
-AndPort.getHost());%0A    
+);%0A
     
@@ -1169,20 +1169,16 @@
 y an IP%0A
-    
     %7D ca
@@ -1208,28 +1208,24 @@
 eption e) %7B%0A
-    
       // %60ad
@@ -1276,20 +1276,16 @@
 olve it%0A
-    
     %7D%0A%0A 
@@ -1285,16 +1285,24 @@
    %7D%0A%0A  
+  try %7B%0A
       St
@@ -1349,30 +1349,11 @@
 host
-AndPort.getHost()
 );%0A
-  
     
@@ -1369,52 +1369,8 @@
 ip;%0A
-        int port = hostAndPort.getPort();%0A  
     
@@ -1397,26 +1397,24 @@
 ) %7B%0A        
-  
 ip = cached;
@@ -1416,26 +1416,24 @@
 ched;%0A      
- 
- 
 %7D else %7B%0A   
@@ -1425,26 +1425,24 @@
    %7D else %7B%0A
-  
         ip =
@@ -1468,33 +1468,16 @@
 ame(host
-AndPort.getHost()
 ).getHos
@@ -1484,26 +1484,24 @@
 tAddress();%0A
-  
         reso
@@ -1521,41 +1521,53 @@
 host
-AndPort.getHost(), ip);%0A        %7D
+, ip);%0A      %7D%0A%0A      if (port.isPresent()) %7B
 %0A   
@@ -1621,16 +1621,22 @@
 ip, port
+.get()
 ));%0A    
@@ -1658,219 +1658,49 @@
     
-try %7B%0A          InetAddresses.forString(address);%0A          return Optional.of(address);%0A        %7D catch (IllegalArgumentException e) %7B %7D%0A        return Optional.of(InetAddress.getByName(address).getHostAddress(
+return Optional.of(String.format(%22%25s%22, ip
 ));%0A

Use the cache even if there's no port on the upstream.