package com.hubspot.blazar.zookeeper;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

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

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.leader.LeaderLatchListener;
import org.apache.curator.utils.ZKPaths;
import org.apache.zookeeper.KeeperException.NoNodeException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.codahale.metrics.Gauge;
import com.codahale.metrics.MetricRegistry;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.Sets;
import com.google.inject.name.Named;
import io.dropwizard.lifecycle.Managed;

@Singleton
public class QueueProcessor implements LeaderLatchListener, Managed {
  private static final Logger LOG = LoggerFactory.getLogger(QueueProcessor.class);

  private final ScheduledExecutorService executorService;
  private final CuratorFramework curatorFramework;
  private final ZooKeeperEventBus eventBus;
  private final ObjectMapper mapper;
  private MetricRegistry metricRegistry;
  private final Set<Object> erroredItems;
  private final AtomicBoolean running;
  private final AtomicBoolean leader;

  @Inject
  public QueueProcessor(@Named("QueueProcessor") ScheduledExecutorService executorService,
                        CuratorFramework curatorFramework,
                        ZooKeeperEventBus eventBus,
                        ObjectMapper mapper,
                        MetricRegistry metricRegistry,
                        Set<Object> erroredItems) {
    this.executorService = executorService;
    this.curatorFramework = curatorFramework;
    this.eventBus = eventBus;
    this.mapper = mapper;
    this.metricRegistry = metricRegistry;
    this.erroredItems = erroredItems;

    this.running = new AtomicBoolean();
    this.leader = new AtomicBoolean();
  }

  @Override
  public void start() {
    running.set(true);
    executorService.scheduleAtFixedRate(new QueueManager(), 0, 10, TimeUnit.SECONDS);

  }

  @Override
  public void stop() {
    running.set(false);
  }

  @Override
  public void isLeader() {
    LOG.info("Now the leader, starting queue processing");
    leader.set(true);
    registerGauges();
  }

  @Override
  public void notLeader() {
    LOG.info("Not the leader, stopping queue processing");
    deRegisterGauges();
    leader.set(false);
  }

  private class QueueManager implements Runnable {
    private final Map<String, Future<?>> activeQueues = new ConcurrentHashMap<>();

    @Override
    public void run() {
      try {
        if (running.get() && leader.get()) {
          Set<String> queues = new HashSet<>(curatorFramework.getChildren().forPath("/queues"));
          LOG.debug("Found queues to process: {}", queues);

          for (String addedQueue : Sets.difference(queues, activeQueues.keySet())) {
            Runnable processQueueRunnable = new ProcessQueueRunnable(addedQueue);
            Future<?> future = executorService.scheduleAtFixedRate(processQueueRunnable, 0, 1, TimeUnit.SECONDS);
            activeQueues.put(addedQueue, future);
          }

          for (String removedQueue : Sets.difference(activeQueues.keySet(), queues)) {
            activeQueues.remove(removedQueue).cancel(false);
          }
        } else {
          LOG.debug("Not checking for queues to process");
          Set<String> queues = new HashSet<>(activeQueues.keySet());
          for (String queue : queues) {
            activeQueues.remove(queue).cancel(false);
          }
        }
      } catch (NoNodeException e) {
        LOG.warn("No node found for path: {}", e.getPath());
      } catch (Throwable t) {
        LOG.error("Error checking for queues to process", t);
      }
    }
  }

  private class ProcessQueueRunnable implements Runnable {
    private final String queue;

    private ProcessQueueRunnable(String queue) {
      this.queue = queue;
    }

    @Override
    public void run() {
      try {
        List<String> items = curatorFramework.getChildren().forPath(ZKPaths.makePath("queues", queue));
        LOG.debug("Found {} items in queue: {}", items.size(), queue);

        for (String item : sort(items)) {
          String path = ZKPaths.makePath("queues", queue, item);

          if (process(path)) {
            LOG.debug("Deleting successfully processed queue item: {}", path);
            curatorFramework.delete().guaranteed().forPath(path);
          } else {
            LOG.info("Failed to process queue item, will retry: {}", path);
          }
        }
      } catch (Throwable t) {
        LOG.error("Error processing queue: {}", queue, t);
      }
    }

    private boolean process(String path) throws Exception {
      try {
        byte[] data = curatorFramework.getData().forPath(path);
        QueueItem item = mapper.readValue(data, QueueItem.class);
        long age = System.currentTimeMillis() - item.getTimestamp();
        if (age < 100) {
          Thread.sleep(100 - age);
        }
        return process(item.getItem());
      } catch (NoNodeException e) {
        LOG.warn("No node found for path: {}", e.getPath());
        return false;
      } catch (Exception e) {
        LOG.error("Error processing queue item: {}", path, e);
        return false;
      }
    }

    private boolean process(Object event) {
      eventBus.dispatch(event);

      return !erroredItems.remove(event);
    }

    private List<String> sort(Collection<String> unsorted) {
      List<String> sorted = new ArrayList<>(unsorted);
      Collections.sort(sorted);
      return sorted;
    }
  }

  private String zkPathToMetricName(String path) {
    return getClass().getName() + path.replace('/', '.');
  }

  private Gauge<Number> makeQueueGauge(final String path){
    return new Gauge<Number>() {
      @Override
      public Number getValue() {
        try {
          return curatorFramework.getChildren().forPath(path).size();
        } catch (Exception e) {
          return 0;
        }
      }
    };
  }

  private void deRegisterGauges() {
    for (Map.Entry<String, Gauge> entry : metricRegistry.getGauges().entrySet()) {
      if (entry.getKey().contains(getClass().getName())) {
        metricRegistry.remove(entry.getKey());
      }
    }
  }

  private void registerGauges() {
    try {
      for (String queue : curatorFramework.getChildren().forPath("/queues")) {
        String queuePath = ZKPaths.makePath("/queues", queue);
        metricRegistry.register(zkPathToMetricName(queuePath), makeQueueGauge(queuePath));
      }
    } catch (Exception e) {
      LOG.error("Could not register the metric gauges", e);
    }
  }
}

@@ -6045,16 +6045,31 @@
 /', '.')
+ + %22.%22 + %22size%22
 ;%0A  %7D%0A%0A 
@@ -6648,24 +6648,111 @@
 rGauges() %7B%0A
+    metricRegistry.register(zkPathToMetricName(%22/queues%22), makeQueueGauge(%22/queues%22));%0A
     try %7B%0A  

add size for queue, add number of queues gauge