/*
 * 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.planner.index.generators;

import java.util.ArrayList;
import java.util.List;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.calcite.plan.RelTrait;
import org.apache.calcite.plan.RelTraitSet;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.InvalidRelException;
import org.apache.calcite.rel.RelCollation;
import org.apache.calcite.rel.RelCollationTraitDef;
import org.apache.calcite.rel.type.RelDataTypeFieldImpl;
import org.apache.calcite.rel.type.RelRecordType;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.drill.common.expression.FieldReference;
import org.apache.drill.common.expression.SchemaPath;
import org.apache.drill.exec.physical.base.DbGroupScan;
import org.apache.drill.exec.planner.index.IndexCallContext;
import org.apache.drill.exec.planner.index.IndexPlanUtils;
import org.apache.drill.exec.planner.logical.DrillFilterRel;
import org.apache.drill.exec.planner.logical.DrillProjectRel;
import org.apache.drill.exec.planner.logical.DrillSortRel;
import org.apache.drill.exec.planner.common.DrillProjectRelBase;
import org.apache.drill.exec.planner.common.OrderedRel;
import org.apache.drill.exec.planner.common.DrillScanRelBase;
import org.apache.drill.exec.planner.physical.SubsetTransformer;
import org.apache.drill.exec.planner.physical.PlannerSettings;
import org.apache.drill.exec.planner.physical.Prule;
import org.apache.drill.exec.planner.physical.DrillDistributionTrait;
import org.apache.drill.exec.planner.physical.Prel;
import org.apache.drill.exec.planner.physical.HashToMergeExchangePrel;
import org.apache.drill.exec.planner.physical.SingleMergeExchangePrel;
import org.apache.drill.exec.planner.physical.PrelUtil;
import org.apache.drill.exec.planner.physical.TopNPrel;
import org.apache.drill.exec.planner.physical.SortPrel;
import org.apache.drill.exec.planner.physical.LimitPrel;
import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.rex.RexBuilder;
import org.apache.calcite.rex.RexNode;

public abstract class AbstractIndexPlanGenerator extends SubsetTransformer<RelNode, InvalidRelException>{

  private static final Logger logger = LoggerFactory.getLogger(AbstractIndexPlanGenerator.class);

  final protected DrillProjectRelBase origProject;
  final protected DrillScanRelBase origScan;
  final protected DrillProjectRelBase upperProject;

  final protected RexNode indexCondition;
  final protected RexNode remainderCondition;
  final protected RexBuilder builder;
  final protected IndexCallContext indexContext;
  final protected PlannerSettings settings;

  public AbstractIndexPlanGenerator(IndexCallContext indexContext,
      RexNode indexCondition,
      RexNode remainderCondition,
      RexBuilder builder,
      PlannerSettings settings) {
    super(indexContext.getCall());
    this.origProject = indexContext.getLowerProject();
    this.origScan = indexContext.getScan();
    this.upperProject = indexContext.getUpperProject();
    this.indexCondition = indexCondition;
    this.remainderCondition = remainderCondition;
    this.indexContext = indexContext;
    this.builder = builder;
    this.settings = settings;
  }

  // Provides the utility functions that don't rely on index(one or multiple) or final plan (covering or not),
  // but those helper functions that focus on serving building index plan (project-filter-indexscan)

  public static int getRowKeyIndex(RelDataType rowType, DrillScanRelBase origScan) {
    List<String> fieldNames = rowType.getFieldNames();
    int idx = 0;
    for (String field : fieldNames) {
      if (field.equalsIgnoreCase(((DbGroupScan)IndexPlanUtils.getGroupScan(origScan)).getRowKeyName())) {
        return idx;
      }
      idx++;
    }
    return -1;
  }

  protected RelDataType convertRowType(RelDataType origRowType, RelDataTypeFactory typeFactory) {
    if ( getRowKeyIndex(origRowType, origScan)>=0 ) { // row key already present
      return origRowType;
    }
    List<RelDataTypeField> fields = new ArrayList<>();

    fields.addAll(origRowType.getFieldList());
    fields.add(new RelDataTypeFieldImpl(
        ((DbGroupScan)IndexPlanUtils.getGroupScan(origScan)).getRowKeyName(), fields.size(),
            typeFactory.createSqlType(SqlTypeName.ANY)));
    return new RelRecordType(fields);
  }

  protected boolean checkRowKey(List<SchemaPath> columns) {
    for (SchemaPath s : columns) {
      if (s.equals(((DbGroupScan)IndexPlanUtils.getGroupScan(origScan)).getRowKeyPath())) {
        return true;
      }
    }
    return false;
  }

  // Range distribute the right side of the join, on row keys using a range partitioning function
  protected RelNode createRangeDistRight(final RelNode rightPrel,
                                         final RelDataTypeField rightRowKeyField,
                                         final DbGroupScan origDbGroupScan) {

    List<DrillDistributionTrait.DistributionField> rangeDistFields =
        Lists.newArrayList(new DrillDistributionTrait.DistributionField(0 /* rowkey ordinal on the right side */));

    FieldReference rangeDistRef = FieldReference.getWithQuotedRef(rightRowKeyField.getName());
    List<FieldReference> rangeDistRefList = Lists.newArrayList();
    rangeDistRefList.add(rangeDistRef);

    final DrillDistributionTrait distRight;
    if (IndexPlanUtils.scanIsPartition(origDbGroupScan)) {
      distRight = new DrillDistributionTrait(
          DrillDistributionTrait.DistributionType.RANGE_DISTRIBUTED,
          ImmutableList.copyOf(rangeDistFields),
          origDbGroupScan.getRangePartitionFunction(rangeDistRefList));
    }
    else {
      distRight = DrillDistributionTrait.SINGLETON;
    }

    RelTraitSet rightTraits = newTraitSet(distRight).plus(Prel.DRILL_PHYSICAL);
    RelNode convertedRight = Prule.convert(rightPrel, rightTraits);

    return convertedRight;
  }

  @Override
  public RelTraitSet newTraitSet(RelTrait... traits) {
    RelTraitSet set = indexContext.getCall().getPlanner().emptyTraitSet();
    for (RelTrait t : traits) {
      if (t != null) {
        set = set.plus(t);
      }
    }
    return set;
  }

  protected static boolean toRemoveSort(RelCollation sortCollation, RelCollation inputCollation) {
    return (inputCollation != null) && inputCollation.satisfies(sortCollation);
  }

  public static RelNode getExchange(RelOptCluster cluster, boolean isSingleton, boolean isExchangeRequired,
                                    RelTraitSet traits, DrillDistributionTrait distributionTrait,
                                    IndexCallContext indexContext, RelNode input) {
    if (!isExchangeRequired) {
      return input;
    }

    if (isSingleton) {
      return new SingleMergeExchangePrel(cluster,
              traits.replace(DrillDistributionTrait.SINGLETON),
              input, indexContext.getCollation());
    } else {
      return new HashToMergeExchangePrel(cluster,
              traits.replace(distributionTrait),
              input, distributionTrait.getFields(), indexContext.getCollation(),
              PrelUtil.getSettings(cluster).numEndPoints());
    }
  }

  private static RelNode getSortOrTopN(IndexCallContext indexContext,
                                       RelNode sortNode, RelNode newRel, RelNode child) {
    if (sortNode instanceof TopNPrel) {
      return new TopNPrel(sortNode.getCluster(),
                    newRel.getTraitSet().replace(Prel.DRILL_PHYSICAL).plus(indexContext.getCollation()),
                    child, ((TopNPrel)sortNode).getLimit(), indexContext.getCollation());
    }
    return new SortPrel(sortNode.getCluster(),
            newRel.getTraitSet().replace(Prel.DRILL_PHYSICAL).plus(indexContext.getCollation()),
            child, indexContext.getCollation());
  }

  public static RelNode getSortNode(IndexCallContext indexContext, RelNode newRel, boolean donotGenerateSort,
                                    boolean isSingleton, boolean isExchangeRequired) {
    OrderedRel rel = indexContext.getSort();
    DrillDistributionTrait hashDistribution =
        new DrillDistributionTrait(DrillDistributionTrait.DistributionType.HASH_DISTRIBUTED,
            ImmutableList.copyOf(indexContext.getDistributionFields()));

    if ( toRemoveSort(indexContext.getCollation(), newRel.getTraitSet().getTrait(RelCollationTraitDef.INSTANCE))) {
      //we are going to remove sort
      logger.debug("Not generating SortPrel since we have the required collation");
      if (IndexPlanUtils.generateLimit(rel)) {
        newRel = new LimitPrel(newRel.getCluster(),
                newRel.getTraitSet().plus(indexContext.getCollation()).plus(Prel.DRILL_PHYSICAL),
                newRel, IndexPlanUtils.getOffset(rel), IndexPlanUtils.getFetch(rel));
      }
      RelTraitSet traits = newRel.getTraitSet().plus(indexContext.getCollation()).plus(Prel.DRILL_PHYSICAL);
      newRel = Prule.convert(newRel, traits);
      newRel = getExchange(newRel.getCluster(), isSingleton, isExchangeRequired,
                                 traits, hashDistribution, indexContext, newRel);
    }
    else {
      if (donotGenerateSort) {
        logger.debug("Not generating SortPrel and index plan, since just picking index for full index scan is not beneficial.");
        return null;
      }
      RelTraitSet traits = newRel.getTraitSet().plus(indexContext.getCollation()).plus(Prel.DRILL_PHYSICAL);
      newRel = getSortOrTopN(indexContext, rel, newRel,
                  Prule.convert(newRel, newRel.getTraitSet().replace(Prel.DRILL_PHYSICAL)));
      newRel = getExchange(newRel.getCluster(), isSingleton, isExchangeRequired,
                                 traits, hashDistribution, indexContext, newRel);
    }
    return newRel;
  }

  @Override
  public abstract RelNode convertChild(RelNode current, RelNode child) throws InvalidRelException;

  @Override
  public boolean forceConvert(){
    return true;
  }

  public void go() throws InvalidRelException {
    RelNode top = indexContext.getCall().rel(0);
    final RelNode input;
    if (top instanceof DrillProjectRel) {
      DrillProjectRel topProject = (DrillProjectRel) top;
      input = topProject.getInput();
    }
    else if (top instanceof DrillFilterRel) {
      DrillFilterRel topFilter = (DrillFilterRel)top;
      input = topFilter.getInput();
    } else if (top instanceof DrillSortRel) {
      DrillSortRel topSort = (DrillSortRel)top;
      input = topSort.getInput();
    }
    else {
      return;
    }
    RelTraitSet traits = input.getTraitSet().plus(Prel.DRILL_PHYSICAL);
    RelNode convertedInput = Prule.convert(input, traits);
    this.go(top, convertedInput);
  }
}