/*
 * 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.carbondata.sdk.file.arrow;

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

import org.apache.carbondata.sdk.file.Schema;

import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.BitVector;
import org.apache.arrow.vector.DateDayVector;
import org.apache.arrow.vector.DecimalVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.Float4Vector;
import org.apache.arrow.vector.Float8Vector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.SmallIntVector;
import org.apache.arrow.vector.TimeStampMicroTZVector;
import org.apache.arrow.vector.TinyIntVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;

public class ArrowWriter {

  private VectorSchemaRoot root;

  private ArrowFieldWriter[] children;

  private int count;

  public void write(Object[] data) {
    int i = 0;
    while (i < children.length) {
      children[i].write(data[i], i);
      i += 1;
    }
    count += 1;
  }

  public void finish() {
    root.setRowCount(count);
    for (int i = 0; i < children.length; i++) {
      children[i].finish();
    }
  }

  public void reset() {
    root.setRowCount(0);
    count = 0;
    for (int i = 0; i < children.length; i++) {
      children[i].reset();
    }
  }

  private ArrowWriter(VectorSchemaRoot root, ArrowFieldWriter[] children) {
    this.root = root;
    this.children = children;
  }

  public static ArrowWriter create(Schema schema, String timeZoneId) {
    org.apache.arrow.vector.types.pojo.Schema arrowSchema =
        ArrowUtils.toArrowSchema(schema, timeZoneId);
    VectorSchemaRoot root = VectorSchemaRoot.create(arrowSchema, ArrowUtils.ROOT_ALLOCATOR);
    return create(root);
  }

  public static ArrowWriter create(VectorSchemaRoot root) {
    final List<FieldVector> fieldVectors = root.getFieldVectors();
    ArrowFieldWriter[] fieldWriters = new ArrowFieldWriter[fieldVectors.size()];
    int i = 0;
    for (FieldVector fieldVector : fieldVectors) {
      fieldWriters[i] = createFieldWriter(fieldVector);
      i++;
    }
    return new ArrowWriter(root, fieldWriters);
  }

  private static ArrowFieldWriter createFieldWriter(ValueVector valueVector) {
    if (valueVector instanceof BitVector) {
      return new BooleanWriter((BitVector) valueVector);
    } else if (valueVector instanceof TinyIntVector) {
      return new ByteWriter((TinyIntVector) valueVector);
    } else if (valueVector instanceof SmallIntVector) {
      return new ShortWriter((SmallIntVector) valueVector);
    } else if (valueVector instanceof IntVector) {
      return new IntWriter((IntVector) valueVector);
    } else if (valueVector instanceof BigIntVector) {
      return new LongWriter((BigIntVector) valueVector);
    } else if (valueVector instanceof DecimalVector) {
      DecimalVector decimalVector = (DecimalVector) valueVector;
      final Field field = decimalVector.getField();
      ArrowType.Decimal c = (ArrowType.Decimal) field.getType();
      return new DecimalWriter((DecimalVector) valueVector, c.getPrecision(), c.getScale());
    } else if (valueVector instanceof VarCharVector) {
      return new StringWriter((VarCharVector) valueVector);
    } else if (valueVector instanceof Float4Vector) {
      return new FloatWriter((Float4Vector) valueVector);
    } else if (valueVector instanceof Float8Vector) {
      return new DoubleWriter((Float8Vector) valueVector);
    } else if (valueVector instanceof ListVector) {
      ArrowFieldWriter elementVector =
          createFieldWriter(((ListVector) valueVector).getDataVector());
      return new ArrayWriter((ListVector) valueVector, elementVector);
    } else if (valueVector instanceof StructVector) {
      StructVector s = (StructVector) valueVector;
      List<ArrowFieldWriter> arrowFieldWriters = new ArrayList<>();
      for (int i = 0; i < s.size(); i++) {
        arrowFieldWriters.add(createFieldWriter(s.getChildByOrdinal(i)));
      }
      return new StructWriter(s,
          arrowFieldWriters.toArray(new ArrowFieldWriter[arrowFieldWriters.size()]));
    } else if (valueVector instanceof VarBinaryVector) {
      return new BinaryWriter((VarBinaryVector) valueVector);
    } else if (valueVector instanceof DateDayVector) {
      return new DateWriter((DateDayVector) valueVector);
    } else if (valueVector instanceof TimeStampMicroTZVector) {
      return new TimeStampWriter((TimeStampMicroTZVector) valueVector);
    } else {
      throw new UnsupportedOperationException("Invalid data type");
    }
  }
}

