/*
 * 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.store.hive.writers;

import java.util.List;
import java.util.function.BiFunction;
import java.util.function.Function;

import org.apache.drill.common.exceptions.UserException;
import org.apache.drill.exec.ops.BufferManager;
import org.apache.drill.exec.store.hive.writers.complex.HiveListWriter;
import org.apache.drill.exec.store.hive.writers.complex.HiveMapWriter;
import org.apache.drill.exec.store.hive.writers.complex.HiveStructWriter;
import org.apache.drill.exec.store.hive.writers.complex.HiveUnionWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveBinaryWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveBooleanWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveByteWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveCharWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveDateWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveDecimalWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveDoubleWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveFloatWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveIntWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveLongWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveShortWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveStringWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveTimestampWriter;
import org.apache.drill.exec.store.hive.writers.primitive.HiveVarCharWriter;
import org.apache.drill.exec.vector.complex.DictVector;
import org.apache.drill.exec.vector.complex.impl.SingleMapWriter;
import org.apache.drill.exec.vector.complex.impl.UnionVectorWriter;
import org.apache.drill.exec.vector.complex.writer.BaseWriter;
import org.apache.drill.exec.vector.complex.writer.BaseWriter.ListWriter;
import org.apache.drill.exec.vector.complex.writer.BaseWriter.MapWriter;
import org.apache.drill.exec.vector.complex.writer.BigIntWriter;
import org.apache.drill.exec.vector.complex.writer.BitWriter;
import org.apache.drill.exec.vector.complex.writer.DateWriter;
import org.apache.drill.exec.vector.complex.writer.Float4Writer;
import org.apache.drill.exec.vector.complex.writer.Float8Writer;
import org.apache.drill.exec.vector.complex.writer.IntWriter;
import org.apache.drill.exec.vector.complex.writer.TimeStampWriter;
import org.apache.drill.exec.vector.complex.writer.VarBinaryWriter;
import org.apache.drill.exec.vector.complex.writer.VarCharWriter;
import org.apache.drill.exec.vector.complex.writer.VarDecimalWriter;
import org.apache.hadoop.hive.serde2.objectinspector.ListObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.MapObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.StructField;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.UnionObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.BinaryObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.BooleanObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.ByteObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.DateObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.DoubleObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.FloatObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.HiveCharObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.HiveDecimalObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.HiveVarcharObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.IntObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.LongObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.ShortObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.StringObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.TimestampObjectInspector;
import org.apache.hadoop.hive.serde2.typeinfo.DecimalTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.ListTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.MapTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.PrimitiveTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoUtils;
import org.apache.hadoop.hive.serde2.typeinfo.UnionTypeInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static org.apache.drill.exec.store.hive.HiveUtilities.throwUnsupportedHiveDataTypeError;

/**
 * Factory used by reader to create Hive writers for columns.
 */
public final class HiveValueWriterFactory {

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

  /**
   * Buffer manager used to create buffers for Hive writers for reading data
   * to buffer than from buffer to vector if needed.
   */
  private final BufferManager bufferManager;

  /**
   * Used to manage and create column writers.
   */
  private final SingleMapWriter rootWriter;

  public HiveValueWriterFactory(BufferManager bufferManager, SingleMapWriter rootWriter) {
    this.bufferManager = bufferManager;
    this.rootWriter = rootWriter;
  }

  /**
   * Method that will be called once for each column in reader to initialize column writer.
   *
   * @param columnName name of column for writer
   * @param fieldRef   metadata about field type
   * @return instance of writer for column
   */
  public HiveValueWriter createHiveColumnValueWriter(String columnName, StructField fieldRef) {
    ObjectInspector objectInspector = fieldRef.getFieldObjectInspector();
    final TypeInfo typeInfo = TypeInfoUtils.getTypeInfoFromTypeString(objectInspector.getTypeName());
    return createHiveValueWriter(columnName, typeInfo, objectInspector, rootWriter);
  }

