/*
 * Copyright 2014 Cask, Inc.
 *
 * 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 co.cask.cdap.metrics.process;

import co.cask.cdap.common.metrics.MetricsScope;
import co.cask.cdap.data2.OperationException;
import co.cask.cdap.metrics.MetricsConstants.ConfigKeys;
import co.cask.cdap.metrics.data.MetricsTableFactory;
import com.google.common.collect.Lists;
import com.google.common.util.concurrent.AbstractIdleService;
import com.google.inject.Inject;
import com.google.inject.assistedinject.Assisted;
import com.google.inject.name.Named;
import org.apache.twill.common.Cancellable;
import org.apache.twill.kafka.client.KafkaClientService;
import org.apache.twill.kafka.client.KafkaConsumer;
import org.apache.twill.kafka.client.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Set;
import java.util.concurrent.TimeUnit;

/**
 * Process metrics by consuming metrics being published to kafka.
 */
public final class KafkaMetricsProcessorService extends AbstractIdleService {

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

  private final KafkaClientService kafkaClient;
  private final MessageCallbackFactory callbackFactory;
  private final String topicPrefix;
  private final Set<Integer> partitions;
  private Cancellable unsubscribe;
  private final MetricsTableFactory metricsTableFactory;

  private volatile boolean stopping = false;

  private KafkaConsumerMetaTable metaTable;

  @Inject
  public KafkaMetricsProcessorService(KafkaClientService kafkaClient,
                                      MetricsTableFactory metricsTableFactory,
                                      MessageCallbackFactory callbackFactory,
                                      @Named(ConfigKeys.KAFKA_TOPIC_PREFIX) String topicPrefix,
                                      @Assisted Set<Integer> partitions) {
    this.kafkaClient = kafkaClient;
    this.callbackFactory = callbackFactory;
    this.topicPrefix = topicPrefix;
    this.partitions = partitions;
    this.metricsTableFactory = metricsTableFactory;
  }

  @Override
  protected void startUp() {
    LOG.info("Starting Metrics Processing for partitions {}.", partitions);
    subscribe();
    LOG.info("Metrics Processing Service started for partition {}.", partitions);
  }

  @Override
  protected void shutDown() {
    stopping = true;
    LOG.info("Stopping Metrics Processing Service.");

    // Cancel kafka subscriptions
    if (unsubscribe != null) {
      unsubscribe.cancel();
    }
    LOG.info("Metrics Processing Service stopped.");
  }

  private KafkaConsumerMetaTable getMetaTable() {
    while (!stopping && metaTable == null) {
      try {
        metaTable = metricsTableFactory.createKafkaConsumerMeta("default");
      } catch (Exception e) {
        LOG.warn("Cannot access kafka consumer metaTable, will retry in 1 sec");
        try {
          TimeUnit.SECONDS.sleep(1);
        } catch (InterruptedException ie) {
          Thread.currentThread().interrupt();
          break;
        }
      }
    }

    return metaTable;
  }

  private void subscribe() {
    List<Cancellable> cancels = Lists.newArrayList();
    for (MetricsScope scope : MetricsScope.values()) {
      // Assuming there is only one process that pulling in all metrics.
      KafkaConsumer.Preparer preparer = kafkaClient.getConsumer().prepare();

      String topic = topicPrefix + "." + scope.name().toLowerCase();
      for (int i : partitions) {
        long offset = getOffset(topic, i);
        if (offset >= 0) {
          preparer.add(topic, i, offset);
        } else {
          preparer.addFromBeginning(topic, i);
        }
      }

      cancels.add(preparer.consume(callbackFactory.create(getMetaTable(), scope)));
      LOG.info("Consumer created for topic {}, partitions {}", topic, partitions);
    }
    unsubscribe = createCancelAll(cancels);
  }

  private long getOffset(String topic, int partition) {
    LOG.info("Retrieve offset for topic: {}, partition: {}", topic, partition);
    try {
      long offset = getMetaTable().get(new TopicPartition(topic, partition));
      LOG.info("Offset for topic: {}, partition: {} is {}", topic, partition, offset);
      return offset;
    } catch (OperationException e) {
      LOG.error("Failed to get offset from meta table. Defaulting to beginning. {}", e.getMessage(), e);
    }
    return -1L;
  }

  private Cancellable createCancelAll(final Iterable<? extends Cancellable> cancels) {
    return new Cancellable() {
      @Override
      public void cancel() {
        for (Cancellable cancel : cancels) {
          cancel.cancel();
        }
      }
    };
  }
}

@@ -914,28 +914,39 @@
 ent.Abstract
-Idle
+ExecutionThread
 Service;%0Aimp
@@ -1552,12 +1552,23 @@
 ract
-Idle
+ExecutionThread
 Serv
@@ -2680,20 +2680,115 @@
 ted 
-void startUp
+String getServiceName() %7B%0A    return this.getClass().getSimpleName();%0A  %7D%0A%0A  @Override%0A  protected void run
 () %7B
@@ -2943,16 +2943,17 @@
 artition
+s
  %7B%7D.%22, p
@@ -2963,16 +2963,241 @@
 itions);
+%0A%0A    while (isRunning()) %7B%0A      try %7B%0A        TimeUnit.SECONDS.sleep(1);%0A      %7D catch (InterruptedException e) %7B%0A        // It's triggered by stop%0A        Thread.currentThread().interrupt();%0A        continue;%0A      %7D%0A    %7D
 %0A  %7D%0A%0A  
@@ -3227,13 +3227,20 @@
 oid 
-shutD
+triggerShutd
 own(
@@ -3260,24 +3260,100 @@
 ing = true;%0A
+    super.triggerShutdown();%0A  %7D%0A%0A  @Override%0A  protected void shutDown() %7B%0A
     LOG.info

fix shutdown of KafkaMetricsProcessorService