/*
 * 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.physical.impl.window;

import org.apache.drill.common.exceptions.DrillException;
import org.apache.drill.exec.exception.SchemaChangeException;
import org.apache.drill.exec.ops.OperatorContext;
import org.apache.drill.exec.physical.config.WindowPOP;
import org.apache.drill.exec.record.VectorAccessible;
import org.apache.drill.exec.record.VectorContainer;
import org.apache.drill.exec.record.VectorWrapper;
import org.apache.drill.exec.vector.BaseDataValueVector;
import org.apache.drill.exec.vector.ValueVector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.inject.Named;
import java.util.List;


/**
 * WindowFramer implementation that supports the FRAME clause.
 * <br>According to the SQL specification, FIRST_VALUE, LAST_VALUE and all aggregate functions support the FRAME clause.
 * This class will handle such functions even if the FRAME clause is not present.
 */
public abstract class FrameSupportTemplate implements WindowFramer {

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

  private VectorContainer container;
  private VectorContainer internal;
  private List<WindowDataBatch> batches;
  private int outputCount; // number of rows in currently/last processed batch

  private WindowDataBatch current;

  private int frameLastRow;

  // true when at least one window function needs to process all batches of a partition before passing any batch downstream
  private boolean requireFullPartition;

  private long remainingRows; // num unprocessed rows in current partition
  private long remainingPeers; // num unprocessed peer rows in current frame
  private boolean partialPartition; // true if we remainingRows only account for the current batch and more batches are expected for the current partition

  private WindowPOP popConfig;

  @Override
  public void setup(final List<WindowDataBatch> batches, final VectorContainer container, final OperatorContext oContext,
                    final boolean requireFullPartition, final WindowPOP popConfig) throws SchemaChangeException {
    this.container = container;
    this.batches = batches;

    internal = new VectorContainer(oContext);
    allocateInternal();

    outputCount = 0;

    this.requireFullPartition = requireFullPartition;
    this.popConfig = popConfig;
  }

  private void allocateInternal() {
    for (VectorWrapper<?> w : container) {
      ValueVector vv = internal.addOrGet(w.getField());
      vv.allocateNew();
    }
  }

  private boolean isPartitionDone() {
    return !partialPartition && remainingRows == 0;
  }

  /**
   * processes all rows of the first batch.
   */
  @Override
  public void doWork() throws DrillException {
    int currentRow = 0;

    this.current = batches.get(0);

    setupSaveFirstValue(current, internal);

    outputCount = current.getRecordCount();

    while (currentRow < outputCount) {
      if (!isPartitionDone()) {
        // we have a pending partition we need to handle from a previous call to doWork()
        assert currentRow == 0 : "pending partitions are only expected at the start of the batch";
        logger.trace("we have a pending partition {}", remainingRows);

        if (!requireFullPartition) {
          // we didn't compute the whole partition length in the previous partition, we need to update the length now
          updatePartitionSize(currentRow);
        }
      } else {
        newPartition(current, currentRow);
      }

      currentRow = processPartition(currentRow);
      if (isPartitionDone()) {
        reset();
      }
    }
  }

  private void newPartition(final WindowDataBatch current, final int currentRow) throws SchemaChangeException {
    remainingRows = 0;
    remainingPeers = 0;
    updatePartitionSize(currentRow);

    setupPartition(current, container);
    saveFirstValue(currentRow);
  }

  private void reset() {
    resetValues();
    for (VectorWrapper<?> vw : internal) {
      if ((vw.getValueVector() instanceof BaseDataValueVector)) {
        ((BaseDataValueVector) vw.getValueVector()).reset();
      }
    }
  }

  /**
   * process all rows (computes and writes aggregation values) of current batch that are part of current partition.
   * @param currentRow first unprocessed row
   * @return index of next unprocessed row
   * @throws DrillException if it can't write into the container
   */
  private int processPartition(final int currentRow) throws DrillException {
    logger.trace("{} rows remaining to process, currentRow: {}, outputCount: {}", remainingRows, currentRow, outputCount);

    setupWriteFirstValue(internal, container);

    if (popConfig.isFrameUnitsRows()) {
      return processROWS(currentRow);
    } else {
      return processRANGE(currentRow);
    }
  }

  private int processROWS(int row) throws DrillException {
    //TODO (DRILL-4413) we only need to call these once per batch
    setupEvaluatePeer(current, container);
    setupReadLastValue(current, container);

    while (row < outputCount && !isPartitionDone()) {
      logger.trace("aggregating row {}", row);
      evaluatePeer(row);

      outputRow(row);
      writeLastValue(row, row);

      remainingRows--;
      row++;
    }

    return row;
  }

  private int processRANGE(int row) throws DrillException {
    while (row < outputCount && !isPartitionDone()) {
      if (remainingPeers == 0) {
        // because all peer rows share the same frame, we only need to compute and aggregate the frame once
        if (popConfig.getStart().isCurrent()) {
          reset();
          saveFirstValue(row);
        }

        remainingPeers = aggregatePeers(row);
      }

      outputRow(row);
      writeLastValue(frameLastRow, row);

      remainingRows--;
      remainingPeers--;
      row++;
    }

    return row;
  }

