package com.hubspot.baragon.client;

import static com.google.common.base.Preconditions.checkNotNull;

import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.stream.Collectors;

import javax.inject.Provider;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.fasterxml.jackson.core.type.TypeReference;
import com.google.common.base.Optional;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.hubspot.baragon.models.BaragonAgentMetadata;
import com.hubspot.baragon.models.BaragonGroup;
import com.hubspot.baragon.models.BaragonGroupAlias;
import com.hubspot.baragon.models.BaragonRequest;
import com.hubspot.baragon.models.BaragonResponse;
import com.hubspot.baragon.models.BaragonService;
import com.hubspot.baragon.models.BaragonServiceState;
import com.hubspot.baragon.models.BaragonServiceStatus;
import com.hubspot.baragon.models.QueuedRequestId;
import com.hubspot.baragon.models.UpstreamInfo;
import com.hubspot.horizon.HttpClient;
import com.hubspot.horizon.HttpRequest;
import com.hubspot.horizon.HttpRequest.Method;
import com.hubspot.horizon.HttpResponse;

public class BaragonServiceClient {

  private static final Logger LOG = LoggerFactory.getLogger(BaragonServiceClient.class);

  private static final String WORKERS_FORMAT = "%s/workers";

  private static final String LOAD_BALANCER_FORMAT = "%s/load-balancer";
  private static final String LOAD_BALANCER_BASE_PATH_FORMAT = LOAD_BALANCER_FORMAT + "/%s/base-path";
  private static final String LOAD_BALANCER_ALL_BASE_PATHS_FORMAT = LOAD_BALANCER_BASE_PATH_FORMAT + "/all";
  private static final String LOAD_BALANCER_AGENTS_FORMAT = LOAD_BALANCER_FORMAT + "/%s/agents";
  private static final String LOAD_BALANCER_KNOWN_AGENTS_FORMAT = LOAD_BALANCER_FORMAT + "/%s/known-agents";
  private static final String LOAD_BALANCER_DELETE_KNOWN_AGENT_FORMAT = LOAD_BALANCER_KNOWN_AGENTS_FORMAT + "/%s";
  private static final String LOAD_BALANCER_GROUP_FORMAT = LOAD_BALANCER_FORMAT + "/%s";
  private static final String ALL_LOAD_BALANCER_GROUPS_FORMAT = LOAD_BALANCER_FORMAT + "/all";
  private static final String LOAD_BALANCER_TRAFFIC_SOURCE_FORMAT = LOAD_BALANCER_GROUP_FORMAT + "/sources";

  private static final String REQUEST_FORMAT = "%s/request";
  private static final String REQUEST_ID_FORMAT = REQUEST_FORMAT + "/%s";

  private static final String UPSTREAM_REQUEST_FORMAT = REQUEST_FORMAT + "/upstreams/%s";

  private static final String STATE_FORMAT = "%s/state";
  private static final String STATE_SERVICE_ID_FORMAT = STATE_FORMAT + "/%s";
  private static final String STATE_RELOAD_FORMAT = STATE_SERVICE_ID_FORMAT + "/reload";

  private static final String STATUS_FORMAT = "%s/status";

  private static final String ALIASES_FORMAT = "%s/aliases/%s";

  private static final TypeReference<Collection<String>> STRING_COLLECTION = new TypeReference<Collection<String>>() {};
  private static final TypeReference<Collection<BaragonGroup>> BARAGON_GROUP_COLLECTION = new TypeReference<Collection<BaragonGroup>>() {};
  private static final TypeReference<Collection<BaragonAgentMetadata>> BARAGON_AGENTS_COLLECTION = new TypeReference<Collection<BaragonAgentMetadata>>() {};
  private static final TypeReference<Collection<QueuedRequestId>> QUEUED_REQUEST_COLLECTION = new TypeReference<Collection<QueuedRequestId>>() {};
  private static final TypeReference<Collection<BaragonServiceState>> BARAGON_SERVICE_STATE_COLLECTION = new TypeReference<Collection<BaragonServiceState>>() {};

  private final Random random;
  private final Provider<List<String>> baseUrlProvider;
  private final Provider<Optional<String>> authkeyProvider;

  private final HttpClient httpClient;

  public BaragonServiceClient(String contextPath, HttpClient httpClient, List<String> hosts, Optional<String> authkey) {
    this(
        httpClient,
        ProviderUtils.of(ImmutableList.copyOf(hosts.stream().map((h) -> String.format("%s/%s", h, contextPath)).collect(Collectors.toList()))),
        ProviderUtils.of(authkey)
    );
  }

  public BaragonServiceClient(HttpClient httpClient, Provider<List<String>> baseUrlProvider, Provider<Optional<String>> authkeyProvider) {
    this.httpClient = httpClient;
    this.baseUrlProvider = baseUrlProvider;
    this.authkeyProvider = authkeyProvider;
    this.random = new Random();
  }

  private String getBaseUrl() {
    final List<String> baseUrls = baseUrlProvider.get();
    String chosenBaseUrl = baseUrls.get(random.nextInt(baseUrls.size()));
    if (chosenBaseUrl.endsWith("/")) {
      return chosenBaseUrl.substring(0, chosenBaseUrl.length() - 1);
    } else {
      return chosenBaseUrl;
    }
  }

  private HttpRequest.Builder buildRequest(String uri) {
    return buildRequest(uri, null);
  }

  private HttpRequest.Builder buildRequest(String uri, Map<String, String> queryParams) {
    final HttpRequest.Builder builder = HttpRequest.newBuilder().setUrl(uri);

    final Optional<String> maybeAuthkey = authkeyProvider.get();

    if (maybeAuthkey.isPresent()) {
      builder.setQueryParam("authkey").to(maybeAuthkey.get());
    }

    if ((queryParams != null) && (!queryParams.isEmpty())) {
      for (Map.Entry<String, String> entry : queryParams.entrySet()) {
        builder.setQueryParam(entry.getKey()).to(entry.getValue());
      }
    }

    return builder;
  }

  private void checkResponse(String type, HttpResponse response) {
    if (response.isError()) {
      throw fail(type, response);
    }
  }

  private BaragonClientException fail(String type, HttpResponse response) {
    String body = "";
    try {
      body = response.getAsString();
    } catch (Exception e) {
      LOG.warn("Unable to read body", e);
    }

    String uri = "";
    try {
      uri = response.getRequest().getUrl().toString();
    } catch (Exception e) {
      LOG.warn("Unable to read uri", e);
    }

    throw new BaragonClientException(String.format("Failed '%s' action on Baragon (%s) - code: %s, %s", type, uri, response.getStatusCode(), body), response.getStatusCode());
  }

  private <T> Optional<T> getSingle(String uri, String type, String id, Class<T> clazz) {
    return getSingle(uri, type, id, clazz, null);
  }

  private <T> Optional<T> getSingle(String uri, String type, String id, Class<T> clazz, Map<String, String> queryParams) {
    checkNotNull(id, String.format("Provide a %s id", type));
    LOG.debug("Getting {} {} from {}", type, id, uri);
    final long start = System.currentTimeMillis();

    HttpResponse response = httpClient.execute(buildRequest(uri, queryParams).build());

    if (response.getStatusCode() == 404) {
      return Optional.absent();
    }

    checkResponse(type, response);
    LOG.debug("Got {} {} in {}ms", type, id, System.currentTimeMillis() - start);
    return Optional.fromNullable(response.getAs(clazz));
  }

  private <T> Collection<T> getCollection(String uri, String type, TypeReference<Collection<T>> typeReference) {
    LOG.debug("Getting all {} from {}", type, uri);
    final long start = System.currentTimeMillis();

    HttpResponse response = httpClient.execute(buildRequest(uri).build());

    if (response.getStatusCode() == 404) {
      throw new BaragonClientException(String.format("%s not found", type), 404);
    }

    checkResponse(type, response);
    LOG.debug("Got {} in {}ms", type, System.currentTimeMillis() - start);
    return response.getAs(typeReference);
  }

  private <T> void delete(String uri, String type, String id, Map<String, String> queryParams) {
    delete(uri, type, id, queryParams, Optional.<Class<T>>absent());
  }

  private <T> Optional<T> delete(String uri, String type, String id, Map<String, String> queryParams, Optional<Class<T>> clazz) {
    LOG.debug("Deleting {} {} from {}", type, id, uri);
    final long start = System.currentTimeMillis();
    HttpRequest.Builder request = buildRequest(uri, queryParams).setMethod(Method.DELETE);
    HttpResponse response = httpClient.execute(request.build());

    if (response.getStatusCode() == 404) {
      LOG.debug("{} ({}) was not found", type, id);
      return Optional.absent();
    }

    checkResponse(type, response);
    LOG.debug("Deleted {} ({}) from Baragon in %sms", type, id, System.currentTimeMillis() - start);

    if (clazz.isPresent()) {
      return Optional.of(response.getAs(clazz.get()));
    }

    return Optional.absent();
  }

  private <T> Optional<T> post(String uri, String type, Optional<?> body, Optional<Class<T>> clazz) {
    return post(uri, type, body, clazz, Collections.<String, String>emptyMap());
  }

  private <T> Optional<T> post(String uri, String type, Optional<?> body, Optional<Class<T>> clazz, Map<String, String> queryParams) {
    try {
      HttpResponse response = post(uri, type, body, queryParams);

      if (clazz.isPresent()) {
        return Optional.of(response.getAs(clazz.get()));
      }
    } catch (Exception e) {
      LOG.warn("Http post failed", e);
    }

    return Optional.absent();
  }

  private HttpResponse post(String uri, String type, Optional<?> body, Map<String, String> params) {
    LOG.debug("Posting {} to {}", type, uri);
    final long start = System.currentTimeMillis();
    HttpRequest.Builder request = buildRequest(uri, params).setMethod(Method.POST);

    if (body.isPresent()) {
      request.setBody(body.get());
    }

    HttpResponse response = httpClient.execute(request.build());
    checkResponse(type, response);
    LOG.debug("Successfully posted {} in {}ms", type, System.currentTimeMillis() - start);
    return response;
  }

  private <T> Optional<T> put(String uri, String type, Optional<Class<T>> clazz, Map<String, String> queryParams) {
    try {
      HttpResponse response = put(uri, type, queryParams);

      if (clazz.isPresent()) {
        return Optional.of(response.getAs(clazz.get()));
      }
    } catch (Exception e) {
      LOG.warn("Http post failed", e);
    }

    return Optional.absent();
  }

  private HttpResponse put(String uri, String type, Map<String, String> params) {
    LOG.debug("Posting {} to {}", type, uri);
    final long start = System.currentTimeMillis();
    HttpRequest.Builder request = buildRequest(uri, params).setMethod(Method.PUT);
    HttpResponse response = httpClient.execute(request.build());
    checkResponse(type, response);
    LOG.debug("Successfully posted {} in {}ms", type, System.currentTimeMillis() - start);
    return response;
  }

  // BaragonService overall status

  public Optional<BaragonServiceStatus> getBaragonServiceStatus(String baseUrl) {
    final String uri = String.format(STATUS_FORMAT, baseUrl);
    return getSingle(uri, "status", "", BaragonServiceStatus.class);
  }

  public Optional<BaragonServiceStatus> getAnyBaragonServiceStatus() {
    return getBaragonServiceStatus(getBaseUrl());
  }

  // BaragonService service states

  public Collection<BaragonServiceState> getGlobalState() {
    final String uri = String.format(STATE_FORMAT, getBaseUrl());
    return getCollection(uri, "global state", BARAGON_SERVICE_STATE_COLLECTION);
  }

  public Optional<BaragonServiceState> getServiceState(String serviceId) {
    final String uri = String.format(STATE_SERVICE_ID_FORMAT, getBaseUrl(), serviceId);
    return getSingle(uri, "service state", serviceId, BaragonServiceState.class);
  }

