/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 org.apache.gobblin.source.extractor.extract.kafka;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Pattern;

import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.codahale.metrics.Timer;
import com.google.common.base.Function;
import com.google.common.base.Joiner;
import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import com.google.common.base.Splitter;
import com.google.common.base.Stopwatch;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.typesafe.config.Config;

import lombok.Getter;
import lombok.Setter;

import org.apache.gobblin.configuration.ConfigurationKeys;
import org.apache.gobblin.configuration.SourceState;
import org.apache.gobblin.configuration.State;
import org.apache.gobblin.configuration.WorkUnitState;
import org.apache.gobblin.dataset.DatasetConstants;
import org.apache.gobblin.dataset.DatasetDescriptor;
import org.apache.gobblin.instrumented.Instrumented;
import org.apache.gobblin.kafka.client.GobblinKafkaConsumerClient;
import org.apache.gobblin.kafka.client.GobblinKafkaConsumerClient.GobblinKafkaConsumerClientFactory;
import org.apache.gobblin.metrics.MetricContext;
import org.apache.gobblin.metrics.event.lineage.LineageInfo;
import org.apache.gobblin.source.extractor.extract.EventBasedSource;
import org.apache.gobblin.source.extractor.extract.kafka.workunit.packer.KafkaTopicGroupingWorkUnitPacker;
import org.apache.gobblin.source.extractor.extract.kafka.workunit.packer.KafkaWorkUnitPacker;
import org.apache.gobblin.source.extractor.extract.kafka.validator.TopicValidators;
import org.apache.gobblin.source.extractor.limiter.LimiterConfigurationKeys;
import org.apache.gobblin.source.workunit.Extract;
import org.apache.gobblin.source.workunit.MultiWorkUnit;
import org.apache.gobblin.source.workunit.WorkUnit;
import org.apache.gobblin.util.ClassAliasResolver;
import org.apache.gobblin.util.ConfigUtils;
import org.apache.gobblin.util.DatasetFilterUtils;
import org.apache.gobblin.util.ExecutorsUtils;
import org.apache.gobblin.util.dataset.DatasetUtils;

import static java.util.stream.Collectors.toSet;


/**
 * A {@link org.apache.gobblin.source.Source} implementation for Kafka source.
 *
 * @author Ziyang Liu
 */
public abstract class KafkaSource<S, D> extends EventBasedSource<S, D> {

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

