package com.hubspot.singularity.managed;

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

import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.leader.LeaderLatch;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.inject.Inject;
import com.google.inject.Singleton;
import com.hubspot.mesos.JavaUtils;
import com.hubspot.singularity.SingularityLeaderController;
import com.hubspot.singularity.SingularityManagedCachedThreadPoolFactory;
import com.hubspot.singularity.SingularityManagedScheduledExecutorServiceFactory;
import com.hubspot.singularity.data.ExecutorIdGenerator;
import com.hubspot.singularity.mesos.SingularityMesosExecutorInfoSupport;
import com.hubspot.singularity.metrics.SingularityGraphiteReporter;
import com.hubspot.singularity.scheduler.SingularityLeaderOnlyPoller;
import com.ning.http.client.AsyncHttpClient;

import io.dropwizard.lifecycle.Managed;

/*
 * All scheduler-related startup/shutdown behavior, other managed classes not represented here:
 * SingularitySmtpSender/SmtpMailer
 * SingularityExceptionNotifierManaged
 * ExecutorIdGenerator
 */
@Singleton
public class SingularityLifecycleManaged implements Managed {
  private static final Logger LOG = LoggerFactory.getLogger(SingularityLifecycleManaged.class);

  private final SingularityManagedCachedThreadPoolFactory cachedThreadPoolFactory;
  private final SingularityManagedScheduledExecutorServiceFactory scheduledExecutorServiceFactory;
  private final AsyncHttpClient asyncHttpClient;
  private final SingularityLeaderController leaderController;
  private final LeaderLatch leaderLatch;
  private final SingularityMesosExecutorInfoSupport executorInfoSupport;
  private final SingularityGraphiteReporter graphiteReporter;
  private final ExecutorIdGenerator executorIdGenerator;
  private final Set<SingularityLeaderOnlyPoller> leaderOnlyPollers;

  private final CuratorFramework curatorFramework;
  private final AtomicBoolean started = new AtomicBoolean(false);
  private final AtomicBoolean stopped = new AtomicBoolean(false);

  @Inject
  public SingularityLifecycleManaged(SingularityManagedCachedThreadPoolFactory cachedThreadPoolFactory,
                                     SingularityManagedScheduledExecutorServiceFactory scheduledExecutorServiceFactory,
                                     AsyncHttpClient asyncHttpClient,
                                     CuratorFramework curatorFramework,
                                     SingularityLeaderController leaderController,
                                     LeaderLatch leaderLatch,
                                     SingularityMesosExecutorInfoSupport executorInfoSupport,
                                     SingularityGraphiteReporter graphiteReporter,
                                     ExecutorIdGenerator executorIdGenerator,
                                     Set<SingularityLeaderOnlyPoller> leaderOnlyPollers) {
    this.cachedThreadPoolFactory = cachedThreadPoolFactory;
    this.scheduledExecutorServiceFactory = scheduledExecutorServiceFactory;
    this.asyncHttpClient = asyncHttpClient;
    this.curatorFramework = curatorFramework;
    this.leaderController = leaderController;
    this.leaderLatch = leaderLatch;
    this.executorInfoSupport = executorInfoSupport;
    this.graphiteReporter = graphiteReporter;
    this.executorIdGenerator = executorIdGenerator;
    this.leaderOnlyPollers = leaderOnlyPollers;
  }

  @Override
  public void start() throws Exception {
    if (!started.getAndSet(true)) {
      startCurator();
      leaderLatch.start();
      leaderController.start(); // start the state poller
      graphiteReporter.start();
      executorIdGenerator.start();
      leaderOnlyPollers.forEach(SingularityLeaderOnlyPoller::start);
    } else {
      LOG.info("Already started, will not call again");
    }
  }

  @Override
  public void stop() throws Exception {
    if (!stopped.getAndSet(true)) {
      stopExecutorsAndPollers();
      stopStatePoller();
      stopDirectoryFetcher();
      stopHttpClients();
      stopLeaderLatch();
      stopCurator();
      stopGraphiteReporter();
    } else {
      LOG.info("Already stopped");
    }
  }

  private void stopDirectoryFetcher() {
    try {
      LOG.info("Stopping directory fetcher");
      executorInfoSupport.stop();
    } catch (Throwable t) {
      LOG.warn("Could not stop task directory fetcher ({})}", t.getMessage());
    }
  }

  private void stopHttpClients() {
    try {
      LOG.info("Stopping http clients");
      asyncHttpClient.close();
    } catch (Throwable t) {
      LOG.warn("Could not stop http clients ({})}", t.getMessage());
    }
  }

  private void stopExecutorsAndPollers() {
    try {
      LOG.info("Stopping pollers and executors");
      leaderOnlyPollers.forEach(SingularityLeaderOnlyPoller::stop); // mostly no-op, custom stop actions per poller
      cachedThreadPoolFactory.stop();
      scheduledExecutorServiceFactory.stop();
    } catch (Throwable t) {
      LOG.warn("Could not stop scheduled executors ({})}", t.getMessage());
    }
  }

  private void stopStatePoller() {
    try {
      LOG.info("Stopping state poller");
      leaderController.stop();
    } catch (Throwable t) {
      LOG.warn("Could not stop state poller ({})}", t.getMessage());
    }
  }

  private void stopGraphiteReporter() {
    try {
      LOG.info("Stopping graphite reporter");
      graphiteReporter.stop();
    } catch (Throwable t) {
      LOG.warn("Could not stop graphite reporter ({})}", t.getMessage());
    }
  }

  private void stopLeaderLatch() {
    try {
      LOG.info("Stopping leader latch");
      leaderLatch.close();
    } catch (Throwable t) {
      LOG.warn("Could not stop leader latch ({})}", t.getMessage());
    }
  }

  private void startCurator() {
    curatorFramework.start();

    final long start = System.currentTimeMillis();

    try {
      checkState(curatorFramework.getZookeeperClient().blockUntilConnectedOrTimedOut(), "did not connect to zookeeper");
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
    }

    LOG.info("Connected to ZK after {}", JavaUtils.duration(start));
  }

  private void stopCurator() {
    try {
      LOG.info("Stopping curator");
      curatorFramework.close();
    } catch (Throwable t) {
      LOG.warn("Could not close curator ({})", t.getMessage());
    }
  }
}

@@ -4063,34 +4063,23 @@
     stop
-ExecutorsAnd
+New
 Poll
-er
 s();%0A   
@@ -4081,33 +4081,38 @@
 ;%0A      stop
-StatePoll
+DirectoryFetch
 er();%0A      
@@ -4119,48 +4119,84 @@
 stop
-DirectoryFetcher();%0A      stopHttpClient
+HttpClients();%0A      stopStatePollerAndMesosConnection();%0A      stopExecutor
 s();
@@ -4827,26 +4827,15 @@
 stop
-ExecutorsAnd
+New
 Poll
-er
 s() 
@@ -4844,66 +4844,62 @@
     
-try %7B%0A      LOG.info(%22Stopping pollers and executors
+LOG.info(%22Marking leader only pollers for shutdown
 %22);%0A
-  
     
@@ -4963,56 +4963,106 @@
 op);
- // mostly no-op, 
+%0A  %7D%0A%0A  private void stopExe
 cu
-s
 to
-m stop actions per poller
+rs() %7B%0A    try %7B%0A      LOG.info(%22Stopping pollers and executors%22);
 %0A   
@@ -5287,16 +5287,34 @@
 tePoller
+AndMesosConnection
 () %7B%0A   

reorder shutdown