/*
 * 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.vector.dictionary;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.util.hash.ArrowBufHasher;
import org.apache.arrow.memory.util.hash.SimpleHasher;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.BaseIntVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.util.TransferPair;

/**
 * Sub fields encoder/decoder for Dictionary encoded {@link StructVector}. Notes that child vectors
 * within struct vector can either be dictionary encodable or not.
 */
public class StructSubfieldEncoder {

  private final BufferAllocator allocator;

  private final DictionaryProvider.MapDictionaryProvider provider;
  private final Map<Long, DictionaryHashTable> dictionaryIdToHashTable;

  /** Construct an instance. */
  public StructSubfieldEncoder(
      BufferAllocator allocator, DictionaryProvider.MapDictionaryProvider provider) {
    this(allocator, provider, SimpleHasher.INSTANCE);
  }

  /** Construct an instance. */
  public StructSubfieldEncoder(
      BufferAllocator allocator,
      DictionaryProvider.MapDictionaryProvider provider,
      ArrowBufHasher hasher) {

    this.allocator = allocator;
    this.provider = provider;

    this.dictionaryIdToHashTable = new HashMap<>();

    provider
        .getDictionaryIds()
        .forEach(
            id ->
                dictionaryIdToHashTable.put(
                    id, new DictionaryHashTable(provider.lookup(id).getVector(), hasher)));
  }

  private static FieldVector getChildVector(StructVector vector, int index) {
    return vector.getChildrenFromFields().get(index);
  }

  private static StructVector cloneVector(StructVector vector, BufferAllocator allocator) {

    final FieldType fieldType = vector.getField().getFieldType();
    StructVector cloned =
        (StructVector)
            fieldType.createNewSingleVector(
                vector.getField().getName(), allocator, /*schemaCallback=*/ null);

    final ArrowFieldNode fieldNode =
        new ArrowFieldNode(vector.getValueCount(), vector.getNullCount());
    cloned.loadFieldBuffers(fieldNode, vector.getFieldBuffers());

    return cloned;
  }

  /**
   * Dictionary encodes subfields for complex vector with a provided dictionary. The dictionary must
   * contain all values in the sub fields vector.
   *
   * @param vector vector to encode
   * @param columnToDictionaryId the mappings between child vector index and dictionary id. A null
   *     dictionary id indicates the child vector is not encodable.
   * @return dictionary encoded vector
   */
  public StructVector encode(StructVector vector, Map<Integer, Long> columnToDictionaryId) {
    final int valueCount = vector.getValueCount();
    final int childCount = vector.getChildrenFromFields().size();

    List<Field> childrenFields = new ArrayList<>();

    // initialize child fields.
    for (int i = 0; i < childCount; i++) {
      FieldVector childVector = getChildVector(vector, i);
      Long dictionaryId = columnToDictionaryId.get(i);
      // A null dictionaryId indicates the child vector shouldn't be encoded.
      if (dictionaryId == null) {
        childrenFields.add(childVector.getField());
      } else {
        Dictionary dictionary = provider.lookup(dictionaryId);
        Preconditions.checkNotNull(dictionary, "Dictionary not found with id:" + dictionaryId);
        FieldType indexFieldType =
            new FieldType(
                childVector.getField().isNullable(),
                dictionary.getEncoding().getIndexType(),
                dictionary.getEncoding());
        childrenFields.add(
            new Field(childVector.getField().getName(), indexFieldType, /*children=*/ null));
      }
    }

    // clone list vector and initialize data vector
    StructVector encoded = cloneVector(vector, allocator);
    try {
      encoded.initializeChildrenFromFields(childrenFields);
      encoded.setValueCount(valueCount);

      for (int index = 0; index < childCount; index++) {
        FieldVector childVector = getChildVector(vector, index);
        FieldVector encodedChildVector = getChildVector(encoded, index);
        Long dictionaryId = columnToDictionaryId.get(index);
        if (dictionaryId != null) {
          BaseIntVector indices = (BaseIntVector) encodedChildVector;
          DictionaryEncoder.buildIndexVector(
              childVector, indices, dictionaryIdToHashTable.get(dictionaryId), 0, valueCount);
        } else {
          childVector.makeTransferPair(encodedChildVector).splitAndTransfer(0, valueCount);
        }
      }

      return encoded;
    } catch (Exception e) {
      AutoCloseables.close(e, encoded);
      throw e;
    }
  }

  /**
   * Decodes a dictionary subfields encoded vector using the provided dictionary.
   *
   * <p>{@link StructSubfieldEncoder#decode(StructVector, DictionaryProvider.MapDictionaryProvider,
   * BufferAllocator)} should be used instead if only decoding is required as it can avoid building
   * the {@link DictionaryHashTable} which only makes sense when encoding.
   *
   * @param vector dictionary encoded vector, its child vector must be int type
   * @return vector with values restored from dictionary
   */
  public StructVector decode(StructVector vector) {
    return decode(vector, provider, allocator);
  }

  /**
   * Decodes a dictionary subfields encoded vector using the provided dictionary.
   *
   * @param vector dictionary encoded vector, its data vector must be int type
   * @param provider dictionary provider used to decode the values
   * @param allocator allocator the decoded values use
   * @return vector with values restored from dictionary
   */
  public static StructVector decode(
      StructVector vector,
      DictionaryProvider.MapDictionaryProvider provider,
      BufferAllocator allocator) {
    final int valueCount = vector.getValueCount();
    final int childCount = vector.getChildrenFromFields().size();

    // clone list vector and initialize child vectors
    StructVector decoded = cloneVector(vector, allocator);
    try {
      List<Field> childFields = new ArrayList<>();
      for (int i = 0; i < childCount; i++) {
        FieldVector childVector = getChildVector(vector, i);
        Dictionary dictionary = getChildVectorDictionary(childVector, provider);
        // childVector is not encoded.
        if (dictionary == null) {
          childFields.add(childVector.getField());
        } else {
          childFields.add(dictionary.getVector().getField());
        }
      }
      decoded.initializeChildrenFromFields(childFields);
      decoded.setValueCount(valueCount);

      for (int index = 0; index < childCount; index++) {
        // get child vector
        FieldVector childVector = getChildVector(vector, index);
        FieldVector decodedChildVector = getChildVector(decoded, index);
        Dictionary dictionary = getChildVectorDictionary(childVector, provider);
        if (dictionary == null) {
          childVector.makeTransferPair(decodedChildVector).splitAndTransfer(0, valueCount);
        } else {
          TransferPair transfer = dictionary.getVector().makeTransferPair(decodedChildVector);
          BaseIntVector indices = (BaseIntVector) childVector;

          DictionaryEncoder.retrieveIndexVector(indices, transfer, valueCount, 0, valueCount);
        }
      }

      return decoded;
    } catch (Exception e) {
      AutoCloseables.close(e, decoded);
      throw e;
    }
  }

  /** Get the child vector dictionary, return null if not dictionary encoded. */
  private static Dictionary getChildVectorDictionary(
      FieldVector childVector, DictionaryProvider.MapDictionaryProvider provider) {
    DictionaryEncoding dictionaryEncoding = childVector.getField().getDictionary();
    if (dictionaryEncoding != null) {
      Dictionary dictionary = provider.lookup(dictionaryEncoding.getId());
      Preconditions.checkNotNull(dictionary, "Dictionary not found with id:" + dictionary);
      return dictionary;
    }
    return null;
  }
}