  public static final String TOPIC_BLACKLIST = "topic.blacklist";
  public static final String TOPIC_WHITELIST = "topic.whitelist";
  public static final String LATEST_OFFSET = "latest";
  public static final String EARLIEST_OFFSET = "earliest";
  public static final String NEAREST_OFFSET = "nearest";
  public static final String OFFSET_LOOKBACK = "offset_lookback";
  public static final String BOOTSTRAP_WITH_OFFSET = "bootstrap.with.offset";
  public static final String KAFKA_OFFSET_LOOKBACK = "kafka.offset.lookback";
  public static final String DEFAULT_BOOTSTRAP_WITH_OFFSET = LATEST_OFFSET;
  public static final String TOPICS_MOVE_TO_LATEST_OFFSET = "topics.move.to.latest.offset";
  public static final String RESET_ON_OFFSET_OUT_OF_RANGE = "reset.on.offset.out.of.range";
  public static final String DEFAULT_RESET_ON_OFFSET_OUT_OF_RANGE = NEAREST_OFFSET;
  public static final String TOPIC_NAME = "topic.name";
  public static final String PARTITION_ID = "partition.id";
  public static final String LEADER_ID = "leader.id";
  public static final String LEADER_HOSTANDPORT = "leader.hostandport";
  public static final Extract.TableType DEFAULT_TABLE_TYPE = Extract.TableType.APPEND_ONLY;
  public static final String DEFAULT_NAMESPACE_NAME = "KAFKA";
  public static final String ALL_TOPICS = "all";
  //A workunit property that contains the number of topic partitions for a given topic. Useful for
  //workunit size estimation to assign weights to a given topic partition.
  public static final String NUM_TOPIC_PARTITIONS = "numTopicPartitions";
  public static final String AVG_RECORD_MILLIS = "avg.record.millis";
  public static final String START_FETCH_EPOCH_TIME = "startFetchEpochTime";
  public static final String STOP_FETCH_EPOCH_TIME = "stopFetchEpochTime";
  public static final String PREVIOUS_START_FETCH_EPOCH_TIME = "previousStartFetchEpochTime";
  public static final String PREVIOUS_STOP_FETCH_EPOCH_TIME = "previousStopFetchEpochTime";
  public static final String PREVIOUS_LOW_WATERMARK = "previousLowWatermark";
  public static final String PREVIOUS_HIGH_WATERMARK = "previousHighWatermark";
  public static final String PREVIOUS_LATEST_OFFSET = "previousLatestOffset";
  public static final String OFFSET_FETCH_EPOCH_TIME = "offsetFetchEpochTime";
  public static final String PREVIOUS_OFFSET_FETCH_EPOCH_TIME = "previousOffsetFetchEpochTime";
  public static final String ALLOW_PERIOD_IN_TOPIC_NAME = "gobblin.kafka.allowPeriodInTopicName";
  public static final String GOBBLIN_KAFKA_CONSUMER_CLIENT_FACTORY_CLASS = "gobblin.kafka.consumerClient.class";
  public static final String GOBBLIN_KAFKA_EXTRACT_ALLOW_TABLE_TYPE_NAMESPACE_CUSTOMIZATION =
      "gobblin.kafka.extract.allowTableTypeAndNamspaceCustomization";
  public static final String DEFAULT_GOBBLIN_KAFKA_CONSUMER_CLIENT_FACTORY_CLASS =
      "org.apache.gobblin.kafka.client.Kafka08ConsumerClient$Factory";
  public static final String GOBBLIN_KAFKA_SHOULD_ENABLE_DATASET_STATESTORE =
      "gobblin.kafka.shouldEnableDatasetStateStore";
  public static final boolean DEFAULT_GOBBLIN_KAFKA_SHOULD_ENABLE_DATASET_STATESTORE = false;
  public static final String OFFSET_FETCH_TIMER = "offsetFetchTimer";
  public static final String RECORD_LEVEL_SLA_MINUTES_KEY = "gobblin.kafka.recordLevelSlaMinutes";
  public static final String MAX_POSSIBLE_OBSERVED_LATENCY_IN_HOURS = "gobblin.kafka.maxobservedLatencyInHours";
  public static final Integer DEFAULT_MAX_POSSIBLE_OBSERVED_LATENCY_IN_HOURS = 24;
  public static final String OBSERVED_LATENCY_PRECISION = "gobblin.kafka.observedLatencyPrecision";
  public static final Integer DEFAULT_OBSERVED_LATENCY_PRECISION = 3;
  public static final String OBSERVED_LATENCY_MEASUREMENT_ENABLED = "gobblin.kafka.observedLatencyMeasurementEnabled";
  public static final Boolean DEFAULT_OBSERVED_LATENCY_MEASUREMENT_ENABLED = false;
  public static final String RECORD_CREATION_TIMESTAMP_FIELD = "gobblin.kafka.recordCreationTimestampField";
  public static final String RECORD_CREATION_TIMESTAMP_UNIT = "gobblin.kafka.recordCreationTimestampUnit";

  private final Set<String> moveToLatestTopics = Sets.newTreeSet(String.CASE_INSENSITIVE_ORDER);
  private final Map<KafkaPartition, Long> previousOffsets = Maps.newConcurrentMap();
  private final Map<KafkaPartition, Long> previousLowWatermarks = Maps.newConcurrentMap();
  private final Map<KafkaPartition, Long> previousExpectedHighWatermarks = Maps.newConcurrentMap();
  private final Map<KafkaPartition, Long> previousOffsetFetchEpochTimes = Maps.newConcurrentMap();
  private final Map<KafkaPartition, Long> previousStartFetchEpochTimes = Maps.newConcurrentMap();
  private final Map<KafkaPartition, Long> previousStopFetchEpochTimes = Maps.newConcurrentMap();

  private final Set<KafkaPartition> partitionsToBeProcessed = Sets.newConcurrentHashSet();

  private final AtomicInteger failToGetOffsetCount = new AtomicInteger(0);
  private final AtomicInteger offsetTooEarlyCount = new AtomicInteger(0);
  private final AtomicInteger offsetTooLateCount = new AtomicInteger(0);