  public Optional<BaragonResponse> deleteService(String serviceId) {
    final String uri = String.format(STATE_SERVICE_ID_FORMAT, getBaseUrl(), serviceId);
    return delete(uri, "service state", serviceId, Collections.emptyMap(), Optional.of(BaragonResponse.class));
  }

  public Optional<BaragonResponse> reloadServiceConfigs(String serviceId){
    final String uri = String.format(STATE_RELOAD_FORMAT, getBaseUrl(), serviceId);
    return post(uri, "service reload",Optional.absent(), Optional.of(BaragonResponse.class));
  }

  // BaragonService Workers

  public Collection<String> getBaragonServiceWorkers() {
    final String requestUri = String.format(WORKERS_FORMAT, getBaseUrl());
    return getCollection(requestUri, "baragon service workers", STRING_COLLECTION);
  }

  // BaragonService load balancer group actions

  public Collection<String> getLoadBalancerGroups() {
    final String requestUri = String.format(LOAD_BALANCER_FORMAT, getBaseUrl());
    return getCollection(requestUri, "load balancer groups", STRING_COLLECTION);
  }

  public Collection<BaragonGroup> getAllLoadBalancerGroups() {
    final String requestUri = String.format(ALL_LOAD_BALANCER_GROUPS_FORMAT, getBaseUrl());
    return getCollection(requestUri, "load balancer groups", BARAGON_GROUP_COLLECTION);
  }

  public Collection<BaragonAgentMetadata> getLoadBalancerGroupAgentMetadata(String loadBalancerGroupName) {
    final String requestUri = String.format(LOAD_BALANCER_AGENTS_FORMAT, getBaseUrl(), loadBalancerGroupName);
    return getCollection(requestUri, "load balancer agent metadata", BARAGON_AGENTS_COLLECTION);
  }

  public Collection<BaragonAgentMetadata> getLoadBalancerGroupKnownAgentMetadata(String loadBalancerGroupName) {
    final String requestUri = String.format(LOAD_BALANCER_KNOWN_AGENTS_FORMAT, getBaseUrl(), loadBalancerGroupName);
    return getCollection(requestUri, "load balancer known agent metadata", BARAGON_AGENTS_COLLECTION);
  }

  public void deleteLoadBalancerGroupKnownAgent(String loadBalancerGroupName, String agentId) {
    final String requestUri = String.format(LOAD_BALANCER_DELETE_KNOWN_AGENT_FORMAT, getBaseUrl(), loadBalancerGroupName, agentId);
    delete(requestUri, "known agent", agentId, Collections.<String, String>emptyMap());
  }

  public BaragonGroup addTrafficSource(String loadBalancerGroupName, String source) {
    final String requestUri = String.format(LOAD_BALANCER_TRAFFIC_SOURCE_FORMAT, getBaseUrl(), loadBalancerGroupName);
    return post(requestUri, "add source", Optional.absent(), Optional.of(BaragonGroup.class), ImmutableMap.of("source", source)).get();
  }

  public Optional<BaragonGroup> removeTrafficSource(String loadBalancerGroupName, String source) {
    final String requestUri = String.format(LOAD_BALANCER_TRAFFIC_SOURCE_FORMAT, getBaseUrl(), loadBalancerGroupName);
    return delete(requestUri, "remove source", source, ImmutableMap.of("source", source), Optional.of(BaragonGroup.class));
  }

  public Optional<BaragonGroup> getGroupDetail(String loadBalancerGroupName) {
    final String requestUri = String.format(LOAD_BALANCER_GROUP_FORMAT, getBaseUrl(), loadBalancerGroupName);
    return getSingle(requestUri, "group detail", loadBalancerGroupName, BaragonGroup.class);
  }

  // BaragonService base path actions

  public Collection<String> getOccupiedBasePaths(String loadBalancerGroupName) {
    final String requestUri = String.format(LOAD_BALANCER_ALL_BASE_PATHS_FORMAT, getBaseUrl(), loadBalancerGroupName);
    return getCollection(requestUri, "occupied base paths", STRING_COLLECTION);
  }

  public Optional<BaragonService> getServiceForBasePath(String loadBalancerGroupName, String basePath) {
    final String requestUri = String.format(LOAD_BALANCER_BASE_PATH_FORMAT, getBaseUrl(), loadBalancerGroupName);
    return getSingle(requestUri, "service for base path", "", BaragonService.class, ImmutableMap.of("basePath", basePath));
  }

  public void clearBasePath(String loadBalancerGroupName, String basePath) {
    final String requestUri = String.format(LOAD_BALANCER_BASE_PATH_FORMAT, getBaseUrl(), loadBalancerGroupName);
    delete(requestUri, "base path", "", ImmutableMap.of("basePath", basePath));
  }

  // BaragonService request actions

  public Optional<BaragonResponse> getRequest(String requestId) {
    final String uri = String.format(REQUEST_ID_FORMAT, getBaseUrl(), requestId);
    return getSingle(uri, "request", requestId, BaragonResponse.class);
  }

  public Optional<BaragonResponse> enqueueRequest(BaragonRequest request) {
    final String uri = String.format(REQUEST_FORMAT, getBaseUrl());
    return post(uri, "request", Optional.of(request), Optional.of(BaragonResponse.class));
  }

  public Optional<BaragonResponse> cancelRequest(String requestId) {
    final String uri = String.format(REQUEST_ID_FORMAT, getBaseUrl(), requestId);
    return delete(uri, "request", requestId, Collections.<String, String>emptyMap(), Optional.of(BaragonResponse.class));
  }

  public Optional<BaragonResponse> addUpstream(String serviceId, UpstreamInfo upstreamInfo) {
    final String uri = String.format(UPSTREAM_REQUEST_FORMAT, getBaseUrl(), serviceId);
    return put(uri, "upstream-add", Optional.of(BaragonResponse.class), ImmutableMap.of("upstream", upstreamInfo.toPath()));
  }

  public Optional<BaragonResponse> setUpstreams(String serviceId, List<UpstreamInfo> upstreamInfo) {
    final String uri = String.format(UPSTREAM_REQUEST_FORMAT, getBaseUrl(), serviceId);
    return post(uri, "upstream-add", Optional.of(upstreamInfo), Optional.of(BaragonResponse.class));
  }

  public Optional<BaragonResponse> removeUpstream(String serviceId, UpstreamInfo upstreamInfo) {
    final String uri = String.format(UPSTREAM_REQUEST_FORMAT, getBaseUrl(), serviceId);
    return delete(uri, "upstream-remove", serviceId, ImmutableMap.of("upstream", upstreamInfo.toPath()), Optional.of(BaragonResponse.class));
  }

  // BaragonService queued request actions

  public Collection<QueuedRequestId> getQueuedRequests() {
    final String uri = String.format(REQUEST_FORMAT, getBaseUrl());
    return getCollection(uri, "queued requests", QUEUED_REQUEST_COLLECTION);
  }

  // Aliases

  public Optional<BaragonGroupAlias> getAlias(String name) {
    final String uri = String.format(ALIASES_FORMAT, getBaseUrl(), name);
    return getSingle(uri, "alias", name, BaragonGroupAlias.class);
  }

  public void createAlias(String name, BaragonGroupAlias alias) {
    final String uri = String.format(ALIASES_FORMAT, getBaseUrl(), name);
    post(uri, "alias", Optional.of(alias), Optional.absent());
  }

  public void deleteAlias(String name) {
    final String uri = String.format(ALIASES_FORMAT, getBaseUrl(), name);
    delete(uri, "alias", name, Collections.emptyMap());
  }
}


package com.hubspot.baragon.service.elb;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.amazonaws.AmazonClientException;
import com.amazonaws.AmazonServiceException;
import com.amazonaws.services.elasticloadbalancing.model.Instance;
import com.amazonaws.services.elasticloadbalancingv2.AmazonElasticLoadBalancingClient;
import com.amazonaws.services.elasticloadbalancingv2.model.AvailabilityZone;
import com.amazonaws.services.elasticloadbalancingv2.model.CreateListenerRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.CreateRuleRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.CreateTargetGroupRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.DeleteListenerRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.DeleteRuleRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.DeregisterTargetsRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeListenersRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeListenersResult;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeLoadBalancersRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeLoadBalancersResult;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeRulesRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeTargetGroupAttributesRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeTargetGroupAttributesResult;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeTargetGroupsRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeTargetGroupsResult;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeTargetHealthRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.DescribeTargetHealthResult;
import com.amazonaws.services.elasticloadbalancingv2.model.Listener;
import com.amazonaws.services.elasticloadbalancingv2.model.LoadBalancer;
import com.amazonaws.services.elasticloadbalancingv2.model.LoadBalancerNotFoundException;
import com.amazonaws.services.elasticloadbalancingv2.model.ModifyListenerRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.ModifyRuleRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.ModifyTargetGroupRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.RegisterTargetsRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.Rule;
import com.amazonaws.services.elasticloadbalancingv2.model.SetSubnetsRequest;
import com.amazonaws.services.elasticloadbalancingv2.model.TargetDescription;
import com.amazonaws.services.elasticloadbalancingv2.model.TargetGroup;
import com.amazonaws.services.elasticloadbalancingv2.model.TargetGroupAttribute;
import com.amazonaws.services.elasticloadbalancingv2.model.TargetGroupNotFoundException;
import com.amazonaws.services.elasticloadbalancingv2.model.TargetHealthDescription;
import com.amazonaws.services.elasticloadbalancingv2.model.TargetHealthStateEnum;
import com.google.common.base.Optional;
import com.google.common.base.Strings;
import com.google.common.base.Throwables;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.SetMultimap;
import com.google.common.collect.Sets;
import com.google.inject.Inject;
import com.google.inject.name.Named;
import com.hubspot.baragon.data.BaragonKnownAgentsDatastore;
import com.hubspot.baragon.data.BaragonLoadBalancerDatastore;
import com.hubspot.baragon.models.AgentCheckInResponse;
import com.hubspot.baragon.models.BaragonAgentMetadata;
import com.hubspot.baragon.models.BaragonGroup;
import com.hubspot.baragon.models.RegisterBy;
import com.hubspot.baragon.models.TrafficSource;
import com.hubspot.baragon.models.TrafficSourceState;
import com.hubspot.baragon.models.TrafficSourceType;
import com.hubspot.baragon.service.BaragonServiceModule;
import com.hubspot.baragon.service.config.ElbConfiguration;
import com.hubspot.baragon.service.exceptions.BaragonExceptionNotifier;

/**
 * Handles interactions with the ApplicationLoadBalancer in AWS.
 *
 * For this class I am working under the assumption that the ID's of Targets can
 * be interchanged with the ID's of instances. That is, if you ask this class
 * for the health of instance 12345, then you will actually expect the health
 * of target 12345.
 */
public class ApplicationLoadBalancer extends ElasticLoadBalancer {
  private static final int MAX_TARGET_GROUP_PAGE_SIZE = 400;

  private static final Logger LOG = LoggerFactory.getLogger(ApplicationLoadBalancer.class);
  private static final String DEREGISTRATION_DELAY_ATTR = "deregistration_delay.timeout_seconds";
  private final AmazonElasticLoadBalancingClient elbClient;

  @Inject
  public ApplicationLoadBalancer(Optional<ElbConfiguration> configuration,
                                 BaragonExceptionNotifier exceptionNotifier,
                                 BaragonLoadBalancerDatastore loadBalancerDatastore,
                                 BaragonKnownAgentsDatastore knownAgentsDatastore,
                                 @Named(BaragonServiceModule.BARAGON_AWS_ELB_CLIENT_V2) AmazonElasticLoadBalancingClient elbClient) {
    super(configuration, exceptionNotifier, loadBalancerDatastore, knownAgentsDatastore);
    this.elbClient = elbClient;
  }

