/*
 * 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.drill.exec.store.kafka;

import org.apache.drill.common.FunctionNames;
import org.apache.drill.common.expression.BooleanOperator;
import org.apache.drill.common.expression.FunctionCall;
import org.apache.drill.common.expression.LogicalExpression;
import org.apache.drill.common.expression.visitors.AbstractExprVisitor;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;

import java.util.List;
import java.util.Map;
import java.util.Set;

public class KafkaPartitionScanSpecBuilder
  extends AbstractExprVisitor<List<KafkaPartitionScanSpec>, Void, RuntimeException>
  implements AutoCloseable {

  private final LogicalExpression le;
  private final KafkaGroupScan groupScan;
  private final KafkaConsumer<?, ?> kafkaConsumer;
  private ImmutableMap<TopicPartition, KafkaPartitionScanSpec> fullScanSpec;

  public KafkaPartitionScanSpecBuilder(KafkaGroupScan groupScan, LogicalExpression conditionExp) {
    this.groupScan = groupScan;
    this.kafkaConsumer = new KafkaConsumer<>(groupScan.getKafkaStoragePluginConfig().getKafkaConsumerProps(),
        new ByteArrayDeserializer(), new ByteArrayDeserializer());
    this.le = conditionExp;
  }

  public List<KafkaPartitionScanSpec> parseTree() {
    ImmutableMap.Builder<TopicPartition, KafkaPartitionScanSpec> builder = ImmutableMap.builder();
    for (KafkaPartitionScanSpec scanSpec : groupScan.getPartitionScanSpecList()) {
      builder.put(new TopicPartition(scanSpec.getTopicName(), scanSpec.getPartitionId()), scanSpec);
    }
    fullScanSpec = builder.build();
    List<KafkaPartitionScanSpec> pushdownSpec = le.accept(this, null);

    /*
    Non-existing / invalid partitions may result in empty scan spec.
    This results in a "ScanBatch" with no reader. DRILL currently requires
    at least one reader to be present in a scan batch.
     */
    if (pushdownSpec != null && pushdownSpec.isEmpty()) {
      TopicPartition firstPartition = new TopicPartition(groupScan.getKafkaScanSpec().getTopicName(), 0);
      KafkaPartitionScanSpec emptySpec =
          new KafkaPartitionScanSpec(firstPartition.topic(),firstPartition.partition(),
              fullScanSpec.get(firstPartition).getEndOffset(), fullScanSpec.get(firstPartition).getEndOffset());
      pushdownSpec.add(emptySpec);
    }
    return pushdownSpec;
  }

  @Override
  public List<KafkaPartitionScanSpec> visitUnknown(LogicalExpression e, Void value) {
    return null;
  }

  @Override
  public List<KafkaPartitionScanSpec> visitBooleanOperator(BooleanOperator op, Void value) {
    Map<TopicPartition, KafkaPartitionScanSpec> specMap = Maps.newHashMap();
    List<LogicalExpression> args = op.args();
    if (op.getName().equals("booleanOr")) {

      for (LogicalExpression expr : args) {
        List<KafkaPartitionScanSpec> parsedSpec = expr.accept(this, null);
        //parsedSpec is null if expression cannot be pushed down
        if (parsedSpec != null) {
          for(KafkaPartitionScanSpec newSpec : parsedSpec) {
            TopicPartition tp = new TopicPartition(newSpec.getTopicName(), newSpec.getPartitionId());
            KafkaPartitionScanSpec existingSpec = specMap.get(tp);

            //If existing spec does not contain topic-partition
            if (existingSpec == null) {
              specMap.put(tp, newSpec); //Add topic-partition to spec for OR
            } else {
              existingSpec.mergeScanSpec(op.getName(), newSpec);
              specMap.put(tp, existingSpec);
            }
          }
        } else {
          return null; //At any level, all arguments of booleanOr should support pushdown, else return null
        }
      }
    } else { //booleanAnd
      specMap.putAll(fullScanSpec);
      for (LogicalExpression expr : args) {
        List<KafkaPartitionScanSpec> parsedSpec = expr.accept(this, null);

        //parsedSpec is null if expression cannot be pushed down
        if (parsedSpec != null) {
          Set<TopicPartition> partitionsInNewSpec = Sets.newHashSet(); //Store topic-partitions returned from new spec.

          for (KafkaPartitionScanSpec newSpec : parsedSpec) {
            TopicPartition tp = new TopicPartition(newSpec.getTopicName(), newSpec.getPartitionId());
            partitionsInNewSpec.add(tp);
            KafkaPartitionScanSpec existingSpec = specMap.get(tp);

            if (existingSpec != null) {
              existingSpec.mergeScanSpec(op.getName(), newSpec);
              specMap.put(tp, existingSpec);
            }
          }

          /*
          For "booleanAnd", handle the case where condition is on `kafkaPartitionId`.
          In this case, we would not want unnecessarily scan all the topic-partitions.
          Hence we remove the unnecessary topic-partitions from the spec.
         */
          specMap.keySet().removeIf(partition -> !partitionsInNewSpec.contains(partition));
        }

      }
    }
    return Lists.newArrayList(specMap.values());
  }

  @Override
  public List<KafkaPartitionScanSpec> visitFunctionCall(FunctionCall call, Void value) {
    String functionName = call.getName();
    if (KafkaNodeProcessor.isPushdownFunction(functionName)) {

      KafkaNodeProcessor kafkaNodeProcessor = KafkaNodeProcessor.process(call);
      if (kafkaNodeProcessor.isSuccess()) {
        switch (kafkaNodeProcessor.getPath()) {
          case "kafkaMsgTimestamp":
            return createScanSpecForTimestamp(kafkaNodeProcessor.getFunctionName(),
                kafkaNodeProcessor.getValue());
          case "kafkaMsgOffset":
            return createScanSpecForOffset(kafkaNodeProcessor.getFunctionName(),
                kafkaNodeProcessor.getValue());
          case "kafkaPartitionId":
            return createScanSpecForPartition(kafkaNodeProcessor.getFunctionName(),
                kafkaNodeProcessor.getValue());
        }
      }
    }
    return null; //Return null, do not pushdown
  }


  private List<KafkaPartitionScanSpec> createScanSpecForTimestamp(String functionName, Long fieldValue) {
    List<KafkaPartitionScanSpec> scanSpec = Lists.newArrayList();
    Map<TopicPartition, Long> timesValMap = Maps.newHashMap();
    ImmutableSet<TopicPartition> topicPartitions = fullScanSpec.keySet();

    for (TopicPartition partitions : topicPartitions) {
      timesValMap.put(partitions, functionName.equals(FunctionNames.GT) ? fieldValue+1 : fieldValue);
    }

    Map<TopicPartition, OffsetAndTimestamp> offsetAndTimestamp = kafkaConsumer.offsetsForTimes(timesValMap);

    for (TopicPartition tp : topicPartitions) {
      OffsetAndTimestamp value = offsetAndTimestamp.get(tp);
      //OffsetAndTimestamp is null if there is no offset greater or equal to requested timestamp
      if (value == null) {
        scanSpec.add(
            new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                fullScanSpec.get(tp).getEndOffset(), fullScanSpec.get(tp).getEndOffset()));
      } else {
        scanSpec.add(
            new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                value.offset(), fullScanSpec.get(tp).getEndOffset()));
      }
    }

    return scanSpec;
  }

  private List<KafkaPartitionScanSpec> createScanSpecForOffset(String functionName, Long fieldValue) {
    List<KafkaPartitionScanSpec> scanSpec = Lists.newArrayList();
    ImmutableSet<TopicPartition> topicPartitions = fullScanSpec.keySet();

    /*
    We should handle the case where the specified offset does not exist in the current context,
    i.e., fieldValue < startOffset or fieldValue > endOffset in a particular topic-partition.
    Else, KafkaConsumer.poll will throw "TimeoutException".
    */

    switch (functionName) {
      case FunctionNames.EQ:
        for (TopicPartition tp : topicPartitions) {
          if (fieldValue < fullScanSpec.get(tp).getStartOffset()) {
            //Offset does not exist
            scanSpec.add(
                new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                    fullScanSpec.get(tp).getEndOffset(), fullScanSpec.get(tp).getEndOffset()));
          } else {
            long val = Math.min(fieldValue, fullScanSpec.get(tp).getEndOffset());
            long nextVal = Math.min(val+1, fullScanSpec.get(tp).getEndOffset());
            scanSpec.add(new KafkaPartitionScanSpec(tp.topic(), tp.partition(), val, nextVal));
          }
        }
        break;
      case FunctionNames.GE:
        for (TopicPartition tp : topicPartitions) {
          //Ensure scan range is between startOffset and endOffset,
          long val = bindOffsetToRange(tp, fieldValue);
          scanSpec.add(
              new KafkaPartitionScanSpec(tp.topic(), tp.partition(), val,
                  fullScanSpec.get(tp).getEndOffset()));
        }
        break;
      case FunctionNames.GT:
        for (TopicPartition tp : topicPartitions) {
          //Ensure scan range is between startOffset and endOffset,
          long val = bindOffsetToRange(tp, fieldValue + 1);
          scanSpec.add(
              new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                  val, fullScanSpec.get(tp).getEndOffset()));
        }
        break;
      case FunctionNames.LE:
        for (TopicPartition tp : topicPartitions) {
          //Ensure scan range is between startOffset and endOffset,
          long val = bindOffsetToRange(tp, fieldValue+1);

          scanSpec.add(
              new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                  fullScanSpec.get(tp).getStartOffset(), val));
        }
        break;
      case FunctionNames.LT:
        for (TopicPartition tp : topicPartitions) {
          //Ensure scan range is between startOffset and endOffset,
          long val = bindOffsetToRange(tp, fieldValue);

          scanSpec.add(
              new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                  fullScanSpec.get(tp).getStartOffset(), val));
        }
        break;
    }
    return scanSpec;
  }

  private List<KafkaPartitionScanSpec> createScanSpecForPartition(String functionName, Long fieldValue) {
    List<KafkaPartitionScanSpec> scanSpecList = Lists.newArrayList();
    ImmutableSet<TopicPartition> topicPartitions = fullScanSpec.keySet();

    switch (functionName) {
      case FunctionNames.EQ:
        for (TopicPartition tp : topicPartitions) {
          if (tp.partition() == fieldValue) {
            scanSpecList.add(
                new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                    fullScanSpec.get(tp).getStartOffset(),
                    fullScanSpec.get(tp).getEndOffset()));
          }
        }
        break;
      case FunctionNames.NE:
        for (TopicPartition tp : topicPartitions) {
          if (tp.partition() != fieldValue) {
            scanSpecList.add(
                new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                    fullScanSpec.get(tp).getStartOffset(),
                    fullScanSpec.get(tp).getEndOffset()));
          }
        }
        break;
      case FunctionNames.GE:
        for (TopicPartition tp : topicPartitions) {
          if (tp.partition() >= fieldValue) {
            scanSpecList.add(
                new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                    fullScanSpec.get(tp).getStartOffset(),
                    fullScanSpec.get(tp).getEndOffset()));
          }
        }
        break;
      case FunctionNames.GT:
        for (TopicPartition tp : topicPartitions) {
          if (tp.partition() > fieldValue) {
            scanSpecList.add(
                new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                    fullScanSpec.get(tp).getStartOffset(),
                    fullScanSpec.get(tp).getEndOffset()));
          }
        }
        break;
      case FunctionNames.LE:
        for (TopicPartition tp : topicPartitions) {
          if (tp.partition() <= fieldValue) {
            scanSpecList.add(
                new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                    fullScanSpec.get(tp).getStartOffset(),
                    fullScanSpec.get(tp).getEndOffset()));
          }
        }
        break;
      case FunctionNames.LT:
        for (TopicPartition tp : topicPartitions) {
          if (tp.partition() < fieldValue) {
            scanSpecList.add(
                new KafkaPartitionScanSpec(tp.topic(), tp.partition(),
                    fullScanSpec.get(tp).getStartOffset(),
                    fullScanSpec.get(tp).getEndOffset()));
          }
        }
        break;
    }
    return scanSpecList;
  }

  @Override
  public void close() {
    groupScan.getStoragePlugin().registerToClose(kafkaConsumer);
  }

  private long bindOffsetToRange(TopicPartition tp, long offset) {
    return Math.max(fullScanSpec.get(tp).getStartOffset(), Math.min(offset, fullScanSpec.get(tp).getEndOffset()));
  }
}