  // sharing the kafka consumer may result in contention, so support thread local consumers
  protected final ConcurrentLinkedQueue<GobblinKafkaConsumerClient> kafkaConsumerClientPool = new ConcurrentLinkedQueue();
  protected static final ThreadLocal<GobblinKafkaConsumerClient> kafkaConsumerClient =
          new ThreadLocal<GobblinKafkaConsumerClient>();
  private GobblinKafkaConsumerClient sharedKafkaConsumerClient = null;
  private final ClassAliasResolver<GobblinKafkaConsumerClientFactory> kafkaConsumerClientResolver =
      new ClassAliasResolver<>(GobblinKafkaConsumerClientFactory.class);

  private volatile boolean doneGettingAllPreviousOffsets = false;
  private Extract.TableType tableType;
  private String extractNamespace;
  private boolean isFullExtract;
  private String kafkaBrokers;
  private boolean shouldEnableDatasetStateStore;
  private AtomicBoolean isDatasetStateEnabled = new AtomicBoolean(false);
  private Set<String> topicsToProcess;

  private MetricContext metricContext;

  protected Optional<LineageInfo> lineageInfo;

  private List<String> getLimiterExtractorReportKeys() {
    List<String> keyNames = new ArrayList<>();
    keyNames.add(KafkaSource.TOPIC_NAME);
    keyNames.add(KafkaSource.PARTITION_ID);
    return keyNames;
  }

  private void setLimiterReportKeyListToWorkUnits(List<WorkUnit> workUnits, List<String> keyNameList) {
    if (keyNameList.isEmpty()) {
      return;
    }
    String keyList = Joiner.on(',').join(keyNameList.iterator());
    for (WorkUnit workUnit : workUnits) {
      workUnit.setProp(LimiterConfigurationKeys.LIMITER_REPORT_KEY_LIST, keyList);
    }
  }

  @Override
  public List<WorkUnit> getWorkunits(SourceState state) {
    return this.getWorkunitsForFilteredPartitions(state, Optional.absent(), Optional.absent());
  }