  private HiveValueWriter createHiveValueWriter(String columnName, TypeInfo typeInfo, ObjectInspector objectInspector, BaseWriter parentWriter) {
    switch (typeInfo.getCategory()) {
      case PRIMITIVE:
        return createPrimitiveHiveValueWriter(columnName, (PrimitiveTypeInfo) typeInfo, objectInspector, parentWriter);
      case LIST: {
        TypeInfo elemTypeInfo = ((ListTypeInfo) typeInfo).getListElementTypeInfo();
        ListObjectInspector listInspector = (ListObjectInspector) objectInspector;
        ObjectInspector elementInspector = listInspector.getListElementObjectInspector();
        ListWriter listWriter = extractWriter(columnName, parentWriter,
            MapWriter::list, ListWriter::list, UnionVectorWriter::list);
        HiveValueWriter elementValueWriter = createHiveValueWriter(null, elemTypeInfo, elementInspector, listWriter);
        return new HiveListWriter(listInspector, listWriter, elementValueWriter);
      }
      case STRUCT: {
        StructObjectInspector structInspector = (StructObjectInspector) objectInspector;
        StructField[] structFields = structInspector.getAllStructFieldRefs().toArray(new StructField[0]);
        MapWriter structWriter = extractWriter(columnName, parentWriter,
            MapWriter::map, ListWriter::map, UnionVectorWriter::map);
        HiveValueWriter[] structFieldWriters = new HiveValueWriter[structFields.length];
        for (int fieldIdx = 0; fieldIdx < structFields.length; fieldIdx++) {
          StructField field = structFields[fieldIdx];
          ObjectInspector fieldObjectInspector = field.getFieldObjectInspector();
          TypeInfo fieldTypeInfo = TypeInfoUtils.getTypeInfoFromTypeString(fieldObjectInspector.getTypeName());
          structFieldWriters[fieldIdx] = createHiveValueWriter(field.getFieldName(), fieldTypeInfo, fieldObjectInspector, structWriter);
        }
        return new HiveStructWriter(structInspector, structFields, structFieldWriters, structWriter);
      }
      case MAP: {
        MapTypeInfo mapTypeInfo = (MapTypeInfo) typeInfo;
        PrimitiveTypeInfo keyTypeInfo = (PrimitiveTypeInfo) mapTypeInfo.getMapKeyTypeInfo();
        TypeInfo valueTypeInfo = mapTypeInfo.getMapValueTypeInfo();
        MapObjectInspector mapObjectInspector = (MapObjectInspector) objectInspector;
        ObjectInspector keyInspector = mapObjectInspector.getMapKeyObjectInspector();
        ObjectInspector valueInspector = mapObjectInspector.getMapValueObjectInspector();
        BaseWriter.DictWriter dictWriter = extractWriter(columnName, parentWriter,
            MapWriter::dict, ListWriter::dict, UnionVectorWriter::dict);
        HiveValueWriter mapKeyWriter = createPrimitiveHiveValueWriter(DictVector.FIELD_KEY_NAME, keyTypeInfo, keyInspector, dictWriter);
        HiveValueWriter mapValueWriter = createHiveValueWriter(DictVector.FIELD_VALUE_NAME, valueTypeInfo, valueInspector, dictWriter);
        return new HiveMapWriter(mapObjectInspector, dictWriter, mapKeyWriter, mapValueWriter);
      }
      case UNION: {
        UnionObjectInspector unionObjectInspector = (UnionObjectInspector) objectInspector;
        UnionTypeInfo unionTypeInfo = (UnionTypeInfo) typeInfo;
        List<TypeInfo> unionFieldsTypeInfo = unionTypeInfo.getAllUnionObjectTypeInfos();
        List<ObjectInspector> objectInspectors = unionObjectInspector.getObjectInspectors();
        UnionVectorWriter unionWriter = extractWriter(columnName, parentWriter,
            MapWriter::union, ListWriter::union, UnionVectorWriter::union);
        HiveValueWriter[] unionFieldWriters = new HiveValueWriter[unionFieldsTypeInfo.size()];
        for (int tag = 0; tag < unionFieldsTypeInfo.size(); tag++) {
          ObjectInspector unionFieldInspector = objectInspectors.get(tag);
          TypeInfo unionFieldTypeInfo = unionFieldsTypeInfo.get(tag);
          HiveValueWriter unionValueWriter = createHiveValueWriter(null, unionFieldTypeInfo, unionFieldInspector, unionWriter);
          unionFieldWriters[tag] = unionValueWriter;
        }
        return new HiveUnionWriter(unionFieldWriters, unionObjectInspector);
      }
    }
    throwUnsupportedHiveDataTypeError(typeInfo.getCategory().toString());
    return null;
  }