  @Override
  public boolean isInstanceHealthy(String instanceId, String elbName) {
    DescribeTargetHealthRequest healthRequest = new DescribeTargetHealthRequest()
        .withTargets(new TargetDescription().withId(instanceId));
    DescribeTargetHealthResult result = elbClient.describeTargetHealth(healthRequest);

    for (TargetHealthDescription health: result.getTargetHealthDescriptions()) {
      if (health.getTargetHealth().getState().equals(TargetHealthStateEnum.Healthy.toString())
          && health.getTarget().getId().equals(instanceId)) {
        return true;
      }
    }

    return false;
  }

  @Override
  public AgentCheckInResponse removeInstance(Instance instance, String id, String trafficSourceName, String agentId) {
    return removeInstance(id, trafficSourceName);
  }

  public AgentCheckInResponse removeInstance(String id, String trafficSourceName) {

    Optional<TargetGroup> maybeTargetGroup = getTargetGroup(trafficSourceName);
    if (maybeTargetGroup.isPresent()) {
      TargetGroup targetGroup = maybeTargetGroup.get();
      TargetDescription targetDescription = new TargetDescription()
          .withId(id);
      if (targetsOn(targetGroup).contains(targetDescription.getId())) {
        DeregisterTargetsRequest deregisterTargetsRequest = new DeregisterTargetsRequest()
            .withTargets(targetDescription)
            .withTargetGroupArn(targetGroup.getTargetGroupArn());
        try {
          elbClient.deregisterTargets(deregisterTargetsRequest);
          LOG.info("De-registered instance {} from target group {}", id, targetGroup);
          return getShutdownResponse(targetGroup.getTargetGroupArn(), targetDescription, true);
        } catch (AmazonServiceException exn) {
          LOG.warn("Failed to de-register instance {} from target group {}", id, targetGroup);
        }
      } else {
        LOG.debug("Instance {} not found at target group {}", id, targetGroup);
      }
    } else {
      LOG.debug("No target group found with name {}", trafficSourceName);
    }

    return new AgentCheckInResponse(TrafficSourceState.DONE, Optional.absent(), 0L);
  }

  private AgentCheckInResponse getShutdownResponse(String targetGroupArn, TargetDescription targetDescription, boolean withDrainTime) {
    Optional<String> maybeException = Optional.absent();
    Optional<Long> maybeDrainTime = Optional.absent();

    TrafficSourceState state = TrafficSourceState.PENDING;
    if (withDrainTime) {
      try {
        DescribeTargetGroupAttributesResult result = elbClient.describeTargetGroupAttributes(new DescribeTargetGroupAttributesRequest().withTargetGroupArn(targetGroupArn));
        LOG.debug("Got target group attributes {}", result.getAttributes());
        for (TargetGroupAttribute attribute : result.getAttributes()) {
          if (attribute.getKey().equals(DEREGISTRATION_DELAY_ATTR)) {
            LOG.info("Target group {} has connection drain time of {}s", targetGroupArn, attribute.getValue());
            maybeDrainTime = Optional.of(TimeUnit.SECONDS.toMillis(Long.parseLong(attribute.getValue())));
          }
        }
      } catch (Exception e) {
        LOG.warn("Error fetching connection drain time, will use default", e);
        maybeException = Optional.of(e.getMessage());
      }
    }

    try {
      DescribeTargetHealthResult healthResult = elbClient.describeTargetHealth(
          new DescribeTargetHealthRequest().withTargetGroupArn(targetGroupArn).withTargets(targetDescription));
      if (!healthResult.getTargetHealthDescriptions().isEmpty()) {
        switch (healthResult.getTargetHealthDescriptions().get(0).getTargetHealth().getState()) {
          case "initial":
          case "healthy":
          case "unhealthy":
          case "draining":
            break;
          case "unused":
          default:
            state = TrafficSourceState.DONE;
        }
      } else {
        state = TrafficSourceState.DONE;
      }
    } catch (Exception e) {
      LOG.error("Error fetching target health", e);
    }
    return new AgentCheckInResponse(state, maybeException, maybeDrainTime.or(configuration.get().getDefaultCheckInWaitTimeMs()));
  }

  @Override
  public AgentCheckInResponse checkRemovedInstance(String id, String trafficSourceName, String agentId) {
    Optional<TargetGroup> maybeTargetGroup = getTargetGroup(trafficSourceName);
    if (maybeTargetGroup.isPresent()) {
      return getShutdownResponse(maybeTargetGroup.get().getTargetGroupArn(), new TargetDescription().withId(id), false);
    } else {
      String message = String.format("Could not find target group %s", trafficSourceName);
      LOG.error(message);
      return new AgentCheckInResponse(TrafficSourceState.ERROR, Optional.of(message), 0L);
    }
  }

  public Collection<TargetHealthDescription> getTargetsOn(TargetGroup targetGroup) {
    DescribeTargetHealthRequest targetHealthRequest = new DescribeTargetHealthRequest()
        .withTargetGroupArn(targetGroup.getTargetGroupArn());

    return elbClient
        .describeTargetHealth(targetHealthRequest)
        .getTargetHealthDescriptions();
  }

  @Override
  public AgentCheckInResponse registerInstance(Instance instance, String id, String trafficSourceName, BaragonAgentMetadata agent) {
    Optional<TargetGroup> maybeTargetGroup = getTargetGroup(trafficSourceName);

    if (!maybeTargetGroup.isPresent()) {
      String message = String.format("Target group with name %s not found", trafficSourceName);
      LOG.debug(message);
      return new AgentCheckInResponse(TrafficSourceState.ERROR, Optional.of(message), 0L);
    } else {
      TargetGroup targetGroup = maybeTargetGroup.get();

      if (!isVpcOkay(agent, targetGroup)) {
        String message = String.format("VPC not configured to accept agent %s", agent);
        LOG.debug(message);
        return new AgentCheckInResponse(TrafficSourceState.ERROR, Optional.of(message), 0L);
      } else {
        return registerInstance(id, targetGroup);
      }
    }
  }

  public AgentCheckInResponse registerInstance(String id, String targetGroup) {
    Optional<TargetGroup> maybeTargetGroup = getTargetGroup(targetGroup);
    if (maybeTargetGroup.isPresent()) {
      return registerInstance(id, maybeTargetGroup.get());
    } else {
      return new AgentCheckInResponse(TrafficSourceState.ERROR, Optional.of("Target group not found"), 0L);
    }
  }

  private AgentCheckInResponse registerInstance(String id, TargetGroup targetGroup) {
    TargetDescription newTarget = new TargetDescription()
        .withId(id);

    if (!targetsOn(targetGroup).contains(newTarget.getId())) {
      try {
        RegisterTargetsRequest registerTargetsRequest = new RegisterTargetsRequest()
            .withTargets(newTarget)
            .withTargetGroupArn(targetGroup.getTargetGroupArn());
        elbClient.registerTargets(registerTargetsRequest);
        LOG.info("Registered instance {} with target group {}", newTarget.getId(), targetGroup);
      } catch (AmazonServiceException exn) {
        LOG.warn("Failed to register instance {} with target group {}", id, targetGroup);
        throw Throwables.propagate(exn);
      }

      return instanceHealthResponse(newTarget, targetGroup, id);
    } else {
      LOG.debug("Instance {} already registered with target group {}", id, targetGroup);
      return new AgentCheckInResponse(TrafficSourceState.DONE, Optional.absent(), 0L);
    }
  }

  private AgentCheckInResponse instanceHealthResponse(TargetDescription targetDescription, TargetGroup targetGroup, String id) {
    TrafficSourceState state = TrafficSourceState.DONE;
    Optional<String> exception = Optional.absent();
    try {
      DescribeTargetHealthResult healthResult = elbClient.describeTargetHealth(
          new DescribeTargetHealthRequest().withTargetGroupArn(targetGroup.getTargetGroupArn()).withTargets(targetDescription));
      if (!healthResult.getTargetHealthDescriptions().isEmpty()) {
        String targetState = healthResult.getTargetHealthDescriptions().get(0).getTargetHealth().getState();
        switch (targetState) {
          case "initial":
            state = TrafficSourceState.PENDING;
            break;
          case "healthy":
            state = TrafficSourceState.DONE;
            break;
          case "unhealthy":
          case "draining":
          case "unused":
          default:
            String message = String.format("Expected agent to be added but was in state %s", targetState);
            exception = Optional.of(message);
            LOG.error(message);
            state = TrafficSourceState.ERROR;
        }
      } else {
        String message = String.format("Instance %s not found in target group", id);
        LOG.error(message);
        exception = Optional.of(message);
        state = TrafficSourceState.ERROR;
      }
    } catch (Exception e) {
      LOG.error("Error fetching target health", e);
    }
    return new AgentCheckInResponse(state, exception, state == TrafficSourceState.PENDING ? configuration.get().getDefaultCheckInWaitTimeMs() : 0L);
  }

  @Override
  public AgentCheckInResponse checkRegisteredInstance(Instance instance, String id, TrafficSource trafficSource, BaragonAgentMetadata agent) {
    Optional<TargetGroup> maybeTargetGroup = getTargetGroup(trafficSource.getName());
    if (maybeTargetGroup.isPresent()) {
      return instanceHealthResponse(
          new TargetDescription().withId(trafficSource.getRegisterBy() == RegisterBy.INSTANCE_ID ? instance.getInstanceId() : agent.getEc2().getPrivateIp().get()),
          maybeTargetGroup.get(),
          id);
    } else {
      String message = String.format("Could not find target group %s", trafficSource.getName());
      LOG.error(message);
      return new AgentCheckInResponse(TrafficSourceState.ERROR, Optional.of(message), 0L);
    }
  }

  @Override
  public void syncAll(Collection<BaragonGroup> baragonGroups) {
    Collection<LoadBalancer> allLoadBalancers = getAllLoadBalancers();
    Map<BaragonGroup, Map<TrafficSource, TargetGroup>> baragonGroupToTrafficSource = getTrafficSources(baragonGroups);
    for (Map.Entry<BaragonGroup, Map<TrafficSource, TargetGroup>> entry : baragonGroupToTrafficSource.entrySet()) {
      BaragonGroup baragonGroup = entry.getKey();
      Collection<LoadBalancer> elbsForBaragonGroup = getLoadBalancersByBaragonGroup(allLoadBalancers, baragonGroup);
      Collection<BaragonAgentMetadata> baragonAgents = getAgentsByBaragonGroup(baragonGroup);

      Map<TrafficSource, TargetGroup> targetGroups = entry.getValue();
      for (Entry<TrafficSource, TargetGroup> targetGroupEntry : targetGroups.entrySet()) {
        try {
          TargetGroup targetGroup = targetGroupEntry.getValue();
          TrafficSource trafficSource = targetGroupEntry.getKey();
          Collection<TargetDescription> targets = targetsInTargetGroup(targetGroup);

          LOG.debug("Registering new instances for target group {}", trafficSource.getName());
          guaranteeRegistered(trafficSource, targetGroup, targets, baragonAgents, elbsForBaragonGroup);

          if (configuration.isPresent() && configuration.get().isDeregisterEnabled()) {
            LOG.debug("De-registering old instances for target group {}", trafficSource.getName());
            deregisterRemovableTargets(trafficSource, baragonGroup, targetGroup, baragonAgents, targets);
          }
        } catch (AmazonClientException exn) {
          LOG.error("Could not retrieve elb information due to ELB client error", exn);
          exceptionNotifier.notify(exn, ImmutableMap.of("baragonGroup", baragonGroup.toString()));
        } catch (Exception exn) {
          LOG.error("Could not process ELB sync", exn);
          exceptionNotifier.notify(exn, ImmutableMap.of("groups", baragonGroup.toString()));
        }
      }
    }
  }