  /**
   * Compute Workunits for Kafka Topics. If filteredTopicPartition present, respect this map and only compute the provided
   * topics and filtered partitions. If not, use state to discover Kafka topics and all available partitions.
   *
   * @param filteredTopicPartition optional parameter to determine if only filtered topic-partitions are needed.
   * @param minContainer give an option to specify a minimum container count. Please be advised that how it being used is
   *                     determined by the implementation of concrete {@link KafkaWorkUnitPacker} class.
   *
   * TODO: Utilize the minContainer in {@link KafkaTopicGroupingWorkUnitPacker#pack(Map, int)}, as the numContainers variable
   *                     is not used currently.
   */
  public List<WorkUnit> getWorkunitsForFilteredPartitions(SourceState state,
      Optional<Map<String, List<Integer>>> filteredTopicPartition, Optional<Integer> minContainer) {
    this.metricContext = Instrumented.getMetricContext(state, KafkaSource.class);
    this.lineageInfo = LineageInfo.getLineageInfo(state.getBroker());

    Map<String, List<Integer>> filteredTopicPartitionMap = filteredTopicPartition.or(new HashMap<>());
    Map<String, List<WorkUnit>> kafkaTopicWorkunitMap = Maps.newConcurrentMap();
    if (state.getPropAsBoolean(KafkaSource.GOBBLIN_KAFKA_EXTRACT_ALLOW_TABLE_TYPE_NAMESPACE_CUSTOMIZATION)) {
      String tableTypeStr =
          state.getProp(ConfigurationKeys.EXTRACT_TABLE_TYPE_KEY, KafkaSource.DEFAULT_TABLE_TYPE.toString());
      tableType = Extract.TableType.valueOf(tableTypeStr);
      extractNamespace =
          state.getProp(ConfigurationKeys.EXTRACT_NAMESPACE_NAME_KEY, KafkaSource.DEFAULT_NAMESPACE_NAME);
    } else {
      // To be compatible, reject table type and namespace configuration keys as previous implementation
      tableType = KafkaSource.DEFAULT_TABLE_TYPE;
      extractNamespace = KafkaSource.DEFAULT_NAMESPACE_NAME;
    }
    isFullExtract = state.getPropAsBoolean(ConfigurationKeys.EXTRACT_IS_FULL_KEY);
    kafkaBrokers = state.getProp(ConfigurationKeys.KAFKA_BROKERS, "");
    this.shouldEnableDatasetStateStore = state.getPropAsBoolean(GOBBLIN_KAFKA_SHOULD_ENABLE_DATASET_STATESTORE,
        DEFAULT_GOBBLIN_KAFKA_SHOULD_ENABLE_DATASET_STATESTORE);

    try {
      Config config = ConfigUtils.propertiesToConfig(state.getProperties());
      GobblinKafkaConsumerClientFactory kafkaConsumerClientFactory = kafkaConsumerClientResolver
          .resolveClass(state.getProp(
              GOBBLIN_KAFKA_CONSUMER_CLIENT_FACTORY_CLASS,
              DEFAULT_GOBBLIN_KAFKA_CONSUMER_CLIENT_FACTORY_CLASS)).newInstance();

      this.kafkaConsumerClient.set(kafkaConsumerClientFactory.create(config));

      Collection<KafkaTopic> topics;
      if(filteredTopicPartition.isPresent()) {
        if(filteredTopicPartition.get().isEmpty()) {
          // return an empty list as filteredTopicPartition is present but contains no valid entry
          return new ArrayList<>();
        } else {
          // If filteredTopicPartition present, use it to construct the whitelist pattern while leave blacklist empty
          topics = this.kafkaConsumerClient.get()
              .getFilteredTopics(Collections.emptyList(),
                  filteredTopicPartitionMap.keySet().stream().map(Pattern::compile).collect(Collectors.toList()));
        }
      } else {
        // get topics based on job level config
        topics = getValidTopics(getFilteredTopics(state), state);
      }
      this.topicsToProcess = topics.stream().map(KafkaTopic::getName).collect(toSet());

      Map<String, State> topicSpecificStateMap =
          DatasetUtils.getDatasetSpecificProps(Iterables.transform(topics, new Function<KafkaTopic, String>() {

            @Override
            public String apply(KafkaTopic topic) {
              return topic.getName();
            }
          }), state);

      int numOfThreads = state.getPropAsInt(ConfigurationKeys.KAFKA_SOURCE_WORK_UNITS_CREATION_THREADS,
          ConfigurationKeys.KAFKA_SOURCE_WORK_UNITS_CREATION_DEFAULT_THREAD_COUNT);
      // No need to allocate more thread than the topic size, but minimum should 1
      numOfThreads = Math.max(Math.min(numOfThreads, topics.size()), 1);
      ExecutorService threadPool =
          Executors.newFixedThreadPool(numOfThreads, ExecutorsUtils.newThreadFactory(Optional.of(LOG)));

      if (state.getPropAsBoolean(ConfigurationKeys.KAFKA_SOURCE_SHARE_CONSUMER_CLIENT,
          ConfigurationKeys.DEFAULT_KAFKA_SOURCE_SHARE_CONSUMER_CLIENT)) {
        this.sharedKafkaConsumerClient = this.kafkaConsumerClient.get();
      } else {
        // preallocate one client per thread
        populateClientPool(numOfThreads, kafkaConsumerClientFactory, config);
      }

      Stopwatch createWorkUnitStopwatch = Stopwatch.createStarted();

      for (KafkaTopic topic : topics) {
        LOG.info("Discovered topic {} with {} number of partitions", topic.getName(), topic.getPartitions().size());
        if (topic.getTopicSpecificState().isPresent()) {
          topicSpecificStateMap.computeIfAbsent(topic.getName(), k -> new State())
              .addAllIfNotExist(topic.getTopicSpecificState().get());
        }
        Optional<Set<Integer>> partitionIDSet = Optional.absent();
        if(filteredTopicPartition.isPresent()) {
          List<Integer> list = java.util.Optional.ofNullable(filteredTopicPartitionMap.get(topic.getName()))
              .orElse(new ArrayList<>());
          partitionIDSet = Optional.of(new HashSet<>(list));
          LOG.info("Compute the workunit for topic {} with num of filtered partitions: {}",
              topic.getName(), list.size());
        }

        threadPool.submit(
            new WorkUnitCreator(topic, state, Optional.fromNullable(topicSpecificStateMap.get(topic.getName())),
                kafkaTopicWorkunitMap, partitionIDSet));
      }

      ExecutorsUtils.shutdownExecutorService(threadPool, Optional.of(LOG), 1L, TimeUnit.HOURS);
      LOG.info(String.format("Created workunits for %d topics in %d seconds", kafkaTopicWorkunitMap.size(),
          createWorkUnitStopwatch.elapsed(TimeUnit.SECONDS)));

      // Create empty WorkUnits for skipped partitions (i.e., partitions that have previous offsets,
      // but aren't processed). When filteredTopicPartition present, only filtered topic-partitions are needed so skip this call
      if(!filteredTopicPartition.isPresent()) {
        createEmptyWorkUnitsForSkippedPartitions(kafkaTopicWorkunitMap, topicSpecificStateMap, state);
      }

      KafkaWorkUnitPacker kafkaWorkUnitPacker = KafkaWorkUnitPacker.getInstance(this, state, Optional.of(this.metricContext));
      int numOfMultiWorkunits = minContainer.or(1);
      if(state.contains(ConfigurationKeys.MR_JOB_MAX_MAPPERS_KEY)) {
        numOfMultiWorkunits = Math.max(numOfMultiWorkunits,
            calculateNumMappersForPacker(state, kafkaWorkUnitPacker, kafkaTopicWorkunitMap));
      }

      addTopicSpecificPropsToWorkUnits(kafkaTopicWorkunitMap, topicSpecificStateMap);
      List<WorkUnit> workUnitList = kafkaWorkUnitPacker.pack(kafkaTopicWorkunitMap, numOfMultiWorkunits);
      setLimiterReportKeyListToWorkUnits(workUnitList, getLimiterExtractorReportKeys());
      return workUnitList;
    } catch (InstantiationException | IllegalAccessException | ClassNotFoundException e) {
      throw new RuntimeException("Checked exception caught", e);
    } catch (Throwable t) {
      throw new RuntimeException("Unexpected throwable caught, ", t);
    } finally {
      try {
        GobblinKafkaConsumerClient consumerClient = this.kafkaConsumerClient.get();
        if (consumerClient != null) {
          consumerClient.close();
        }
        // cleanup clients from pool
        for (GobblinKafkaConsumerClient client: kafkaConsumerClientPool) {
          client.close();
        }
      } catch (Throwable t) {
        //Swallow any exceptions in the finally{..} block to allow potential exceptions from the main try{..} block to be
        //propagated
        LOG.error("Exception {} encountered closing GobblinKafkaConsumerClient ", t);
      }
    }

  }