  /**
   * updates partition's length after computing the number of rows for the current the partition starting at the specified
   * row of the first batch. If !requiresFullPartition, this method will only count the rows in the current batch
   */
  private void updatePartitionSize(final int start) {
    logger.trace("compute partition size starting from {} on {} batches", start, batches.size());

    long length = 0;
    int row = start;

    // count all rows that are in the same partition of start
    // keep increasing length until we find first row of next partition or we reach the very last batch

    outer:
    for (WindowDataBatch batch : batches) {
      final int recordCount = batch.getRecordCount();

      // check first container from start row, and subsequent containers from first row
      for (; row < recordCount; row++, length++) {
        if (!isSamePartition(start, current, row, batch)) {
          break outer;
        }
      }

      if (!requireFullPartition) {
        // we are only interested in the first batch's records
        break;
      }

      row = 0;
    }

    if (!requireFullPartition) {
      // this is the last batch of current partition if
      boolean lastBatch = row < outputCount                     // partition ends before the end of the batch
        || batches.size() == 1                                  // it's the last available batch
        || !isSamePartition(start, current, 0, batches.get(1)); // next batch contains a different partition

      partialPartition = !lastBatch;
    } else {
      partialPartition = false;
    }

    remainingRows += length;
  }

  /**
   * aggregates all peer rows of current row
   * @param start starting row of the current frame
   * @return num peer rows for current row
   * @throws SchemaChangeException
   */
  private long aggregatePeers(final int start) throws SchemaChangeException {
    logger.trace("aggregating rows starting from {}", start);

    final boolean unboundedFollowing = popConfig.getEnd().isUnbounded();
    VectorAccessible last = current;
    long length = 0;

    // a single frame can include rows from multiple batches
    // start processing first batch and, if necessary, move to next batches
    for (WindowDataBatch batch : batches) {
      setupEvaluatePeer(batch, container);
      final int recordCount = batch.getRecordCount();

      // for every remaining row in the partition, count it if it's a peer row
      for (int row = (batch == current) ? start : 0; row < recordCount; row++, length++) {
        if (unboundedFollowing) {
          if (length >= remainingRows) {
            break;
          }
        } else {
          if (!isPeer(start, current, row, batch)) {
            break;
          }
        }

        evaluatePeer(row);
        last = batch;
        frameLastRow = row;
      }
    }

    setupReadLastValue(last, container);

    return length;
  }

  @Override
  public int getOutputCount() {
    return outputCount;
  }

  // we need this abstract method for code generation
  @Override
  public void cleanup() {
    logger.trace("clearing internal");
    internal.clear();
  }

  @Override
  public String toString() {
    return "FrameSupportTemplate[internal=" + internal
        + ", outputCount=" + outputCount
        + ", current=" + current
        + ", frameLastRow=" + frameLastRow
        + ", remainingRows=" + remainingRows
        + ", partialPartition=" + partialPartition
        + "]";
  }

  /**
   * called once for each peer row of the current frame.
   * @param index of row to aggregate
   */
  public abstract void evaluatePeer(@Named("index") int index);
  public abstract void setupEvaluatePeer(@Named("incoming") VectorAccessible incoming, @Named("outgoing") VectorAccessible outgoing) throws SchemaChangeException;

  public abstract void setupReadLastValue(@Named("incoming") VectorAccessible incoming, @Named("outgoing") VectorAccessible outgoing) throws SchemaChangeException;
  public abstract void writeLastValue(@Named("index") int index, @Named("outIndex") int outIndex);

  public abstract void setupSaveFirstValue(@Named("incoming") VectorAccessible incoming, @Named("outgoing") VectorAccessible outgoing) throws SchemaChangeException;
  public abstract void saveFirstValue(@Named("index") int index);
  public abstract void setupWriteFirstValue(@Named("incoming") VectorAccessible incoming, @Named("outgoing") VectorAccessible outgoing);

  /**
   * called once for each row after we evaluate all peer rows. Used to write a value in the row
   *
   * @param outIndex index of row
   */
  public abstract void outputRow(@Named("outIndex") int outIndex);

  /**
   * Called once per partition, before processing the partition. Used to setup read/write vectors
   * @param incoming batch we will read from
   * @param outgoing batch we will be writing to
   *
   * @throws SchemaChangeException
   */
  public abstract void setupPartition(@Named("incoming") WindowDataBatch incoming,
                                      @Named("outgoing") VectorAccessible outgoing) throws SchemaChangeException;

  /**
   * reset all window functions
   */
  public abstract boolean resetValues();

  /**
   * compares two rows from different batches (can be the same), if they have the same value for the partition by
   * expression
   * @param b1Index index of first row
   * @param b1 batch for first row
   * @param b2Index index of second row
   * @param b2 batch for second row
   * @return true if the rows are in the same partition
   */
  public abstract boolean isSamePartition(@Named("b1Index") int b1Index, @Named("b1") VectorAccessible b1,
                                          @Named("b2Index") int b2Index, @Named("b2") VectorAccessible b2);

  /**
   * compares two rows from different batches (can be the same), if they have the same value for the order by
   * expression
   * @param b1Index index of first row
   * @param b1 batch for first row
   * @param b2Index index of second row
   * @param b2 batch for second row
   * @return true if the rows are in the same partition
   */
  public abstract boolean isPeer(@Named("b1Index") int b1Index, @Named("b1") VectorAccessible b1,
                                 @Named("b2Index") int b2Index, @Named("b2") VectorAccessible b2);
}