  private Map<BaragonGroup, Map<TrafficSource, TargetGroup>> getTrafficSources(Collection<BaragonGroup> baragonGroups) {
    SetMultimap<BaragonGroup, TrafficSource> trafficSources = HashMultimap.create();
    for (BaragonGroup baragonGroup : baragonGroups) {
      Set<TrafficSource> albSources = baragonGroup
          .getTrafficSources()
          .stream()
          .filter(trafficSource -> trafficSource.getType() == TrafficSourceType.ALB_TARGET_GROUP)
          .collect(Collectors.toSet());
      trafficSources.putAll(baragonGroup, albSources);
    }

    Set<TargetGroup> targetGroups = getAllTargetGroups();
    Map<String, TargetGroup> namedTargetGroups = targetGroups.stream()
        .collect(Collectors.toMap(TargetGroup::getTargetGroupName, Function.identity()));

    Map<BaragonGroup, Map<TrafficSource, TargetGroup>> result = new HashMap<>();

    for (BaragonGroup baragonGroup : trafficSources.keySet()) {
      HashMap<TrafficSource, TargetGroup> targetGroupsInBaragonGroup = new HashMap<>();
      for (TrafficSource trafficSource: trafficSources.get(baragonGroup)) {
        targetGroupsInBaragonGroup.put(trafficSource, namedTargetGroups.get(trafficSource.getName()));
      }
      result.put(baragonGroup, targetGroupsInBaragonGroup);
    }

    return result;
  }

  public Collection<LoadBalancer> getAllLoadBalancers() {
    Collection<LoadBalancer> loadBalancers = new HashSet<>();
    DescribeLoadBalancersRequest loadBalancersRequest = new DescribeLoadBalancersRequest();
    DescribeLoadBalancersResult result = elbClient.describeLoadBalancers(loadBalancersRequest);
    String nextPage = result.getNextMarker();
    loadBalancers.addAll(result.getLoadBalancers());

    while (!Strings.isNullOrEmpty(nextPage)) {
      loadBalancersRequest = new DescribeLoadBalancersRequest()
          .withMarker(nextPage);
      result = elbClient.describeLoadBalancers(loadBalancersRequest);
      nextPage = result.getNextMarker();
      loadBalancers.addAll(result.getLoadBalancers());
    }

    return loadBalancers;
  }

  public Optional<LoadBalancer> getLoadBalancer(String loadBalancer) {
    DescribeLoadBalancersRequest request = new DescribeLoadBalancersRequest()
        .withNames(loadBalancer);

    try {
      List<LoadBalancer> maybeLoadBalancer = elbClient
          .describeLoadBalancers(request)
          .getLoadBalancers();

      if (maybeLoadBalancer.size() > 0) {
        return Optional.of(maybeLoadBalancer.get(0));
      } else {
        return Optional.absent();
      }
    } catch (LoadBalancerNotFoundException notFound) {
      LOG.warn("Could not find load balancer with name {}", loadBalancer);
      return Optional.absent();
    }
  }

  public Set<TargetGroup> getAllTargetGroups() {
    Set<TargetGroup> targetGroups = new HashSet<>();
    DescribeTargetGroupsRequest request = new DescribeTargetGroupsRequest()
        .withPageSize(MAX_TARGET_GROUP_PAGE_SIZE);
    DescribeTargetGroupsResult result = elbClient.describeTargetGroups(request);
    String nextMarker = result.getNextMarker();
    targetGroups.addAll(result.getTargetGroups());

    while (!Strings.isNullOrEmpty(nextMarker)) {
      DescribeTargetGroupsRequest nextRequest = new DescribeTargetGroupsRequest()
          .withMarker(nextMarker)
          .withPageSize(MAX_TARGET_GROUP_PAGE_SIZE);
      DescribeTargetGroupsResult nextResult = elbClient.describeTargetGroups(nextRequest);
      nextMarker = nextResult.getNextMarker();
      targetGroups.addAll(nextResult.getTargetGroups());
    }

    return targetGroups;
  }

  public Optional<TargetGroup> getTargetGroup(String trafficSourceName) {
    DescribeTargetGroupsRequest targetGroupsRequest = new DescribeTargetGroupsRequest()
        .withNames(trafficSourceName);

    try {
      List<TargetGroup> maybeTargetGroup = elbClient
          .describeTargetGroups(targetGroupsRequest)
          .getTargetGroups();

      if (maybeTargetGroup.size() > 0) {
        return Optional.of(maybeTargetGroup.get(0));
      } else {
        return Optional.absent();
      }
    } catch (TargetGroupNotFoundException exn) {
      LOG.warn("Could not find target group with name {}", trafficSourceName);
      return Optional.absent();
    }
  }

  public Collection<Listener> getListenersForElb(String elbName) {
    Optional<LoadBalancer> maybeLoadBalancer = getLoadBalancer(elbName);
    if (maybeLoadBalancer.isPresent()) {
      Collection<Listener> listeners = new HashSet<>();
      DescribeListenersRequest listenersRequest = new DescribeListenersRequest()
          .withLoadBalancerArn(maybeLoadBalancer.get().getLoadBalancerArn());
      DescribeListenersResult result = elbClient.describeListeners(listenersRequest);
      String nextMarker = result.getNextMarker();
      listeners.addAll(result.getListeners());

      while (! Strings.isNullOrEmpty(nextMarker)) {
        listenersRequest = new DescribeListenersRequest()
            .withMarker(nextMarker);
        result = elbClient.describeListeners(listenersRequest);
        nextMarker = result.getNextMarker();
        listeners.addAll(result.getListeners());
      }

      return listeners;
    } else {
      return Collections.emptySet();
    }
  }

  public Collection<Rule> getRulesByListener(String listenerArn) {
    DescribeRulesRequest rulesRequest = new DescribeRulesRequest()
        .withListenerArn(listenerArn);

    return elbClient
        .describeRules(rulesRequest)
        .getRules();
  }

  public Listener createListener(CreateListenerRequest createListenerRequest) {
    return elbClient
        .createListener(createListenerRequest)
        .getListeners()
        .get(0);
  }

  public Listener modifyListener(ModifyListenerRequest modifyListenerRequest) {
    return elbClient
        .modifyListener(modifyListenerRequest)
        .getListeners()
        .get(0);
  }

  public void deleteListener(String listenerArn) {
    DeleteListenerRequest deleteListenerRequest = new DeleteListenerRequest()
        .withListenerArn(listenerArn);

    elbClient
        .deleteListener(deleteListenerRequest);
  }

  public Rule createRule(CreateRuleRequest createRuleRequest) {
    return elbClient
        .createRule(createRuleRequest)
        .getRules()
        .get(0);
  }

  public Rule modifyRule(ModifyRuleRequest modifyRuleRequest) {
    return elbClient
        .modifyRule(modifyRuleRequest)
        .getRules()
        .get(0);
  }

  public void deleteRule(String ruleArn) {
    DeleteRuleRequest deleteRuleRequest = new DeleteRuleRequest()
        .withRuleArn(ruleArn);

    elbClient.deleteRule(deleteRuleRequest);
  }

  public TargetGroup createTargetGroup(CreateTargetGroupRequest targetGroupRequest) {
    return elbClient
        .createTargetGroup(targetGroupRequest)
        .getTargetGroups()
        .get(0);
  }

  public TargetGroup modifyTargetGroup(ModifyTargetGroupRequest modifyTargetGroupRequest) {
    return elbClient
        .modifyTargetGroup(modifyTargetGroupRequest)
        .getTargetGroups()
        .get(0);
  }

  /**
   * @param targetGroup Target group to check
   * @return the instance IDs of the targets in the given target group
   */
  private Collection<String> targetsOn(TargetGroup targetGroup) {
    DescribeTargetHealthRequest targetHealthRequest = new DescribeTargetHealthRequest()
        .withTargetGroupArn(targetGroup.getTargetGroupArn());
    return elbClient
        .describeTargetHealth(targetHealthRequest)
        .getTargetHealthDescriptions().stream()
        .map((t) -> t.getTarget().getId())
        .collect(Collectors.toSet());
  }

  private boolean isVpcOkay(BaragonAgentMetadata agent, TargetGroup targetGroup) {
    if (configuration.isPresent() && configuration.get().isCheckForCorrectVpc()) {
      if (agent.getEc2().getVpcId().isPresent()) {
        String vpcId = agent.getEc2().getVpcId().get();
        return vpcId.equals(targetGroup.getVpcId());
      } else {
        return false;
      }
    } else {
      return true;
    }
  }

  private Collection<LoadBalancer> getLoadBalancersByBaragonGroup(Collection<LoadBalancer> allLoadBalancers, BaragonGroup baragonGroup) {
    Set<TrafficSource> trafficSources = baragonGroup.getTrafficSources();
    Set<String> trafficSourceNames = new HashSet<>();
    Collection<LoadBalancer> loadBalancersForGroup = new HashSet<>();

    for (TrafficSource trafficSource : trafficSources) {
      trafficSourceNames.add(trafficSource.getName());
    }

    for (LoadBalancer loadBalancer : allLoadBalancers) {
      if (trafficSourceNames.contains(loadBalancer.getLoadBalancerName())) {
        loadBalancersForGroup.add(loadBalancer);
      }
    }

    return loadBalancersForGroup;
  }

  /**
   * Ensure that the given baragon agent is attached to the given target group. When this function
   * completes, the baragon agent will be attached to the load balancer, whether or not it originally
   * was.
   *
   * @param baragonAgents BaragonAgent to register with given load balancer
   * @param loadBalancers Load balancer to register with
   */
  private void guaranteeRegistered(TrafficSource trafficSource,
                                   TargetGroup targetGroup,
                                   Collection<TargetDescription> targets,
                                   Collection<BaragonAgentMetadata> baragonAgents,
                                   Collection<LoadBalancer> loadBalancers) {
    /*
    - Check that load balancers, baragon agents, target groups are on same VPC
    - Check that load balancers, targets are on same subnet (== AZ)
    - Check that all baragon agents are associated with a target on target group
    - Check that load balancers has listeners, rules to make talk to target group
     */

    if (configuration.isPresent() && configuration.get().isCheckForCorrectVpc()) {
      guaranteeSameVPC(targetGroup, baragonAgents, loadBalancers);
    }

    guaranteeAzEnabled(baragonAgents, loadBalancers);
    guaranteeHasAllTargets(trafficSource, targetGroup, targets, baragonAgents);
    //guaranteeListenersPresent(targetGroup, loadBalancers);
  }

  /**
   * De-register any targets representing agents that are not known to the BaragonService,
   * or which otherwise need to be removed.
   *
   * @param targetGroup TargetGroup to check for old agents
   * @param agents Known agents, to be used as a reference sheet
   */
  private void deregisterRemovableTargets(TrafficSource trafficSource,
                                          BaragonGroup baragonGroup,
                                          TargetGroup targetGroup,
                                          Collection<BaragonAgentMetadata> agents,
                                          Collection<TargetDescription> targets) {
    Collection<TargetDescription> removableTargets = listRemovableTargets(trafficSource, baragonGroup, targets, agents);

    for (TargetDescription removableTarget : removableTargets) {
      try {
        if (configuration.isPresent()
            && !configuration.get().isRemoveLastHealthyEnabled()
            && isLastHealthyInstance(removableTarget, targetGroup)) {
          LOG.info("Will not de-register target {} because it is last healthy instance in {}", removableTarget, targetGroup);
        } else {
          elbClient.deregisterTargets(new DeregisterTargetsRequest()
              .withTargetGroupArn(targetGroup.getTargetGroupArn())
              .withTargets(removableTarget));
          LOG.info("De-registered target {} from target group {}", removableTarget, targetGroup);
        }
      } catch (AmazonClientException acexn) {
        LOG.error("Could not de-register target {} from target group {} due to error",
            removableTarget, targetGroup, acexn);
        exceptionNotifier.notify(acexn, ImmutableMap
            .of("targetGroup", targetGroup.getTargetGroupName()));
      }
    }
  }