  protected void populateClientPool(int count,
      GobblinKafkaConsumerClientFactory kafkaConsumerClientFactory,
      Config config) {
    // Clear the pool first as clients within may already be close
    kafkaConsumerClientPool.clear();
    for (int i = 0; i < count; i++) {
      kafkaConsumerClientPool.offer(kafkaConsumerClientFactory.create(config));
    }
  }

  private void addTopicSpecificPropsToWorkUnits(Map<String, List<WorkUnit>> workUnits, Map<String, State> topicSpecificStateMap) {
    for (List<WorkUnit> workUnitList : workUnits.values()) {
      for (WorkUnit workUnit : workUnitList) {
        addTopicSpecificPropsToWorkUnit(workUnit, topicSpecificStateMap);
      }
    }
  }

  private void addTopicSpecificPropsToWorkUnit(WorkUnit workUnit, Map<String, State> topicSpecificStateMap) {
    if (workUnit instanceof MultiWorkUnit) {
      for (WorkUnit wu : ((MultiWorkUnit) workUnit).getWorkUnits()) {
        addTopicSpecificPropsToWorkUnit(wu, topicSpecificStateMap);
      }
    } else if (!workUnit.contains(TOPIC_NAME)) {
      return;
    } else {
      addDatasetUrnOptionally(workUnit);
      if (topicSpecificStateMap == null) {
        return;
      } else if (!topicSpecificStateMap.containsKey(workUnit.getProp(TOPIC_NAME))) {
        return;
      } else {
        workUnit.addAll(topicSpecificStateMap.get(workUnit.getProp(TOPIC_NAME)));
      }
    }
  }

  private void addDatasetUrnOptionally(WorkUnit workUnit) {
    if (!this.shouldEnableDatasetStateStore) {
      return;
    }
    workUnit.setProp(ConfigurationKeys.DATASET_URN_KEY, workUnit.getProp(TOPIC_NAME));
  }

