/*
 * 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.arrow.gandiva.evaluator;

import java.math.BigDecimal;
import java.time.Instant;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import org.apache.arrow.gandiva.exceptions.GandivaException;
import org.apache.arrow.gandiva.expression.Condition;
import org.apache.arrow.gandiva.expression.ExpressionTree;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.DecimalVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;

class BaseEvaluatorTest {

  interface BaseEvaluator {

    void evaluate(ArrowRecordBatch recordBatch, BufferAllocator allocator) throws GandivaException;

    long getElapsedMillis();
  }

  class ProjectEvaluator implements BaseEvaluator {

    private Projector projector;
    private DataAndVectorGenerator generator;
    private int numExprs;
    private int maxRowsInBatch;
    private long elapsedTime = 0;
    private List<ValueVector> outputVectors = new ArrayList<>();

    public ProjectEvaluator(
        Projector projector, DataAndVectorGenerator generator, int numExprs, int maxRowsInBatch) {
      this.projector = projector;
      this.generator = generator;
      this.numExprs = numExprs;
      this.maxRowsInBatch = maxRowsInBatch;
    }

    @Override
    public void evaluate(ArrowRecordBatch recordBatch, BufferAllocator allocator)
        throws GandivaException {
      // set up output vectors
      // for each expression, generate the output vector
      for (int i = 0; i < numExprs; i++) {
        ValueVector valueVector = generator.generateOutputVector(maxRowsInBatch);
        outputVectors.add(valueVector);
      }

      try {
        long start = System.nanoTime();
        projector.evaluate(recordBatch, outputVectors);
        long finish = System.nanoTime();
        elapsedTime += (finish - start);
      } finally {
        for (ValueVector valueVector : outputVectors) {
          valueVector.close();
        }
      }
      outputVectors.clear();
    }

    @Override
    public long getElapsedMillis() {
      return TimeUnit.NANOSECONDS.toMillis(elapsedTime);
    }
  }

  class FilterEvaluator implements BaseEvaluator {

    private Filter filter;
    private long elapsedTime = 0;

    public FilterEvaluator(Filter filter) {
      this.filter = filter;
    }

    @Override
    public void evaluate(ArrowRecordBatch recordBatch, BufferAllocator allocator)
        throws GandivaException {
      ArrowBuf selectionBuffer = allocator.buffer(recordBatch.getLength() * 2);
      SelectionVectorInt16 selectionVector = new SelectionVectorInt16(selectionBuffer);

      try {
        long start = System.nanoTime();
        filter.evaluate(recordBatch, selectionVector);
        long finish = System.nanoTime();
        elapsedTime += (finish - start);
      } finally {
        selectionBuffer.close();
      }
    }

    @Override
    public long getElapsedMillis() {
      return TimeUnit.NANOSECONDS.toMillis(elapsedTime);
    }
  }

  interface DataAndVectorGenerator {

    void writeData(ArrowBuf buffer);

    ValueVector generateOutputVector(int numRowsInBatch);
  }

  class Int32DataAndVectorGenerator implements DataAndVectorGenerator {

    protected final BufferAllocator allocator;
    protected final Random rand;

    Int32DataAndVectorGenerator(BufferAllocator allocator) {
      this.allocator = allocator;
      this.rand = new Random();
    }

    @Override
    public void writeData(ArrowBuf buffer) {
      buffer.writeInt(rand.nextInt());
    }

    @Override
    public ValueVector generateOutputVector(int numRowsInBatch) {
      IntVector intVector = new IntVector(BaseEvaluatorTest.EMPTY_SCHEMA_PATH, allocator);
      intVector.allocateNew(numRowsInBatch);
      return intVector;
    }
  }

  class BoundedInt32DataAndVectorGenerator extends Int32DataAndVectorGenerator {

    private final int upperBound;

    BoundedInt32DataAndVectorGenerator(BufferAllocator allocator, int upperBound) {
      super(allocator);
      this.upperBound = upperBound;
    }

    @Override
    public void writeData(ArrowBuf buffer) {
      buffer.writeInt(rand.nextInt(upperBound));
    }
  }

  protected static final int THOUSAND = 1000;
  protected static final int MILLION = THOUSAND * THOUSAND;

  protected static final String EMPTY_SCHEMA_PATH = "";

  protected BufferAllocator allocator;
  protected ArrowType boolType;
  protected ArrowType int8;
  protected ArrowType int32;
  protected ArrowType int64;
  protected ArrowType float64;

  @BeforeEach
  public void init() {
    allocator = new RootAllocator(Long.MAX_VALUE);
    boolType = new ArrowType.Bool();
    int8 = new ArrowType.Int(8, true);
    int32 = new ArrowType.Int(32, true);
    int64 = new ArrowType.Int(64, true);
    float64 = new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE);
  }

  @AfterEach
  public void tearDown() {
    allocator.close();
  }

  ArrowBuf buf(int length) {
    ArrowBuf buffer = allocator.buffer(length);
    return buffer;
  }

  ArrowBuf buf(byte[] bytes) {
    ArrowBuf buffer = allocator.buffer(bytes.length);
    buffer.writeBytes(bytes);
    return buffer;
  }

  ArrowBuf arrowBufWithAllValid(int size) {
    int bufLen = (size + 7) / 8;
    ArrowBuf buffer = allocator.buffer(bufLen);
    for (int i = 0; i < bufLen; i++) {
      buffer.writeByte(255);
    }

    return buffer;
  }

  ArrowBuf intBuf(int[] ints) {
    ArrowBuf buffer = allocator.buffer(ints.length * 4);
    for (int i = 0; i < ints.length; i++) {
      buffer.writeInt(ints[i]);
    }
    return buffer;
  }

  DecimalVector decimalVector(String[] values, int precision, int scale) {
    DecimalVector vector =
        new DecimalVector("decimal" + Math.random(), allocator, precision, scale);
    vector.allocateNew();
    for (int i = 0; i < values.length; i++) {
      BigDecimal decimal = new BigDecimal(values[i]).setScale(scale);
      vector.setSafe(i, decimal);
    }

    vector.setValueCount(values.length);
    return vector;
  }

  Set decimalSet(String[] values, Integer scale) {
    Set<BigDecimal> decimalSet = new HashSet<>();
    for (int i = 0; i < values.length; i++) {
      decimalSet.add(new BigDecimal(values[i]).setScale(scale));
    }

    return decimalSet;
  }

  VarCharVector varcharVector(String[] values) {
    VarCharVector vector = new VarCharVector("VarCharVector" + Math.random(), allocator);
    vector.allocateNew();
    for (int i = 0; i < values.length; i++) {
      vector.setSafe(i, values[i].getBytes(), 0, values[i].length());
    }

    vector.setValueCount(values.length);
    return vector;
  }

  ArrowBuf longBuf(long[] longs) {
    ArrowBuf buffer = allocator.buffer(longs.length * 8);
    for (int i = 0; i < longs.length; i++) {
      buffer.writeLong(longs[i]);
    }
    return buffer;
  }

  ArrowBuf doubleBuf(double[] data) {
    ArrowBuf buffer = allocator.buffer(data.length * 8);
    for (int i = 0; i < data.length; i++) {
      buffer.writeDouble(data[i]);
    }

    return buffer;
  }

  ArrowBuf stringToMillis(String[] dates) {
    ArrowBuf buffer = allocator.buffer(dates.length * 8);
    for (int i = 0; i < dates.length; i++) {
      Instant instant = Instant.parse(dates[i]);
      buffer.writeLong(instant.toEpochMilli());
    }

    return buffer;
  }

  ArrowBuf stringToDayInterval(String[] values) {
    ArrowBuf buffer = allocator.buffer(values.length * 8);
    for (int i = 0; i < values.length; i++) {
      buffer.writeInt(Integer.parseInt(values[i].split(" ")[0])); // days
      buffer.writeInt(Integer.parseInt(values[i].split(" ")[1])); // millis
    }
    return buffer;
  }

  void releaseRecordBatch(ArrowRecordBatch recordBatch) {
    // There are 2 references to the buffers
    // One in the recordBatch - release that by calling close()
    // One in the allocator - release that explicitly
    List<ArrowBuf> buffers = recordBatch.getBuffers();
    recordBatch.close();
    for (ArrowBuf buf : buffers) {
      buf.getReferenceManager().release();
    }
  }

  void releaseValueVectors(List<ValueVector> valueVectors) {
    for (ValueVector valueVector : valueVectors) {
      valueVector.close();
    }
  }

  void generateData(DataAndVectorGenerator generator, int numRecords, ArrowBuf buffer) {
    for (int i = 0; i < numRecords; i++) {
      generator.writeData(buffer);
    }
  }

  private void generateDataAndEvaluate(
      DataAndVectorGenerator generator,
      BaseEvaluator evaluator,
      int numFields,
      int numRows,
      int maxRowsInBatch,
      int inputFieldSize)
      throws GandivaException, Exception {
    int numRemaining = numRows;
    List<ArrowBuf> inputData = new ArrayList<ArrowBuf>();
    List<ArrowFieldNode> fieldNodes = new ArrayList<ArrowFieldNode>();

    // set the bitmap
    while (numRemaining > 0) {
      int numRowsInBatch = maxRowsInBatch;
      if (numRowsInBatch > numRemaining) {
        numRowsInBatch = numRemaining;
      }

      // generate data
      for (int i = 0; i < numFields; i++) {
        ArrowBuf buf = allocator.buffer(numRowsInBatch * inputFieldSize);
        ArrowBuf validity = arrowBufWithAllValid(maxRowsInBatch);
        generateData(generator, numRowsInBatch, buf);

        fieldNodes.add(new ArrowFieldNode(numRowsInBatch, 0));
        inputData.add(validity);
        inputData.add(buf);
      }

      // create record batch
      ArrowRecordBatch recordBatch = new ArrowRecordBatch(numRowsInBatch, fieldNodes, inputData);

      evaluator.evaluate(recordBatch, allocator);

      // fix numRemaining
      numRemaining -= numRowsInBatch;

      // release refs
      releaseRecordBatch(recordBatch);

      inputData.clear();
      fieldNodes.clear();
    }
  }

  long timedProject(
      DataAndVectorGenerator generator,
      Schema schema,
      List<ExpressionTree> exprs,
      int numRows,
      int maxRowsInBatch,
      int inputFieldSize)
      throws GandivaException, Exception {
    Projector projector = Projector.make(schema, exprs);
    try {
      ProjectEvaluator evaluator =
          new ProjectEvaluator(projector, generator, exprs.size(), maxRowsInBatch);
      generateDataAndEvaluate(
          generator, evaluator, schema.getFields().size(), numRows, maxRowsInBatch, inputFieldSize);
      return evaluator.getElapsedMillis();
    } finally {
      projector.close();
    }
  }

  long timedFilter(
      DataAndVectorGenerator generator,
      Schema schema,
      Condition condition,
      int numRows,
      int maxRowsInBatch,
      int inputFieldSize)
      throws GandivaException, Exception {

    Filter filter = Filter.make(schema, condition);
    try {
      FilterEvaluator evaluator = new FilterEvaluator(filter);
      generateDataAndEvaluate(
          generator, evaluator, schema.getFields().size(), numRows, maxRowsInBatch, inputFieldSize);
      return evaluator.getElapsedMillis();
    } finally {
      filter.close();
    }
  }
}
