package com.hubspot.singularity.scheduler;

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

import javax.inject.Inject;
import javax.inject.Singleton;

import com.google.common.base.Optional;
import com.hubspot.baragon.models.BaragonRequest;
import com.hubspot.baragon.models.BaragonService;
import com.hubspot.baragon.models.UpstreamInfo;
import com.hubspot.singularity.LoadBalancerRequestType;
import com.hubspot.singularity.LoadBalancerRequestType.LoadBalancerRequestId;
import com.hubspot.singularity.SingularityDeploy;
import com.hubspot.singularity.SingularityLoadBalancerUpdate;
import com.hubspot.singularity.SingularityRequest;
import com.hubspot.singularity.SingularityRequestWithState;
import com.hubspot.singularity.SingularityTask;
import com.hubspot.singularity.SingularityTaskId;
import com.hubspot.singularity.data.DeployManager;
import com.hubspot.singularity.data.RequestManager;
import com.hubspot.singularity.data.TaskManager;
import com.hubspot.singularity.hooks.LoadBalancerClientImpl;

@Singleton
public class SingularityUpstreamChecker {

  private final LoadBalancerClientImpl lbClient;
  private final TaskManager taskManager;
  private final RequestManager requestManager;
  private final DeployManager deployManager;

  @Inject
  public SingularityUpstreamChecker(LoadBalancerClientImpl lbClient, TaskManager taskManager, RequestManager requestManager, DeployManager deployManager) {
    this.lbClient = lbClient;
    this.taskManager = taskManager;
    this.requestManager = requestManager;
    this.deployManager = deployManager;
  }

  private List<SingularityTask> getActiveSingularityTasksForRequest(String requestId) {
    List<SingularityTask> activeSingularityTasksForRequest = new ArrayList<>();
    for (SingularityTaskId taskId: taskManager.getActiveTaskIdsForRequest(requestId)){
      Optional<SingularityTask> maybeTask = taskManager.getTask(taskId);
      if (maybeTask.isPresent()) {
        activeSingularityTasksForRequest.add(maybeTask.get());
      }
    }
    return activeSingularityTasksForRequest;
  }

  private Collection<UpstreamInfo> getUpstreamsFromActiveTasks(String requestId, Optional<String> loadBalancerUpstreamGroup) {
    return lbClient.getUpstreamsForTasks(getActiveSingularityTasksForRequest(requestId), requestId, loadBalancerUpstreamGroup);
  }

  private SingularityLoadBalancerUpdate syncUpstreamsForService(SingularityRequest request, String requestId, SingularityDeploy deploy, Optional<String> loadBalancerUpstreamGroup) {
    Collection<UpstreamInfo> upstreamsInBaragonForRequest = lbClient.getBaragonUpstreamsForRequest(requestId);
    Collection<UpstreamInfo> upstreamsInSingularityForRequest = getUpstreamsFromActiveTasks(requestId, loadBalancerUpstreamGroup);
    upstreamsInBaragonForRequest.removeAll(upstreamsInSingularityForRequest);
    LoadBalancerRequestId loadBalancerRequestId = new LoadBalancerRequestId(requestId, LoadBalancerRequestType.REMOVE, Optional.absent());
    return lbClient.makeAndSendBaragonRequest(loadBalancerRequestId, Collections.emptyList(), new ArrayList<>(upstreamsInBaragonForRequest), deploy, request);
  }

  public void syncUpstreams() {
    for (SingularityRequestWithState singularityRequestWithState: requestManager.getActiveRequests()){
      SingularityRequest request = singularityRequestWithState.getRequest();
      if (request.isLoadBalanced()) {
        //TODO: lock on the requestId
        String requestId = singularityRequestWithState.getRequest().getId();
        String deployId = deployManager.getInUseDeployId(requestId).get(); //TODO: handle when it's absent
        SingularityDeploy deploy = deployManager.getDeploy(requestId, deployId).get();
        Optional<String> loadBalancerUpstreamGroup = deploy.getLoadBalancerUpstreamGroup(); // TODO: handle when absent
        syncUpstreamsForService(request, requestId, deploy, loadBalancerUpstreamGroup);
      }
    }
  }
}

@@ -168,11 +168,69 @@
 til.
-Set
+Map;%0Aimport java.util.Set;%0Aimport java.util.stream.Collectors
 ;%0A%0Ai
@@ -1738,35 +1738,24 @@
 k%3E getActive
-Singularity
 TasksForRequ
@@ -1782,20 +1782,25 @@
 ) %7B%0A    
-List
+final Map
 %3CSingula
@@ -1811,16 +1811,12 @@
 Task
-%3E active
+Id, 
 Sing
@@ -1818,32 +1818,44 @@
  SingularityTask
+%3E activeTask
 sForRequest = ne
@@ -1856,62 +1856,29 @@
 t = 
-new ArrayList%3C%3E();%0A    for (SingularityTaskId taskId: 
+taskManager.getTasks(
 task
@@ -1927,238 +1927,68 @@
 Id))
-%7B
+;
 %0A    
-  Optional%3CSingularityTask%3E maybeTask = taskManager.getTask(taskId);%0A      if (maybeTask.isPresent()) %7B%0A        activeSingularityTasksForRequest.add(maybeTask.get());%0A      %7D%0A    %7D%0A    return activeSingularityTasksForRequest
+return new ArrayList%3C%3E(activeTasksForRequest.values())
 ;%0A  
@@ -2052,16 +2052,26 @@
 iveTasks
+ForRequest
 (String 
@@ -2177,27 +2177,16 @@
 etActive
-Singularity
 TasksFor
@@ -2633,16 +2633,26 @@
 iveTasks
+ForRequest
 (request
@@ -2765,16 +2765,115 @@
 t);%0A    
+final List%3CUpstreamInfo%3E extraUpstreams = new ArrayList%3C%3E(upstreamsInBaragonForRequest);%0A    final 
 LoadBala
@@ -3097,53 +3097,22 @@
 (), 
-new ArrayList%3C%3E(upstreamsInBaragonForRequest)
+extraUpstreams
 , de
@@ -3386,46 +3386,8 @@
 ) %7B%0A
-        //TODO: lock on the requestId%0A
     
@@ -3458,16 +3458,46 @@
 getId();
+ //TODO: lock on the requestId
 %0A       

improve getting active tasks for a request and their upstreams