  private void createEmptyWorkUnitsForSkippedPartitions(Map<String, List<WorkUnit>> workUnits,
      Map<String, State> topicSpecificStateMap, SourceState state) {

    // in case the previous offset not been set
    getAllPreviousOffsetState(state);

    // For each partition that has a previous offset, create an empty WorkUnit for it if
    // it is not in this.partitionsToBeProcessed.
    for (Map.Entry<KafkaPartition, Long> entry : this.previousOffsets.entrySet()) {
      KafkaPartition partition = entry.getKey();
      if (!this.partitionsToBeProcessed.contains(partition)) {
        String topicName = partition.getTopicName();
        if (!this.isDatasetStateEnabled.get() || this.topicsToProcess.contains(topicName)) {
          long previousOffset = entry.getValue();
          WorkUnit emptyWorkUnit = createEmptyWorkUnit(partition, previousOffset,
              this.previousOffsetFetchEpochTimes.get(partition),
              Optional.fromNullable(topicSpecificStateMap.get(partition.getTopicName())));

          if (workUnits.containsKey(topicName)) {
            workUnits.get(topicName).add(emptyWorkUnit);
          } else {
            workUnits.put(topicName, Lists.newArrayList(emptyWorkUnit));
          }
        }
      }
    }
  }

  //determine the number of mappers/containers for workunit packer
  private int calculateNumMappersForPacker(SourceState state,
      KafkaWorkUnitPacker kafkaWorkUnitPacker, Map<String, List<WorkUnit>> workUnits) {
    int maxMapperNum =
        state.getPropAsInt(ConfigurationKeys.MR_JOB_MAX_MAPPERS_KEY, ConfigurationKeys.DEFAULT_MR_JOB_MAX_MAPPERS);
    int numContainers = maxMapperNum;
    if(state.contains(ConfigurationKeys.MR_TARGET_MAPPER_SIZE)) {
      double totalEstDataSize = kafkaWorkUnitPacker.setWorkUnitEstSizes(workUnits);
      LOG.info(String.format("The total estimated data size is %.2f", totalEstDataSize));
      double targetMapperSize = state.getPropAsDouble(ConfigurationKeys.MR_TARGET_MAPPER_SIZE);
      numContainers = (int) (totalEstDataSize / targetMapperSize) + 1;
      numContainers = Math.min(numContainers, maxMapperNum);
    }
    return numContainers;
  }

  /*
   * This function need to be thread safe since it is called in the Runnable
   */
  public List<WorkUnit> getWorkUnitsForTopic(KafkaTopic topic, SourceState state,
      Optional<State> topicSpecificState, Optional<Set<Integer>> filteredPartitions) {
    Timer.Context context = this.metricContext.timer("isTopicQualifiedTimer").time();
    boolean topicQualified = isTopicQualified(topic);
    context.close();

    final List<WorkUnit> workUnits = Lists.newArrayList();
    final List<KafkaPartition> topicPartitions = topic.getPartitions();
    Map<KafkaPartition, WorkUnit> workUnitMap;

    if (filteredPartitions.isPresent()) {
      LOG.info("Filtered partitions for topic {} are {}", topic.getName(), filteredPartitions.get());
      final List<KafkaPartition> filteredPartitionsToBeProcessed = topicPartitions.stream()
          .filter(partition -> filteredPartitions.get().contains(partition.getId()))
          .collect(Collectors.toList());
      workUnitMap = getWorkUnits(filteredPartitionsToBeProcessed, state, topicSpecificState);
    } else {
      workUnitMap = getWorkUnits(topicPartitions, state, topicSpecificState);
    }

    if (!topicQualified) {
      workUnitMap.values().forEach(KafkaSource::skipWorkUnit);
    }

    for (WorkUnit workUnit : workUnitMap.values()) {
      workUnit.setProp(NUM_TOPIC_PARTITIONS, topicPartitions.size());
      workUnits.add(workUnit);
    }

    this.partitionsToBeProcessed.addAll(topic.getPartitions());
    return workUnits;
  }

  /**
   * Whether a {@link KafkaTopic} is qualified to be pulled.
   *
   * This method can be overridden by subclasses for verifying topic eligibility, e.g., one may want to
   * skip a topic if its schema cannot be found in the schema registry.
   */
  protected boolean isTopicQualified(KafkaTopic topic) {
    return true;
  }