  /**
   * Creates writer for primitive type value.
   *
   * @param name         column name or null if nested
   * @param typeInfo     column type used to distinguish returned writers
   * @param inspector    inspector of column values
   * @param parentWriter parent writer used to extract writer for primitive
   * @return appropriate instance of HiveValueWriter for column containing primitive scalar
   */
  private HiveValueWriter createPrimitiveHiveValueWriter(String name, PrimitiveTypeInfo typeInfo, ObjectInspector inspector, BaseWriter parentWriter) {
    switch (typeInfo.getPrimitiveCategory()) {
      case BINARY: {
        VarBinaryWriter writer = extractWriter(name, parentWriter,
            MapWriter::varBinary, ListWriter::varBinary, UnionVectorWriter::varBinary);
        return new HiveBinaryWriter((BinaryObjectInspector) inspector, writer, bufferManager.getManagedBuffer());
      }
      case BOOLEAN: {
        BitWriter writer = extractWriter(name, parentWriter,
            MapWriter::bit, ListWriter::bit, UnionVectorWriter::bit);
        return new HiveBooleanWriter((BooleanObjectInspector) inspector, writer);
      }
      case BYTE: {
        IntWriter writer = extractWriter(name, parentWriter,
            MapWriter::integer, ListWriter::integer, UnionVectorWriter::integer);
        return new HiveByteWriter((ByteObjectInspector) inspector, writer);
      }
      case DOUBLE: {
        Float8Writer writer = extractWriter(name, parentWriter,
            MapWriter::float8, ListWriter::float8, UnionVectorWriter::float8);
        return new HiveDoubleWriter((DoubleObjectInspector) inspector, writer);
      }
      case FLOAT: {
        Float4Writer writer = extractWriter(name, parentWriter,
            MapWriter::float4, ListWriter::float4, UnionVectorWriter::float4);
        return new HiveFloatWriter((FloatObjectInspector) inspector, writer);
      }
      case INT: {
        IntWriter writer = extractWriter(name, parentWriter,
            MapWriter::integer, ListWriter::integer, UnionVectorWriter::integer);
        return new HiveIntWriter((IntObjectInspector) inspector, writer);
      }
      case LONG: {
        BigIntWriter writer = extractWriter(name, parentWriter,
            MapWriter::bigInt, ListWriter::bigInt, UnionVectorWriter::bigInt);
        return new HiveLongWriter((LongObjectInspector) inspector, writer);
      }
      case SHORT: {
        IntWriter writer = extractWriter(name, parentWriter,
            MapWriter::integer, ListWriter::integer, UnionVectorWriter::integer);
        return new HiveShortWriter((ShortObjectInspector) inspector, writer);
      }
      case STRING: {
        VarCharWriter writer = extractWriter(name, parentWriter,
            MapWriter::varChar, ListWriter::varChar, UnionVectorWriter::varChar);
        return new HiveStringWriter((StringObjectInspector) inspector, writer, bufferManager.getManagedBuffer());
      }
      case VARCHAR: {
        VarCharWriter writer = extractWriter(name, parentWriter,
            MapWriter::varChar, ListWriter::varChar, UnionVectorWriter::varChar);
        return new HiveVarCharWriter((HiveVarcharObjectInspector) inspector, writer, bufferManager.getManagedBuffer());
      }
      case TIMESTAMP: {
        TimeStampWriter writer = extractWriter(name, parentWriter,
            MapWriter::timeStamp, ListWriter::timeStamp, UnionVectorWriter::timeStamp);
        return new HiveTimestampWriter((TimestampObjectInspector) inspector, writer);
      }
      case DATE: {
        DateWriter writer = extractWriter(name, parentWriter,
            MapWriter::date, ListWriter::date, UnionVectorWriter::date);
        return new HiveDateWriter((DateObjectInspector) inspector, writer);
      }
      case CHAR: {
        VarCharWriter writer = extractWriter(name, parentWriter,
            MapWriter::varChar, ListWriter::varChar, UnionVectorWriter::varChar);
        return new HiveCharWriter((HiveCharObjectInspector) inspector, writer, bufferManager.getManagedBuffer());
      }
      case DECIMAL: {
        DecimalTypeInfo decimalType = (DecimalTypeInfo) typeInfo;
        int scale = decimalType.getScale();
        int precision = decimalType.getPrecision();
        VarDecimalWriter writer = extractWriter(name, parentWriter,
            (mapWriter, key) -> mapWriter.varDecimal(key, precision, scale),
            listWriter -> listWriter.varDecimal(precision, scale),
            unionWriter -> unionWriter.varDecimal(precision, scale));
        return new HiveDecimalWriter((HiveDecimalObjectInspector) inspector, writer, scale);
      }
      default:
        throw UserException.unsupportedError()
            .message("Unsupported primitive data type '%s'", typeInfo.getPrimitiveCategory())
            .build(logger);
    }
  }

  /**
   * Used to extract child writer from parent writer, assuming that parent writer may be instance of
   * {@link MapWriter} or {@link ListWriter}
   *
   * @param name         column or struct field name
   * @param parentWriter parent writer used for getting child writer
   * @param fromMap      function for extracting writer from map parent writer
   * @param fromList     function for extracting writer from list parent writer
   * @param fromUnion    function for extracting writer from union parent writer
   * @param <T>          type of extracted writer
   * @return writer extracted using either fromMap or fromList function
   */
  private static <T> T extractWriter(String name, BaseWriter parentWriter,
                                     BiFunction<MapWriter, String, T> fromMap,
                                     Function<ListWriter, T> fromList,
                                     Function<UnionVectorWriter, T> fromUnion) {
    if (parentWriter instanceof MapWriter && name != null) {
      return fromMap.apply((MapWriter) parentWriter, name);
    } else if (parentWriter instanceof UnionVectorWriter) {
      return fromUnion.apply((UnionVectorWriter) parentWriter);
    } else if (parentWriter instanceof ListWriter) {
      return fromList.apply((ListWriter) parentWriter);
    } else {
      throw new IllegalStateException(String.format("Parent writer with type [%s] is unsupported", parentWriter.getClass()));
    }
  }

}