  /**
   * When this method completes, the target group, the agents, and the loadBalancers are
   * all on the same VPC.
   * The target group, each of the agents, and each of the load balancers should think
   * that they are on the same VPC, otherwise they won't be able to talk to each other.
   *
   *
   * @param targetGroup Group - and consequently all targets - to check
   * @param agents Agents to check
   * @param loadBalancers Load balances to check
   */
  private void guaranteeSameVPC(TargetGroup targetGroup,
                                Collection<BaragonAgentMetadata> agents,
                                Collection<LoadBalancer> loadBalancers) {
    String vpcId = targetGroup.getVpcId();

    for (BaragonAgentMetadata agent : agents) {
      if (agent.getEc2().getVpcId().isPresent()) {
        if (! agent.getEc2().getVpcId().get().equals(vpcId)) {
          LOG.error("Agent {} not on same VPC as its target group {}", agent, targetGroup);
          throw new IllegalStateException(String.format("Agent %s not on same VPC as its target group %s", agent, targetGroup));
        }
      } else {
        LOG.error("Agent {} not assigned to a VPC", agent);
        throw new IllegalStateException(String.format("Agent %s not assigned to a VPC", agent));
      }
    }

    for (LoadBalancer loadBalancer : loadBalancers) {
      if (!vpcId.equals(loadBalancer.getVpcId())) {
        LOG.error("Load balancer {} on different VPC from target group {}", loadBalancer, targetGroup);
        throw new IllegalStateException(String.format("Load balancer %s on different VPC from target group %s", loadBalancer, targetGroup));
      }
    }
  }

  private void guaranteeAzEnabled(Collection<BaragonAgentMetadata> agents,
                                  Collection<LoadBalancer> loadBalancers) {
    for (LoadBalancer loadBalancer : loadBalancers) {
      Collection<String> azNames = new HashSet<>();
      for (AvailabilityZone availabilityZone : loadBalancer.getAvailabilityZones()) {
        azNames.add(availabilityZone.getZoneName());
      }

      for (BaragonAgentMetadata agent : agents) {
        if (agent.getEc2().getAvailabilityZone().isPresent()
            && ! azNames.contains(agent.getEc2().getAvailabilityZone().get())) {
          guaranteeHasAllSubnets(
              agent.getEc2().getSubnetId().asSet(),
              loadBalancer);
        }
      }
    }
  }

  /**
   *
   * @param target Target to check
   * @param targetGroup Group to check in
   * @return if the given target is the last healthy target in the given target group
   */
  private boolean isLastHealthyInstance(TargetDescription target, TargetGroup targetGroup) {
    DescribeTargetHealthRequest targetHealthRequest = new DescribeTargetHealthRequest()
        .withTargetGroupArn(targetGroup.getTargetGroupArn());
    List<TargetHealthDescription> targetHealthDescriptions = elbClient
        .describeTargetHealth(targetHealthRequest)
        .getTargetHealthDescriptions();

    boolean instanceIsHealthy = false;
    int healthyCount = 0;

    for (TargetHealthDescription targetHealthDescription : targetHealthDescriptions) {
      if (targetHealthDescription.getTargetHealth().getState()
          .equals(TargetHealthStateEnum.Healthy.toString())) {
        healthyCount += 1;
        if (targetHealthDescription.getTarget().equals(target)) {
          instanceIsHealthy = true;
        }
      }
    }

    return instanceIsHealthy && healthyCount == 1;
  }

  private void guaranteeHasAllSubnets(Collection<String> subnetIds, LoadBalancer loadBalancer) {
    Collection<String> subnetsOnLoadBalancer = getSubnetsFromLoadBalancer(loadBalancer);
    Set<String> subnetsToAdd = new HashSet<>();
    for (String subnetId : subnetIds) {
      if (! subnetsOnLoadBalancer.contains(subnetId)) {
        subnetsToAdd.add(subnetId);
      }
    }

    subnetsToAdd = Sets.union(new HashSet<>(subnetsOnLoadBalancer), subnetsToAdd);

    try {
      SetSubnetsRequest subnetsRequest = new SetSubnetsRequest()
          .withLoadBalancerArn(loadBalancer.getLoadBalancerArn())
          .withSubnets(subnetsToAdd);
      elbClient.setSubnets(subnetsRequest);
    } catch (AmazonClientException acexn) {
      LOG.error("Could not attach subnets {} to load balancer {} due to error",
          subnetsToAdd, loadBalancer.getLoadBalancerName(), acexn);
      exceptionNotifier.notify(acexn, ImmutableMap.of(
          "elb", loadBalancer.getLoadBalancerName(),
          "subnets", subnetsToAdd.toString()));
    }
  }

  /**
   * After this method completes, every agent in baragonAgents should be associated with
   * a target in the given target group.
   *
   * @param targetGroup group to register in
   * @param baragonAgents agents to be registered
   */
  private void guaranteeHasAllTargets(TrafficSource trafficSource,
                                      TargetGroup targetGroup,
                                      Collection<TargetDescription> targets,
                                      Collection<BaragonAgentMetadata> baragonAgents) {
    Collection<TargetDescription> targetDescriptions = new HashSet<>();
    for (BaragonAgentMetadata agent : baragonAgents) {
      try {
        if ((trafficSource.getRegisterBy() == RegisterBy.INSTANCE_ID && agent.getEc2().getInstanceId().isPresent())
            || (trafficSource.getRegisterBy() == RegisterBy.PRIVATE_IP && agent.getEc2().getPrivateIp().isPresent())) {
          String id = trafficSource.getRegisterBy() == RegisterBy.INSTANCE_ID ? agent.getEc2().getInstanceId().get() : agent.getEc2().getPrivateIp().get();
          if (agentShouldRegisterInTargetGroup(id, targets)) {
            targetDescriptions.add(new TargetDescription()
                .withId(id));
            LOG.info("Will register agent {} to target in group {}", agent, targetGroup);
          } else {
            LOG.debug("Agent {} is already registered", agent);
          }
        }
      } catch (Exception exn) {
        LOG.error("Could not create request to register agent {} due to error", agent, exn);
        exceptionNotifier.notify(exn, ImmutableMap.of("agent", agent.toString()));
      }
    }

    if (targetDescriptions.isEmpty()) {
      LOG.debug("No new instances to register with target group");
    } else {
      try {
        RegisterTargetsRequest registerTargetsRequest = new RegisterTargetsRequest()
            .withTargetGroupArn(targetGroup.getTargetGroupArn())
            .withTargets(targetDescriptions);
        elbClient.registerTargets(registerTargetsRequest);
        LOG.info("Registered targets {} onto target group {}", targetDescriptions, targetGroup);
      } catch (AmazonClientException acexn) {
        LOG.error("Failed to register targets {} onto target group {}", targetDescriptions, targetGroup);
        exceptionNotifier.notify(acexn, ImmutableMap.of(
            "targets", targetDescriptions.toString(),
            "targetGroup", targetGroup.toString()));
      }
    }
  }

  private boolean agentShouldRegisterInTargetGroup(String id, Collection<TargetDescription> targets) {
    boolean shouldRegister = true;
    for (TargetDescription target : targets) {
      if (target.getId().equals(id)) {
        shouldRegister = false;
      }
    }
    return shouldRegister;
  }

  private Collection<String> getSubnetsFromLoadBalancer(LoadBalancer loadBalancer) {
    List<AvailabilityZone> availabilityZones = loadBalancer.getAvailabilityZones();
    Set<String> subnetIds = new HashSet<>();

    for (AvailabilityZone availabilityZone : availabilityZones) {
      subnetIds.add(availabilityZone.getSubnetId());
    }

    return subnetIds;
  }

  private Collection<TargetDescription> targetsInTargetGroup(TargetGroup targetGroup) {
    DescribeTargetHealthRequest targetHealthRequest = new DescribeTargetHealthRequest()
        .withTargetGroupArn(targetGroup.getTargetGroupArn());
    List<TargetHealthDescription> targetGroupsResult = elbClient
        .describeTargetHealth(targetHealthRequest)
        .getTargetHealthDescriptions();

    Collection<TargetDescription> targetDescriptions = new HashSet<>();
    for (TargetHealthDescription targetHealthDescription : targetGroupsResult) {
      targetDescriptions.add(targetHealthDescription.getTarget());
    }

    return targetDescriptions;
  }

  private Collection<TargetDescription> listRemovableTargets(TrafficSource trafficSource,
                                                             BaragonGroup baragonGroup,
                                                             Collection<TargetDescription> targetsOnGroup,
                                                             Collection<BaragonAgentMetadata> agentsInBaragonGroup) {
    Collection<String> agentIds = trafficSource.getRegisterBy() == RegisterBy.INSTANCE_ID ? instanceIds(agentsInBaragonGroup) : privateIps(agentsInBaragonGroup);

    Collection<TargetDescription> removableTargets = new HashSet<>();
    for (TargetDescription targetDescription : targetsOnGroup) {
      if (!agentIds.contains(targetDescription.getId()) && canDeregisterAgent(baragonGroup, targetDescription.getId())) {
        LOG.info("Will attempt to deregister target {}", targetDescription.getId());
        removableTargets.add(targetDescription);
      }
    }

    return removableTargets;
  }

  private Collection<String> instanceIds(Collection<BaragonAgentMetadata> agents) {
    Collection<String> instanceIds = new HashSet<>();
    for (BaragonAgentMetadata agent : agents) {
      if (agent.getEc2().getInstanceId().isPresent()) {
        instanceIds.add(agent.getEc2().getInstanceId().get());
      }
    }
    return instanceIds;
  }

  private Collection<String> privateIps(Collection<BaragonAgentMetadata> agents) {
    Collection<String> privateIps = new HashSet<>();
    for (BaragonAgentMetadata agent : agents) {
      if (agent.getEc2().getPrivateIp().isPresent()) {
        privateIps.add(agent.getEc2().getPrivateIp().get());
      }
    }
    return privateIps;
  }

  private Collection<BaragonAgentMetadata> getAgentsByBaragonGroup(BaragonGroup baragonGroup) {
    return loadBalancerDatastore.getAgentMetadata(baragonGroup.getName());
  }
}


package com.hubspot.baragon.data;

import java.util.Collection;
import java.util.Collections;
import java.util.List;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.api.transaction.CuratorTransactionResult;
import org.apache.curator.utils.ZKPaths;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException.NodeExistsException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.codahale.metrics.annotation.Timed;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Optional;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.inject.Inject;
import com.google.inject.Singleton;
import com.hubspot.baragon.config.ZooKeeperConfiguration;
import com.hubspot.baragon.models.BaragonRequest;
import com.hubspot.baragon.models.InternalRequestStates;
import com.hubspot.baragon.models.QueuedRequestId;

@Singleton
public class BaragonRequestDatastore extends AbstractDataStore {
  private static final Logger LOG = LoggerFactory.getLogger(BaragonRequestDatastore.class);

  public static final String REQUESTS_FORMAT = "/request";
  public static final String REQUEST_FORMAT = REQUESTS_FORMAT + "/%s";
  public static final String REQUEST_STATE_FORMAT = REQUEST_FORMAT + "/status";
  public static final String REQUEST_MESSAGE_FORMAT = REQUEST_FORMAT + "/message";

  public static final String REQUEST_QUEUE_FORMAT = "/queue";
  public static final String REQUEST_ENQUEUE_FORMAT = REQUEST_QUEUE_FORMAT + "/%s|%s|";
  public static final String REQUEST_QUEUE_ITEM_FORMAT = REQUEST_QUEUE_FORMAT + "/%s";