  @SuppressWarnings("deprecation")
  private static void skipWorkUnit(WorkUnit workUnit) {
    workUnit.setProp(ConfigurationKeys.WORK_UNIT_HIGH_WATER_MARK_KEY, workUnit.getLowWaterMark());
  }

  /**
   * Get the workunits of all the partitions passed, this method fetches all the offsets for the partitions
   * at once from kafka, and for each partiton creates a workunit.
   * @param partitions
   * @param state
   * @param topicSpecificState
   * @return
   */
  private Map<KafkaPartition, WorkUnit> getWorkUnits(Collection<KafkaPartition> partitions, SourceState state,
      Optional<State> topicSpecificState) {
    final Map<KafkaPartition, Offsets> partitionOffsetMap = Maps.newHashMap();
    final Set<KafkaPartition> failedOffsetsGetList = Sets.newHashSet();
    try (final Timer.Context context = this.metricContext.timer(OFFSET_FETCH_TIMER).time()) {
      // Fetch the offsets for all the partitions at once
      final Map<KafkaPartition, Long> earliestOffsetMap = this.kafkaConsumerClient.get().getEarliestOffsets(partitions);
      final Map<KafkaPartition, Long> latestOffsetMap = this.kafkaConsumerClient.get().getLatestOffsets(partitions);
      for (KafkaPartition partition : partitions) {
        final Offsets offsets = new Offsets();
        offsets.setOffsetFetchEpochTime(System.currentTimeMillis());
        // Check if both earliest and latest offset are fetched for the partition, then set the offsets
        if (earliestOffsetMap.containsKey(partition) && latestOffsetMap.containsKey(partition)) {
          offsets.setEarliestOffset(earliestOffsetMap.get(partition));
          offsets.setLatestOffset(latestOffsetMap.get(partition));
          offsets.setOffsetFetchEpochTime(System.currentTimeMillis());
          partitionOffsetMap.put(partition, offsets);
          // If either is not available, put it in the failed offsets list
        } else {
          failedOffsetsGetList.add(partition);
        }
      }
      LOG.info("Time taken to fetch offset for partitions {} is {} ms", partitions,
          TimeUnit.NANOSECONDS.toMillis(context.stop()));
    } catch (KafkaOffsetRetrievalFailureException e) {
      // When exception occurred  while fetching earliest or latest offset for all the partitions,
      // add all the partitions to fetchOffsetsFailedPartitions
      failedOffsetsGetList.addAll(partitions);
      LOG.error("Caught error in creating work unit for {}", partitions, e);
    }
    if (!failedOffsetsGetList.isEmpty()) {
      LOG.error("Failed to fetch offsets for partitions {}", failedOffsetsGetList);
    }
    final Map<KafkaPartition, WorkUnit> workUnitMap = Maps.newHashMap();
    for (Map.Entry<KafkaPartition, Offsets> partitionOffset : partitionOffsetMap.entrySet()) {
      WorkUnit workUnit =
          getWorkUnitForTopicPartition(partitionOffset.getKey(), state, topicSpecificState, partitionOffset.getValue(),
              failedOffsetsGetList.contains(partitionOffset.getKey()));
      if (workUnit != null) {
        workUnitMap.put(partitionOffset.getKey(), workUnit);
      }
    }
    return workUnitMap;
  }

