package org.apache.arrow.gandiva.evaluator;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import org.apache.arrow.gandiva.exceptions.EvaluatorClosedException;
import org.apache.arrow.gandiva.exceptions.GandivaException;
import org.apache.arrow.gandiva.exceptions.UnsupportedTypeException;
import org.apache.arrow.gandiva.expression.ArrowTypeHelper;
import org.apache.arrow.gandiva.expression.ExpressionTree;
import org.apache.arrow.gandiva.ipc.GandivaTypes;
import org.apache.arrow.gandiva.ipc.GandivaTypes.ExpressionList;
import org.apache.arrow.gandiva.ipc.GandivaTypes.SelectionVectorType;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.vector.BaseVariableWidthVector;
import org.apache.arrow.vector.FixedWidthVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VariableWidthVector;
import org.apache.arrow.vector.ipc.message.ArrowBuffer;
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
import org.apache.arrow.vector.types.pojo.Schema;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class Projector {
    private static final Logger logger = LoggerFactory.getLogger(Projector.class);
    private JniWrapper wrapper;
    private final long moduleId;
    private final Schema schema;
    private final int numExprs;
    private boolean closed;

    private Projector(JniWrapper wrapper, long moduleId, Schema schema, int numExprs) {
        this.wrapper = wrapper;
        this.moduleId = moduleId;
        this.schema = schema;
        this.numExprs = numExprs;
        this.closed = false;
    }

    public static Projector make(Schema schema, List<ExpressionTree> exprs) throws GandivaException {
        return make(schema, exprs, SelectionVectorType.SV_NONE, JniLoader.getDefaultConfiguration());
    }

    public static Projector make(Schema schema, List<ExpressionTree> exprs, ConfigurationBuilder.ConfigOptions configOptions) throws GandivaException {
        return make(schema, exprs, SelectionVectorType.SV_NONE, JniLoader.getConfiguration(configOptions));
    }

    /** @deprecated */
    @Deprecated
    public static Projector make(Schema schema, List<ExpressionTree> exprs, boolean optimize) throws GandivaException {
        return make(schema, exprs, SelectionVectorType.SV_NONE, JniLoader.getConfiguration((new ConfigurationBuilder.ConfigOptions()).withOptimize(optimize)));
    }

    public static Projector make(Schema schema, List<ExpressionTree> exprs, GandivaTypes.SelectionVectorType selectionVectorType) throws GandivaException {
        return make(schema, exprs, selectionVectorType, JniLoader.getDefaultConfiguration());
    }

    public static Projector make(Schema schema, List<ExpressionTree> exprs, GandivaTypes.SelectionVectorType selectionVectorType, ConfigurationBuilder.ConfigOptions configOptions) throws GandivaException {
        return make(schema, exprs, selectionVectorType, JniLoader.getConfiguration(configOptions));
    }

    /** @deprecated */
    @Deprecated
    public static Projector make(Schema schema, List<ExpressionTree> exprs, GandivaTypes.SelectionVectorType selectionVectorType, boolean optimize) throws GandivaException {
        return make(schema, exprs, selectionVectorType, JniLoader.getConfiguration((new ConfigurationBuilder.ConfigOptions()).withOptimize(optimize)));
    }

    public static Projector make(Schema schema, List<ExpressionTree> exprs, GandivaTypes.SelectionVectorType selectionVectorType, long configurationId) throws GandivaException {
        return make(schema, exprs, selectionVectorType, configurationId, (JavaSecondaryCacheInterface)null);
    }

    public static Projector make(Schema schema, List<ExpressionTree> exprs, ConfigurationBuilder.ConfigOptions configOptions, JavaSecondaryCacheInterface secondaryCache) throws GandivaException {
        return make(schema, exprs, SelectionVectorType.SV_NONE, JniLoader.getConfiguration(configOptions), secondaryCache);
    }

    public static Projector make(Schema schema, List<ExpressionTree> exprs, GandivaTypes.SelectionVectorType selectionVectorType, long configurationId, JavaSecondaryCacheInterface secondaryCache) throws GandivaException {
        GandivaTypes.ExpressionList.Builder builder = ExpressionList.newBuilder();
        Iterator var7 = exprs.iterator();

        while(var7.hasNext()) {
            ExpressionTree expr = (ExpressionTree)var7.next();
            builder.addExprs(expr.toProtobuf());
        }

        GandivaTypes.Schema schemaBuf = ArrowTypeHelper.arrowSchemaToProtobuf(schema);
        JniWrapper wrapper = JniLoader.getInstance().getWrapper();
        long moduleId = wrapper.buildProjector(secondaryCache, schemaBuf.toByteArray(), builder.build().toByteArray(), selectionVectorType.getNumber(), configurationId);
        logger.debug("Created module for the projector with id {}", moduleId);
        return new Projector(wrapper, moduleId, schema, exprs.size());
    }

    public void evaluate(ArrowRecordBatch recordBatch, List<ValueVector> outColumns) throws GandivaException {
        this.evaluate(recordBatch.getLength(), recordBatch.getBuffers(), recordBatch.getBuffersLayout(), SelectionVectorType.SV_NONE.getNumber(), recordBatch.getLength(), 0L, 0L, outColumns);
    }

    public void evaluate(int numRows, List<ArrowBuf> buffers, List<ValueVector> outColumns) throws GandivaException {
        List<ArrowBuffer> buffersLayout = new ArrayList();
        long offset = 0L;

        long size;
        for(Iterator var7 = buffers.iterator(); var7.hasNext(); offset += size) {
            ArrowBuf arrowBuf = (ArrowBuf)var7.next();
            size = arrowBuf.readableBytes();
            buffersLayout.add(new ArrowBuffer(offset, size));
        }

        this.evaluate(numRows, buffers, buffersLayout, SelectionVectorType.SV_NONE.getNumber(), numRows, 0L, 0L, outColumns);
    }

    public void evaluate(ArrowRecordBatch recordBatch, SelectionVector selectionVector, List<ValueVector> outColumns) throws GandivaException {
        this.evaluate(recordBatch.getLength(), recordBatch.getBuffers(), recordBatch.getBuffersLayout(), selectionVector.getType().getNumber(), selectionVector.getRecordCount(), selectionVector.getBuffer().memoryAddress(), selectionVector.getBuffer().capacity(), outColumns);
    }

    public void evaluate(int numRows, List<ArrowBuf> buffers, SelectionVector selectionVector, List<ValueVector> outColumns) throws GandivaException {
        List<ArrowBuffer> buffersLayout = new ArrayList();
        long offset = 0L;

        long size;
        for(Iterator var8 = buffers.iterator(); var8.hasNext(); offset += size) {
            ArrowBuf arrowBuf = (ArrowBuf)var8.next();
            size = arrowBuf.readableBytes();
            buffersLayout.add(new ArrowBuffer(offset, size));
        }

        this.evaluate(numRows, buffers, buffersLayout, selectionVector.getType().getNumber(), selectionVector.getRecordCount(), selectionVector.getBuffer().memoryAddress(), selectionVector.getBuffer().capacity(), outColumns);
    }

    private void evaluate(int numRows, List<ArrowBuf> buffers, List<ArrowBuffer> buffersLayout, int selectionVectorType, int selectionVectorRecordCount, long selectionVectorAddr, long selectionVectorSize, List<ValueVector> outColumns) throws GandivaException {
        if (this.closed) {
            throw new EvaluatorClosedException();
        } else if (this.numExprs != outColumns.size()) {
            logger.info("Expected " + this.numExprs + " columns, got " + outColumns.size());
            throw new GandivaException("Incorrect number of columns for the output vector");
        } else {
            long[] bufAddrs = new long[buffers.size()];
            long[] bufSizes = new long[buffers.size()];
            int idx = 0;

            Iterator var14;
            ArrowBuf buf;
            for(var14 = buffers.iterator(); var14.hasNext(); bufAddrs[idx++] = buf.memoryAddress()) {
                buf = (ArrowBuf)var14.next();
            }

            idx = 0;

            ArrowBuffer bufLayout;
            for(var14 = buffersLayout.iterator(); var14.hasNext(); bufSizes[idx++] = bufLayout.getSize()) {
                bufLayout = (ArrowBuffer)var14.next();
            }

            boolean hasVariableWidthColumns = false;
            BaseVariableWidthVector[] resizableVectors = new BaseVariableWidthVector[outColumns.size()];
            long[] outAddrs = new long[3 * outColumns.size()];
            long[] outSizes = new long[3 * outColumns.size()];
            idx = 0;
            int outColumnIdx = 0;

            for(Iterator var19 = outColumns.iterator(); var19.hasNext(); ++outColumnIdx) {
                ValueVector valueVector = (ValueVector)var19.next();
                boolean isFixedWith = valueVector instanceof FixedWidthVector;
                boolean isVarWidth = valueVector instanceof VariableWidthVector;
                if (!isFixedWith && !isVarWidth) {
                    throw new UnsupportedTypeException("Unsupported value vector type " + valueVector.getField().getFieldType());
                }

                outAddrs[idx] = valueVector.getValidityBuffer().memoryAddress();
                outSizes[idx++] = valueVector.getValidityBuffer().capacity();
                if (isVarWidth) {
                    outAddrs[idx] = valueVector.getOffsetBuffer().memoryAddress();
                    outSizes[idx++] = valueVector.getOffsetBuffer().capacity();
                    hasVariableWidthColumns = true;
                    resizableVectors[outColumnIdx] = (BaseVariableWidthVector)valueVector;
                }

                outAddrs[idx] = valueVector.getDataBuffer().memoryAddress();
                outSizes[idx++] = valueVector.getDataBuffer().capacity();
                valueVector.setValueCount(selectionVectorRecordCount);
            }

            this.wrapper.evaluateProjector(hasVariableWidthColumns ? new VectorExpander(resizableVectors) : null, this.moduleId, numRows, bufAddrs, bufSizes, selectionVectorType, selectionVectorRecordCount, selectionVectorAddr, selectionVectorSize, outAddrs, outSizes);
        }
    }

    public void close() throws GandivaException {
        if (!this.closed) {
            this.wrapper.closeProjector(this.moduleId);
            this.closed = true;
        }
    }
}