  @Inject
  public BaragonRequestDatastore(CuratorFramework curatorFramework, ObjectMapper objectMapper, ZooKeeperConfiguration zooKeeperConfiguration) {
    super(curatorFramework, objectMapper, zooKeeperConfiguration);
  }

  //
  // REQUEST DATA
  //
  @Timed
  public Optional<BaragonRequest> getRequest(String requestId) {
    return readFromZk(String.format(REQUEST_FORMAT, requestId), BaragonRequest.class);
  }

  @Timed
  public Optional<BaragonRequest> deleteRequest(String requestId) {
    final Optional<BaragonRequest> maybeRequest = getRequest(requestId);

    if (maybeRequest.isPresent()) {
      deleteNode(String.format(REQUEST_FORMAT, requestId), true);
    }

    return maybeRequest;
  }

  @Timed
  public List<String> getAllRequestIds() {
    return getChildren(REQUESTS_FORMAT);
  }

  @Timed
  public Optional<Long> getRequestUpdatedAt(String requestId) {
    return getUpdatedAt(String.format(String.format(REQUEST_STATE_FORMAT, requestId)));
  }

  //
  // REQUEST STATE
  //
  public boolean activeRequestExists(String requestId) {
    return nodeExists(String.format(REQUEST_FORMAT, requestId));
  }

  @Timed
  public Optional<InternalRequestStates> getRequestState(String requestId) {
    return readFromZk(String.format(REQUEST_STATE_FORMAT, requestId), InternalRequestStates.class);
  }

  @Timed
  public void setRequestState(String requestId, InternalRequestStates state) {
    writeToZk(String.format(REQUEST_STATE_FORMAT, requestId), state);
  }

  // REQUEST MESSAGE
  @Timed
  public Optional<String> getRequestMessage(String requestId) {
    return readFromZk(String.format(REQUEST_MESSAGE_FORMAT, requestId), String.class);
  }

  @Timed
  public void setRequestMessage(String requestId, String message) {
    writeToZk(String.format(REQUEST_MESSAGE_FORMAT, requestId), message);
  }

  //
  // REQUEST QUEUING
  //
  @Timed
  public QueuedRequestId enqueueRequest(BaragonRequest request, InternalRequestStates state) throws NodeExistsException {
    final long start = System.currentTimeMillis();

    final String queuedRequestPath = String.format(REQUEST_ENQUEUE_FORMAT, request.getLoadBalancerService().getServiceId(), request.getLoadBalancerRequestId());
    final String requestPath = String.format(REQUEST_FORMAT, request.getLoadBalancerRequestId());
    final String requestStatePath = String.format(REQUEST_STATE_FORMAT, request.getLoadBalancerRequestId());

    try {
      if (!nodeExists(REQUESTS_FORMAT)) {
        createNode(REQUESTS_FORMAT);
      }
      if (!nodeExists(REQUEST_QUEUE_FORMAT)) {
        createNode(REQUEST_QUEUE_FORMAT);
      }

      byte[] requestBytes = objectMapper.writeValueAsBytes(request);
      byte[] stateBytes = objectMapper.writeValueAsBytes(state);

      Collection<CuratorTransactionResult> results = curatorFramework.inTransaction()
          .create().forPath(requestPath, requestBytes).and()
          .create().forPath(requestStatePath, stateBytes).and()
          .create().withMode(CreateMode.PERSISTENT_SEQUENTIAL).forPath(queuedRequestPath)
          .and().commit();

      log(OperationType.WRITE, Optional.of(3), Optional.of(requestBytes.length + stateBytes.length), start, String.format("Transaction Paths [%s + %s + %s]", requestPath, requestStatePath, queuedRequestPath));

      return QueuedRequestId.fromString(ZKPaths.getNodeFromPath(Iterables.find(results, CuratorTransactionResult.ofTypeAndPath(org.apache.curator.framework.api.transaction.OperationType.CREATE, queuedRequestPath))
          .getResultPath()));
    } catch (NodeExistsException nee) {
      throw nee;
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }

  @Timed
  public List<QueuedRequestId> getQueuedRequestIds() {
    final List<String> nodes = getChildren(REQUEST_QUEUE_FORMAT);

    Collections.sort(nodes, SEQUENCE_NODE_COMPARATOR_LOW_TO_HIGH);

    final List<QueuedRequestId> queuedRequestIds = Lists.newArrayListWithCapacity(nodes.size());

    for (String node : nodes) {
      queuedRequestIds.add(QueuedRequestId.fromString(node));
    }

    return queuedRequestIds;
  }

  @Timed
  public int getQueuedRequestCount() {
    return getChildren(REQUEST_QUEUE_FORMAT).size();
  }

  @Timed
  public void removeQueuedRequest(QueuedRequestId queuedRequestId) {
    deleteNode(String.format(REQUEST_QUEUE_ITEM_FORMAT, queuedRequestId.buildZkPath()));
  }
}


/**
 * Copyright © 2016 Mercateo AG (http://www.mercateo.com)
 *
 * 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.mercateo.rest.hateoas.client.schema;

import java.io.UnsupportedEncodingException;
import java.net.URLDecoder;
import java.util.HashMap;
import java.util.Map;

import org.glassfish.jersey.uri.UriTemplate;

import com.fasterxml.jackson.annotation.JsonAnyGetter;
import com.fasterxml.jackson.annotation.JsonAnySetter;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.databind.JsonNode;

import lombok.Data;

/**
 * the client representation for links. Note that href is an String because it
 * could be an uri or uri Template
 *
 * @author joerg.adler
 *
 */
@Data
@JsonInclude(Include.NON_NULL)
public class SchemaLink {

	private UriTemplate href;

	private String rel;

	private Map<String, String> map = new HashMap<>();

	private JsonNode schema;

	private JsonNode targetSchema;

	private String mediaType;

	// backwards compatibility, would be canged in next version
	public void setHref(String href) {
		if (href.contains("%7B")) {
			try {
				href = URLDecoder.decode(href, "UTF-8");
			} catch (UnsupportedEncodingException e) {
				throw new RuntimeException(e);
			}
		}
		this.href = new UriTemplate(href);
	}

	@JsonAnyGetter
	public Map<String, String> getMap() {
		return map;
	}

	@JsonAnySetter
	public void setMap(String key, String value) {
		map.put(key, value);
	}
}


/*
 * Copyright 2017 Trimou Team
 *
 * 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 org.trimou.extension.spring.starter;

import java.nio.charset.StandardCharsets;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.autoconfigure.template.AbstractTemplateViewResolverProperties;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.trimou.engine.MustacheEngineBuilder;
import org.trimou.engine.config.EngineConfigurationKey;
import org.trimou.spring4.web.SpringResourceTemplateLocator;
import org.trimou.util.Strings;

/**
 * {@link ConfigurationProperties} for Trimou
 */
@ConfigurationProperties(prefix = TrimouProperties.PROPERTY_PREFIX)
public final class TrimouProperties extends AbstractTemplateViewResolverProperties {

    static final String PROPERTY_PREFIX = "trimou";
    private static final Logger LOGGER = LoggerFactory.getLogger(TrimouProperties.class);

    /**
     * Prefix to apply to template names.
     */
    private String prefix = SpringResourceTemplateLocator.DEFAULT_PREFIX;

    /**
     * Suffix to apply to template names.
     */
    private String suffix = SpringResourceTemplateLocator.DEFAULT_SUFFIX;

    /**
     * Start delimiter
     */
    private String startDelimiter = getDefaultStringValue(EngineConfigurationKey.START_DELIMITER);

    /**
     * End delimiter
     */
    private String endDelimiter = getDefaultStringValue(EngineConfigurationKey.END_DELIMITER);

    /**
     * Precompile all templates
     */
    private boolean precompileTemplates = getDefaultBooleanValue(EngineConfigurationKey.PRECOMPILE_ALL_TEMPLATES);

    /**
     * Remove standalone lines
     */
    private boolean removeStandaloneLines = getDefaultBooleanValue(EngineConfigurationKey.REMOVE_STANDALONE_LINES);

    /**
     * Remove unnecessary segments
     */
    private boolean removeUnnecessarySegments =
            getDefaultBooleanValue(EngineConfigurationKey.REMOVE_UNNECESSARY_SEGMENTS);

    /**
     * Debug mode
     */
    private boolean debugMode = getDefaultBooleanValue(EngineConfigurationKey.DEBUG_MODE);

    /**
     * Cache section-based literal blocks
     */
    private boolean cacheSectionLiteralBlock =
            getDefaultBooleanValue(EngineConfigurationKey.CACHE_SECTION_LITERAL_BLOCK);

    /**
     * Template recursive invocation limit
     */
    private int templateRecursiveInvocationLimit =
            getDefaultIntegerValue(EngineConfigurationKey.TEMPLATE_RECURSIVE_INVOCATION_LIMIT);

    /**
     * Skip value escaping
     */
    private boolean skipValueEscaping = getDefaultBooleanValue(EngineConfigurationKey.SKIP_VALUE_ESCAPING);

    /**
     * Template cache expiration timeout
     */
    private long templateCacheExpirationTimeout =
            getDefaultLongValue(EngineConfigurationKey.TEMPLATE_CACHE_EXPIRATION_TIMEOUT);

    /**
     * Enable helper
     */
    private boolean enableHelper = getDefaultBooleanValue(EngineConfigurationKey.HANDLEBARS_SUPPORT_ENABLED);

    /**
     * Reuse line separator segments
     */
    private boolean reuseLineSeparatorSegments =
            getDefaultBooleanValue(EngineConfigurationKey.REUSE_LINE_SEPARATOR_SEGMENTS);

    /**
     * Iteration meta-data alias
     */
    private String iterationMetadataAlias = getDefaultStringValue(EngineConfigurationKey.ITERATION_METADATA_ALIAS);

    /**
     * Enable resolver hints
     */
    private boolean enableResolverHints = getDefaultBooleanValue(EngineConfigurationKey.RESOLVER_HINTS_ENABLED);

    /**
     * Enable nested template support
     */
    private boolean enableNestedTemplates =
            getDefaultBooleanValue(EngineConfigurationKey.NESTED_TEMPLATE_SUPPORT_ENABLED);

    /**
     * Cache template sources
     */
    private boolean cacheTemplateSources =
            getDefaultBooleanValue(EngineConfigurationKey.TEMPLATE_CACHE_USED_FOR_SOURCE);

    public TrimouProperties() {
        super(SpringResourceTemplateLocator.DEFAULT_PREFIX, SpringResourceTemplateLocator.DEFAULT_SUFFIX);
        final boolean cacheEnabled = getDefaultBooleanValue(EngineConfigurationKey.TEMPLATE_CACHE_ENABLED);
        setCache(cacheEnabled);
        if (getCharset() == null) {
            setCharset(StandardCharsets.UTF_8);
        }
    }

    private static boolean getDefaultBooleanValue(final EngineConfigurationKey key) {
        final Object object = key.getDefaultValue();
        if (object != null) {
            final String value = String.valueOf(object);
            if (!Strings.isEmpty(value)) {
                return Boolean.valueOf(value);
            }
        }
        LOGGER.warn("Unable to decode key '{}' as boolean. Return false", key);
        return false;
    }

    private static String getDefaultStringValue(final EngineConfigurationKey key) {
        final Object object = key.getDefaultValue();
        if (object != null) {
            return String.valueOf(object);
        }
        LOGGER.warn("Unable to decode key '{}' as string. Return null", key);
        return null;
    }

    private static int getDefaultIntegerValue(final EngineConfigurationKey key) {
        final Object object = key.getDefaultValue();
        if (object != null) {
            final String value = String.valueOf(object);
            if (!Strings.isEmpty(value)) {
                return Integer.parseInt(value);
            }
        }
        LOGGER.warn("Unable to decode key '{}' as int. Return 0", key);
        return 0;
    }