  private WorkUnit getWorkUnitForTopicPartition(KafkaPartition partition, SourceState state,
      Optional<State> topicSpecificState, Offsets offsets, boolean failedToGetKafkaOffsets) {

    long previousOffset = 0;
    long previousOffsetFetchEpochTime = 0;
    boolean previousOffsetNotFound = false;
    try {
      previousOffset = getPreviousOffsetForPartition(partition, state);
      offsets.setPreviousEndOffset(previousOffset);
      offsets.setPreviousStartOffset(getPreviousLowWatermark(partition, state));
      offsets.setPreviousStartFetchEpochTime(getPreviousStartFetchEpochTimeForPartition(partition, state));
      offsets.setPreviousStopFetchEpochTime(getPreviousStopFetchEpochTimeForPartition(partition, state));
      offsets.setPreviousLatestOffset(getPreviousExpectedHighWatermark(partition, state));
      previousOffsetFetchEpochTime = getPreviousOffsetFetchEpochTimeForPartition(partition, state);
      offsets.setPreviousOffsetFetchEpochTime(previousOffsetFetchEpochTime);
    } catch (PreviousOffsetNotFoundException e) {
      previousOffsetNotFound = true;
    }

    if (failedToGetKafkaOffsets) {

      // Increment counts, which will be reported as job metrics
      this.failToGetOffsetCount.incrementAndGet();

      // When unable to get earliest/latest offsets from Kafka, skip the partition and create an empty workunit,
      // so that previousOffset is persisted.
      LOG.warn(String
          .format("Failed to retrieve earliest and/or latest offset for partition %s. This partition will be skipped.",
              partition));
      return previousOffsetNotFound ? null : createEmptyWorkUnit(partition, previousOffset, previousOffsetFetchEpochTime,
          topicSpecificState);
    }

    if (shouldMoveToLatestOffset(partition, state)) {
      offsets.startAtLatestOffset();
    } else if (previousOffsetNotFound) {

      /**
       * When previous offset cannot be found, either start at earliest offset, latest offset, go back with (latest - lookback)
       * (long value to be deducted from latest offset in order to avoid data loss) or skip the partition
       * (no need to create an empty workunit in this case since there's no offset to persist).
       * In case of no previous state OFFSET_LOOKBACK will make sure to avoid consuming huge amount of data (earlist) and data loss (latest offset)
       * lookback can be set to any long value where (latest-lookback) is nearest offset for each partition. If computed offset is out of range then
       * partition will be consumed from latest offset
       **/
      String offsetNotFoundMsg = String.format("Previous offset for partition %s does not exist. ", partition);
      String offsetOption = state.getProp(BOOTSTRAP_WITH_OFFSET, DEFAULT_BOOTSTRAP_WITH_OFFSET).toLowerCase();
      if (offsetOption.equals(LATEST_OFFSET)) {
        LOG.warn(offsetNotFoundMsg + "This partition will start from the latest offset: " + offsets.getLatestOffset());
        offsets.startAtLatestOffset();
      } else if (offsetOption.equals(EARLIEST_OFFSET)) {
        LOG.warn(
            offsetNotFoundMsg + "This partition will start from the earliest offset: " + offsets.getEarliestOffset());
        offsets.startAtEarliestOffset();
      } else if (offsetOption.equals(OFFSET_LOOKBACK)) {
        long lookbackOffsetRange = state.getPropAsLong(KAFKA_OFFSET_LOOKBACK , 0L);
        long latestOffset = offsets.getLatestOffset();
        long offset = latestOffset - lookbackOffsetRange;
        LOG.warn(offsetNotFoundMsg + "This partition will start from latest-lookback [ " + latestOffset + " - " + lookbackOffsetRange + " ]  start offset: " + offset);
        try {
          offsets.startAt(offset);
        } catch (StartOffsetOutOfRangeException e) {
          // Increment counts, which will be reported as job metrics
          if (offsets.getStartOffset() <= offsets.getLatestOffset()) {
            this.offsetTooEarlyCount.incrementAndGet();
          } else {
            this.offsetTooLateCount.incrementAndGet();
          }

          // When above computed offset (latest-lookback) is out of range, either start at earliest, latest or nearest offset, or skip the
          // partition. If skipping, need to create an empty workunit so that previousOffset is persisted.
          String offsetOutOfRangeMsg = String.format(
                  "Start offset for partition %s is out of range. Start offset = %d, earliest offset = %d, latest offset = %d.",
                  partition, offsets.getStartOffset(), offsets.getEarliestOffset(), offsets.getLatestOffset());
          offsetOption =
                  state.getProp(RESET_ON_OFFSET_OUT_OF_RANGE, DEFAULT_RESET_ON_OFFSET_OUT_OF_RANGE).toLowerCase();
          if (offsetOption.equals(LATEST_OFFSET) || (offsetOption.equals(NEAREST_OFFSET)
                  && offsets.getStartOffset() >= offsets.getLatestOffset())) {
            LOG.warn(
                    offsetOutOfRangeMsg + "This partition will start from the latest offset: " + offsets.getLatestOffset());
            offsets.startAtLatestOffset();
          } else if (offsetOption.equals(EARLIEST_OFFSET) || offsetOption.equals(NEAREST_OFFSET)) {
