package com.hubspot.singularity.scheduler;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import javax.inject.Singleton;

import org.apache.mesos.v1.Protos.Offer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.inject.Inject;
import com.hubspot.singularity.SingularityAction;
import com.hubspot.singularity.config.SingularityConfiguration;
import com.hubspot.singularity.data.DisasterManager;
import com.hubspot.singularity.mesos.OfferCache;
import com.hubspot.singularity.mesos.SingularityMesosOfferScheduler;
import com.hubspot.singularity.mesos.SingularityMesosSchedulerClient;
import com.hubspot.singularity.mesos.SingularityOfferCache.CachedOffer;
import com.hubspot.singularity.mesos.SingularityOfferHolder;
import com.hubspot.singularity.mesos.SingularitySchedulerLock;

@Singleton
public class SingularitySchedulerPoller extends SingularityLeaderOnlyPoller {

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

  private final OfferCache offerCache;
  private final SingularityMesosSchedulerClient schedulerClient;
  private final SingularityMesosOfferScheduler offerScheduler;
  private final DisasterManager disasterManager;
  private final SingularitySchedulerLock lock;

  @Inject
  SingularitySchedulerPoller(SingularityMesosOfferScheduler offerScheduler, OfferCache offerCache, SingularityMesosSchedulerClient schedulerClient,
                             SingularityConfiguration configuration, SingularitySchedulerLock lock, DisasterManager disasterManager) {
    super(configuration.getCheckSchedulerEverySeconds(), TimeUnit.SECONDS, true);

    this.offerCache = offerCache;
    this.offerScheduler = offerScheduler;
    this.schedulerClient = schedulerClient;
    this.disasterManager = disasterManager;
    this.lock = lock;
  }

  @Override
  public void runActionOnPoll() {
    if (disasterManager.isDisabled(SingularityAction.RUN_SCHEDULER_POLLER)) {
      LOG.warn("Scheduler poller is disabled");
      return;
    }


    lock.runWithOffersLock(() -> {
      List<CachedOffer> cachedOffers = offerCache.checkoutOffers();
      Map<String, CachedOffer> offerIdToCachedOffer = new HashMap<>(cachedOffers.size());
      List<Offer> offers = new ArrayList<>(cachedOffers.size());

      for (CachedOffer cachedOffer : cachedOffers) {
        offerIdToCachedOffer.put(cachedOffer.getOfferId(), cachedOffer);
        offers.add(cachedOffer.getOffer());
      }

      Collection<SingularityOfferHolder> offerHolders = offerScheduler.checkOffers(offers);

      if (offerHolders.isEmpty()) {
        return;
      }

      int acceptedOffers = 0;
      int launchedTasks = 0;

      for (SingularityOfferHolder offerHolder : offerHolders) {
        List<CachedOffer> cachedOffersFromHolder = offerHolder.getOffers().stream().map((o) -> offerIdToCachedOffer.remove(o.getId().getValue())).collect(Collectors.toList());

        if (!offerHolder.getAcceptedTasks().isEmpty()) {
          List<Offer> unusedOffers = offerHolder.launchTasksAndGetUnusedOffers(schedulerClient);
          launchedTasks += offerHolder.getAcceptedTasks().size();
          acceptedOffers += cachedOffersFromHolder.size() - unusedOffers.size();

          // Return to the cache those offers which we checked out of the cache, but didn't end up using.
          List<CachedOffer> unusedCachedOffers = unusedOffers.stream().map((o) -> offerIdToCachedOffer.remove(o.getId().getValue())).collect(Collectors.toList());
          unusedCachedOffers.forEach(offerCache::returnOffer);

          // Notify the cache of the cached offers that we did use.
          cachedOffersFromHolder.removeAll(unusedCachedOffers);
          cachedOffersFromHolder.forEach(offerCache::useOffer);
        } else {
          cachedOffersFromHolder.forEach(offerCache::returnOffer);
        }
      }

      LOG.info("{} remaining offers not accounted for in offer check", offerIdToCachedOffer.size());
      offerIdToCachedOffer.values().forEach(offerCache::returnOffer);

      LOG.info("Launched {} tasks on {} cached offers (returned {})", launchedTasks, acceptedOffers, offerHolders.size() - acceptedOffers);
    }, getClass().getSimpleName());
  }
}

@@ -2991,38 +2991,35 @@
 IdToCachedOffer.
-remove
+get
 (o.getId().getVa
@@ -3568,22 +3568,19 @@
 edOffer.
-remove
+get
 (o.getId
@@ -3662,39 +3662,162 @@
 forEach(
+(cachedOffer) -%3E %7B%0A            
 offer
+IdTo
 Cache
-::returnOffer
+dOffer.remove(cachedOffer.getOfferId());%0A            offerCache.returnOffer(cachedOffer);%0A          %7D
 );%0A%0A    
@@ -3989,28 +3989,151 @@
 ach(
+(cachedOffer) -%3E %7B%0A            
 offer
+IdTo
 Cache
-::useOffer
+dOffer.remove(cachedOffer.getOfferId());%0A            offerCache.useOffer(cachedOffer);%0A          %7D
 );%0A 
@@ -4189,39 +4189,162 @@
 forEach(
+(cachedOffer) -%3E %7B%0A            
 offer
+IdTo
 Cache
-::returnOffer
+dOffer.remove(cachedOffer.getOfferId());%0A            offerCache.returnOffer(cachedOffer);%0A          %7D
 );%0A     

fix removal from offer cache map