    private static long getDefaultLongValue(final EngineConfigurationKey key) {
        final Object object = key.getDefaultValue();
        if (object != null) {
            final String value = String.valueOf(object);
            if (!Strings.isEmpty(value)) {
                return Long.parseLong(value);
            }
        }
        LOGGER.warn("Unable to decode key '{}' as long. Return 0L", key);
        return 0L;
    }

    @Override
    public String getPrefix() {
        return this.prefix;
    }

    @Override
    public void setPrefix(final String prefix) {
        this.prefix = prefix;
    }

    @Override
    public String getSuffix() {
        return this.suffix;
    }

    @Override
    public void setSuffix(final String suffix) {
        this.suffix = suffix;
    }

    public String getStartDelimiter() {
        return startDelimiter;
    }

    public void setStartDelimiter(final String startDelimiter) {
        this.startDelimiter = startDelimiter;
    }

    public String getEndDelimiter() {
        return endDelimiter;
    }

    public void setEndDelimiter(final String endDelimiter) {
        this.endDelimiter = endDelimiter;
    }

    public boolean isPrecompileTemplates() {
        return precompileTemplates;
    }

    public void setPrecompileTemplates(final boolean precompileTemplates) {
        this.precompileTemplates = precompileTemplates;
    }

    public boolean isRemoveStandaloneLines() {
        return removeStandaloneLines;
    }

    public void setRemoveStandaloneLines(final boolean removeStandaloneLines) {
        this.removeStandaloneLines = removeStandaloneLines;
    }

    public boolean isRemoveUnnecessarySegments() {
        return removeUnnecessarySegments;
    }

    public void setRemoveUnnecessarySegments(final boolean removeUnnecessarySegments) {
        this.removeUnnecessarySegments = removeUnnecessarySegments;
    }

    public boolean isDebugMode() {
        return debugMode;
    }

    public void setDebugMode(final boolean debugMode) {
        this.debugMode = debugMode;
    }

    public boolean isCacheSectionLiteralBlock() {
        return cacheSectionLiteralBlock;
    }

    public void setCacheSectionLiteralBlock(final boolean cacheSectionLiteralBlock) {
        this.cacheSectionLiteralBlock = cacheSectionLiteralBlock;
    }

    public int getTemplateRecursiveInvocationLimit() {
        return templateRecursiveInvocationLimit;
    }

    public void setTemplateRecursiveInvocationLimit(final int templateRecursiveInvocationLimit) {
        this.templateRecursiveInvocationLimit = templateRecursiveInvocationLimit;
    }

    public boolean isSkipValueEscaping() {
        return skipValueEscaping;
    }

    public void setSkipValueEscaping(final boolean skipValueEscaping) {
        this.skipValueEscaping = skipValueEscaping;
    }

    public long getTemplateCacheExpirationTimeout() {
        return templateCacheExpirationTimeout;
    }

    public void setTemplateCacheExpirationTimeout(final long templateCacheExpirationTimeout) {
        this.templateCacheExpirationTimeout = templateCacheExpirationTimeout;
    }

    public boolean isEnableHelper() {
        return enableHelper;
    }

    public void setEnableHelper(final boolean enableHelper) {
        this.enableHelper = enableHelper;
    }

    public boolean isReuseLineSeparatorSegments() {
        return reuseLineSeparatorSegments;
    }

    public void setReuseLineSeparatorSegments(final boolean reuseLineSeparatorSegments) {
        this.reuseLineSeparatorSegments = reuseLineSeparatorSegments;
    }

    public String getIterationMetadataAlias() {
        return iterationMetadataAlias;
    }

    public void setIterationMetadataAlias(final String iterationMetadataAlias) {
        this.iterationMetadataAlias = iterationMetadataAlias;
    }

    public boolean isEnableResolverHints() {
        return enableResolverHints;
    }

    public void setEnableResolverHints(final boolean enableResolverHints) {
        this.enableResolverHints = enableResolverHints;
    }

    public boolean isEnableNestedTemplates() {
        return enableNestedTemplates;
    }

    public void setEnableNestedTemplates(final boolean enableNestedTemplates) {
        this.enableNestedTemplates = enableNestedTemplates;
    }

    public boolean isCacheTemplateSources() {
        return cacheTemplateSources;
    }

    public void setCacheTemplateSources(final boolean cacheTemplateSources) {
        this.cacheTemplateSources = cacheTemplateSources;
    }

    /**
     * Apply the {@link TrimouProperties} to a {@link MustacheEngineBuilder}.
     *
     * @param engineBuilder the Trimou mustache engine builder to apply the properties to
     */
    public void applyToTrimouMustacheEngineBuilder(final MustacheEngineBuilder engineBuilder) {
        engineBuilder
                .setProperty(EngineConfigurationKey.START_DELIMITER, getStartDelimiter())
                .setProperty(EngineConfigurationKey.END_DELIMITER, getEndDelimiter())
                .setProperty(EngineConfigurationKey.PRECOMPILE_ALL_TEMPLATES, isPrecompileTemplates())
                .setProperty(EngineConfigurationKey.REMOVE_STANDALONE_LINES, isRemoveStandaloneLines())
                .setProperty(EngineConfigurationKey.REMOVE_UNNECESSARY_SEGMENTS, isRemoveUnnecessarySegments())
                .setProperty(EngineConfigurationKey.DEBUG_MODE, isDebugMode())
                .setProperty(EngineConfigurationKey.CACHE_SECTION_LITERAL_BLOCK, isCacheSectionLiteralBlock())
                .setProperty(EngineConfigurationKey.TEMPLATE_RECURSIVE_INVOCATION_LIMIT,
                        getTemplateRecursiveInvocationLimit())
                .setProperty(EngineConfigurationKey.SKIP_VALUE_ESCAPING, isSkipValueEscaping())
                .setProperty(EngineConfigurationKey.DEFAULT_FILE_ENCODING, getCharset().name())
                .setProperty(EngineConfigurationKey.TEMPLATE_CACHE_ENABLED, isCache())
                .setProperty(EngineConfigurationKey.TEMPLATE_CACHE_EXPIRATION_TIMEOUT,
                        getTemplateCacheExpirationTimeout())
                .setProperty(EngineConfigurationKey.HANDLEBARS_SUPPORT_ENABLED, isEnableHelper())
                .setProperty(EngineConfigurationKey.REUSE_LINE_SEPARATOR_SEGMENTS, isReuseLineSeparatorSegments())
                .setProperty(EngineConfigurationKey.ITERATION_METADATA_ALIAS, getIterationMetadataAlias())
                .setProperty(EngineConfigurationKey.RESOLVER_HINTS_ENABLED, isEnableResolverHints())
                .setProperty(EngineConfigurationKey.NESTED_TEMPLATE_SUPPORT_ENABLED, isEnableNestedTemplates())
                .setProperty(EngineConfigurationKey.TEMPLATE_CACHE_USED_FOR_SOURCE, isCacheTemplateSources());
    }
}


/*
 * Copyright 2013 Martin Kouba
 *
 * 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 org.trimou.engine.resolver;

import org.trimou.util.Reflections;

/**
 *
 * @author Martin Kouba
 */
final class MemberKey {

    private final Class<?> clazz;

    private final String name;

    public Class<?> getClazz() {
        return clazz;
    }

    public String getName() {
        return name;
    }

    MemberKey(Class<?> clazz, String name) {
        this.clazz = clazz;
        this.name = name;
    }

    @Override
    public int hashCode() {
        final int prime = 31;
        int result = 1;
        result = prime * result + ((clazz == null) ? 0 : clazz.hashCode());
        result = prime * result + ((name == null) ? 0 : name.hashCode());
        return result;
    }

    @Override
    public boolean equals(Object obj) {
        if (this == obj)
            return true;
        if (obj == null)
            return false;
        if (!(obj instanceof MemberKey))
            return false;
        MemberKey other = (MemberKey) obj;
        if (clazz == null) {
            if (other.clazz != null)
                return false;
        } else if (!clazz.equals(other.clazz))
            return false;
        if (name == null) {
            if (other.name != null)
                return false;
        } else if (!name.equals(other.name))
            return false;
        return true;
    }

    @Override
    public String toString() {
        return String.format("MemberKey [clazz: %s, name: %s]", clazz, name);
    }

    /**
     * Simple factory method to obtain a new instance for the given context
     * object and name.
     *
     * @param contextObject
     * @param name
     * @return a new instance
     */
    static MemberKey newInstance(Object contextObject, String name) {
        if (contextObject instanceof Class<?>) {
            Class<?> clazz = (Class<?>) contextObject;
            if (clazz.isEnum() && ("values".equals(name) || Reflections.isConstantName(clazz, name))) {
                // Special handling for enums - allows to access values() and constants
                return new MemberKey(clazz, name);
            }
        }
        return new MemberKey(contextObject.getClass(), name);
    }

}


/*
 * Copyright 2017 Trimou Team
 *
 * 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 org.trimou.spring4.web;

import java.util.Locale;

import org.springframework.web.servlet.View;
import org.springframework.web.servlet.ViewResolver;
import org.springframework.web.servlet.view.AbstractTemplateViewResolver;
import org.trimou.Mustache;
import org.trimou.engine.MustacheEngine;
import org.trimou.engine.MustacheEngineBuilder;
import org.trimou.engine.config.EngineConfigurationKey;

/**
 * Spring MVC {@link ViewResolver} for Trimou.
 */
public final class TrimouViewResolver extends AbstractTemplateViewResolver {

    static final String DEFAULT_CONTENT_TYPE = "text/html;charset=UTF-8";
    private SpringResourceTemplateLocator loader = new SpringResourceTemplateLocator();
    private MustacheEngine engine;

    public TrimouViewResolver() {
        init();
        this.engine = MustacheEngineBuilder.newBuilder()
                .addTemplateLocator(loader)
                .build();
    }

    public TrimouViewResolver(final MustacheEngine engine) {
        init();
        this.engine = engine;
    }

    private void init() {
        setViewClass(requiredViewClass());
        setContentType(DEFAULT_CONTENT_TYPE);
        setPrefix(loader.getPrefix());
        setSuffix(loader.getSuffix());
    }

    /**
     * Returns true, if the cache of the {@link MustacheEngine} is enabled and the debug mode is disabled. Also Spring's
     * view resolution caching must be enabled.
     */
    @Override
    public boolean isCache() {
        return engine.getConfiguration().getBooleanPropertyValue(EngineConfigurationKey.TEMPLATE_CACHE_ENABLED)
                && !engine.getConfiguration().getBooleanPropertyValue(EngineConfigurationKey.DEBUG_MODE)
                && super.isCache();
    }

    @Override
    protected Class<?> requiredViewClass() {
        return TrimouView.class;
    }

    /**
     * Set the mustache engine
     *
     * @param engine the mustache engine
     */
    public void setEngine(final MustacheEngine engine) {
        this.engine = engine;
    }

    @Override
    protected View loadView(final String viewName, final Locale locale) throws Exception {
        final Mustache mustache = engine.getMustache(viewName);
        if (mustache != null) {
            final TrimouView trimouView = (TrimouView) super.loadView(viewName, locale);
            trimouView.setMustache(mustache);
            return trimouView;
        }
        return null;
    }

    @Override
    public void setPrefix(final String prefix) {
        super.setPrefix(prefix);
        loader.setPrefix(prefix);
    }

    @Override
    public void setSuffix(final String suffix) {
        super.setSuffix(suffix);
        loader.setSuffix(suffix);
    }
}


/*
 * Copyright 2013 Martin Kouba
 *
 * 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 org.trimou.engine.parser;

import static org.trimou.engine.config.EngineConfigurationKey.END_DELIMITER;
import static org.trimou.engine.config.EngineConfigurationKey.START_DELIMITER;
import static org.trimou.util.Checker.checkArgumentNotEmpty;
import static org.trimou.util.Checker.checkArgumentsNotNull;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.Reader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.trimou.engine.MustacheEngine;
import org.trimou.engine.MustacheTagType;
import org.trimou.engine.config.EngineConfigurationKey;
import org.trimou.exception.MustacheException;
import org.trimou.exception.MustacheProblem;
import org.trimou.util.ImmutableSet;
import org.trimou.util.Strings;

/**
 * The default parser. It's not thread-safe and may not be reused.
 *
 * @author Martin Kouba
 */
class DefaultParser implements Parser {

    private static final Logger LOGGER = LoggerFactory
            .getLogger(DefaultParser.class);

    private MustacheEngine engine;

    private State state;

    private int line;

    private int delimiterIdx;

    private boolean triple;

    private StringBuilder buffer;

    private final Delimiters delimiters;

    private ParsingHandler handler;

    private final Set<String> supportedSeparators;

    private List<String> lastMatchedSeparators;

    private Set<Character> zeroIndexNonSeparatorCharacters;

    private int separatorIdx;

    /**
     *
     * @param engine
     */
    public DefaultParser(MustacheEngine engine) {
        this.state = State.TEXT;
        this.line = 1;
        this.delimiterIdx = 0;
        this.triple = false;
        this.buffer = new StringBuilder();
        this.separatorIdx = 0;
        this.engine = engine;
        this.delimiters = new Delimiters(engine.getConfiguration()
                .getStringPropertyValue(START_DELIMITER), engine
                .getConfiguration().getStringPropertyValue(END_DELIMITER));
        this.supportedSeparators = ImmutableSet.of(Strings.LINE_SEPARATOR_LF,
                Strings.LINE_SEPARATOR_CR, Strings.LINE_SEPARATOR_CRLF);
        this.zeroIndexNonSeparatorCharacters = new HashSet<>();
    }

    public void parse(String name, Reader reader, ParsingHandler handler) {
        checkArgumentNotEmpty(name);
        checkArgumentsNotNull(reader, handler);
        this.handler = handler;
        reader = ensureBufferedReader(reader);

        try {

            // Start of document
            handler.startTemplate(name, delimiters, engine);

            int val;
            while ((val = reader.read()) != -1) {
                processCharacter((char) val);
            }

            if (buffer.length() > 0) {
                if (state == State.TEXT) {
                    // Flush the last text segment
                    flushText();
                } else {
                    throw new MustacheException(
                            MustacheProblem.COMPILE_INVALID_TEMPLATE,
                            "Unexpected non-text buffer at the end of the document (probably unterminated tag): %s",
                            buffer);
                }
            }

            if (state == State.LINE_SEPARATOR) {
                // Flush the last line separator
                lineSeparatorFound(lastMatchedSeparators.get(0));
            }

            // End of document
            handler.endTemplate();

        } catch (IOException e) {
            throw new MustacheException(MustacheProblem.COMPILE_IO_ERROR, e);
        }
    }

    private void processCharacter(char character) {
        switch (state) {
        case TEXT:
            text(character);
            break;
        case START_TAG:
            startTag(character);
            break;
        case TAG:
            tag(character);
            break;
        case END_TAG:
            endTag(character);
            break;
        case LINE_SEPARATOR:
            lineSeparator(character);
            break;
        default:
            throw new IllegalStateException("Unknown parsing state");
        }
    }

    private void text(char character) {
        if (character == delimiters.getStart(0)) {
            if (delimiters.isStartOver(delimiterIdx)) {
                tagStartFound();
            } else {
                // Probably multi-char start tag
                state = State.START_TAG;
                delimiterIdx = 1;
            }
        } else if ((lastMatchedSeparators = findMatchingSeparators(character, 0))
                .size() > 0) {
            if (lastMatchedSeparators.size() == 1
                    && lastMatchedSeparators.get(0).length() == 1) {
                // Single-char separator
                lineSeparatorFound(lastMatchedSeparators.get(0));
            } else if ((lastMatchedSeparators.size() > 1)
                    || (lastMatchedSeparators.size() == 1 && lastMatchedSeparators
                            .get(0).length() > 1)) {
                // Multiple separators or multi-char separator
                state = State.LINE_SEPARATOR;
                separatorIdx = 1;
            }
        } else {
            buffer.append(character);
        }
    }

    private void startTag(char character) {
        if (character == delimiters.getStart(delimiterIdx)) {
            if (delimiters.isStartOver(delimiterIdx)) {
                tagStartFound();
            } else {
                delimiterIdx++;
            }
        } else {
            // False alarm - not a start delimiter
            state = State.TEXT;
            buffer.append(delimiters.getStartPart(delimiterIdx));
            delimiterIdx = 0;
            processCharacter(character);
        }
    }

    private void tag(char character) {
        if (character == delimiters.getEnd(0)) {
            if (triple) {
                // Triple mustache detected - skip first ending
                // mustache
                buffer.append(character);
                triple = false;
            } else if (delimiters.isEndOver(delimiterIdx)) {
                // One char delimiter
                flushTag();
            } else {
                // Ending tag
                state = State.END_TAG;
                delimiterIdx = 1;
            }
        } else {
            if (character == delimiters.getStart(0) && buffer.length() == 0) {
                // Most likely a triple mustache
                triple = true;
            }
            buffer.append(character);
        }
    }

    private void endTag(char character) {
        if (character == delimiters.getEnd(delimiterIdx)) {
            if (delimiters.isEndOver(delimiterIdx)) {
                flushTag();
            } else {
                delimiterIdx++;
            }
        } else {
            // False alarm - not an end delimiter
            LOGGER.info(
                    "Tag contains a part of the end delimiter - most probably an invalid key [part: {}, line: {}]",
                    delimiters.getStartPart(delimiterIdx), line);
            state = State.TAG;
            buffer.append(delimiters.getEndPart(delimiterIdx));
            buffer.append(character);
            delimiterIdx = 0;
        }
    }

    private void lineSeparator(char character) {

        List<String> matched = findMatchingSeparators(character, separatorIdx);

        if (matched.isEmpty()) {
            // Single-char separator
            for (String separator : lastMatchedSeparators) {
                if (separator.length() == separatorIdx) {
                    lineSeparatorFound(separator);
                    processCharacter(character);
                }
            }
        } else if (matched.size() == 1) {
            // Multi-char separator
            lineSeparatorFound(matched.get(0));
        } else if (matched.size() > 1) {
            lastMatchedSeparators = matched;
            separatorIdx++;
        }
    }

    /**
     * Line separator end - flush.
     *
     * @param lineSeparator
     */
    private void lineSeparatorFound(String lineSeparator) {
        flushText();
        flushLineSeparator(lineSeparator);
        line++;
        state = State.TEXT;
        separatorIdx = 0;
    }

    /**
     * Real tag start, flush text if any.
     */
    private void tagStartFound() {
        state = State.TAG;
        delimiterIdx = 0;
        flushText();
    }

    private void flushText() {
        if (buffer.length() > 0) {
            handler.text(buffer.toString());
            clearBuffer();
        }
    }

    /**
     * Real tag end - flush.
     */
    private void flushTag() {
        state = State.TEXT;
        handler.tag(deriveTag(buffer.toString()));
        delimiterIdx = 0;
        clearBuffer();
    }

    private void flushLineSeparator(String separator) {
        handler.lineSeparator(separator);
    }

    private Reader ensureBufferedReader(Reader reader) {
        return reader instanceof BufferedReader ? reader : new BufferedReader(
                reader);
    }

    private ParsedTag deriveTag(String buffer) {
        MustacheTagType type = identifyTagType(buffer);
        String key = extractContent(type, buffer);
        return new ParsedTag(key, type);
    }

    /**
     * Identify the tag type (variable, comment, etc.).
     *
     * @param buffer
     * @param delimiters
     * @return the tag type
     */
    private MustacheTagType identifyTagType(String buffer) {

        if (buffer.length() == 0) {
            return MustacheTagType.VARIABLE;
        }

        // Triple mustache is supported for default delimiters only
        if (delimiters.hasDefaultDelimitersSet()
                && buffer.charAt(0) == ((String) EngineConfigurationKey.START_DELIMITER
                        .getDefaultValue()).charAt(0)
                && buffer.charAt(buffer.length() - 1) == ((String) EngineConfigurationKey.END_DELIMITER
                        .getDefaultValue()).charAt(0)) {
            return MustacheTagType.UNESCAPE_VARIABLE;
        }

        Character command = buffer.charAt(0);

        for (MustacheTagType type : MustacheTagType.values()) {
            if (command.equals(type.getCommand())) {
                return type;
            }
        }
        return MustacheTagType.VARIABLE;
    }

    /**
     * Extract the tag content.
     *
     * @param buffer
     * @return
     */
    private String extractContent(MustacheTagType tagType, String buffer) {

        switch (tagType) {
        case VARIABLE:
            return buffer.trim();
        case UNESCAPE_VARIABLE:
            return (buffer.charAt(0) == ((String) EngineConfigurationKey.START_DELIMITER
                    .getDefaultValue()).charAt(0) ? buffer.substring(1,
                    buffer.length() - 1).trim() : buffer.substring(1).trim());
        case SECTION:
        case INVERTED_SECTION:
        case PARTIAL:
        case EXTEND:
        case EXTEND_SECTION:
        case SECTION_END:
        case NESTED_TEMPLATE:
        case COMMENT:
            return buffer.substring(1).trim();
        case DELIMITER:
            return buffer.trim();
        default:
            return null;
        }
    }

    private void clearBuffer() {
        this.buffer = new StringBuilder();
    }

    /**
     *
     * @param character
     * @param atIndex
     * @return the list of matching line separators
     */
    private List<String> findMatchingSeparators(char character, int atIndex) {

        if (atIndex == 0 && zeroIndexNonSeparatorCharacters.contains(character)) {
            return Collections.emptyList();
        }

        List<String> matchedSeparators = new ArrayList<>(
                supportedSeparators.size());

        for (String separator : supportedSeparators) {
            if (separator.length() > atIndex
                    && separator.charAt(atIndex) == character) {
                matchedSeparators.add(separator);
            }
        }
        if (atIndex == 0 && matchedSeparators.isEmpty()) {
            zeroIndexNonSeparatorCharacters.add(character);
        }
        return matchedSeparators;
    }

    private enum State {

        TEXT,
        START_TAG,
        TAG,
        END_TAG,
        LINE_SEPARATOR;

    }

}


/*
 * Copyright 2013 Martin Kouba
 *
 * 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 org.trimou.engine.resource;

import java.util.LinkedList;
import java.util.List;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.trimou.annotations.Internal;

/**
 * Abstract {@link ReleaseCallbackContainer}. It's not thread-safe.
 *
 * @author Martin Kouba
 */
@Internal
public abstract class AbstractReleaseCallbackContainer
        implements ReleaseCallbackContainer {

    private static final Logger LOGGER = LoggerFactory
            .getLogger(AbstractReleaseCallbackContainer.class);

    private List<ReleaseCallback> releaseCallbacks = null;

    /**
     * Release all the resources, i.e. invoke {@link ReleaseCallback#release()}
     * on each callback. All the callbacks are always invoked, even if one of
     * the invocation fails (throws unchecked exception). Callbacks are invoked
     * in the order in which they were registered.
     */
    public void release() {
        if (releaseCallbacks == null) {
            return;
        }
        for (ReleaseCallback callback : releaseCallbacks) {
            try {
                callback.release();
            } catch (Exception e) {
                LOGGER.warn(
                        "Exception occured during release callback invocation:",
                        e);
            }
        }
    }

    @Override
    public void registerReleaseCallback(ReleaseCallback callback) {
        if (releaseCallbacks == null) {
            releaseCallbacks = new LinkedList<>();
        }
        releaseCallbacks.add(callback);
    }

}