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

import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;

import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.BitVectorHelper;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.util.CallBack;
import org.apache.arrow.vector.util.PromotableMultiMapWithOrdinal;
import org.apache.arrow.vector.util.ValueVectorUtility;

/**
 * Base class for StructVectors. Currently used by NonNullableStructVector
 */
public abstract class AbstractStructVector extends AbstractContainerVector {
  private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(AbstractContainerVector.class);
  private static final String STRUCT_CONFLICT_POLICY_ENV = "ARROW_STRUCT_CONFLICT_POLICY";
  private static final String STRUCT_CONFLICT_POLICY_JVM = "arrow.struct.conflict.policy";
  private static final ConflictPolicy DEFAULT_CONFLICT_POLICY;
  // Maintains a map with key as field name and value is the vector itself
  private final PromotableMultiMapWithOrdinal<String, FieldVector> vectors;
  protected final boolean allowConflictPolicyChanges;
  private ConflictPolicy conflictPolicy;


  static {
    String conflictPolicyStr = System.getProperty(STRUCT_CONFLICT_POLICY_JVM,
        ConflictPolicy.CONFLICT_REPLACE.toString());
    if (conflictPolicyStr == null) {
      conflictPolicyStr = System.getenv(STRUCT_CONFLICT_POLICY_ENV);
    }
    ConflictPolicy conflictPolicy;
    try {
      conflictPolicy = ConflictPolicy.valueOf(conflictPolicyStr.toUpperCase());
    } catch (Exception e) {
      conflictPolicy = ConflictPolicy.CONFLICT_REPLACE;
    }
    DEFAULT_CONFLICT_POLICY = conflictPolicy;
  }

  /**
   * Policy to determine how to react when duplicate columns are encountered.
   */
  public enum ConflictPolicy {
    // Ignore the conflict and append the field. This is the default behaviour
    CONFLICT_APPEND,
    // Keep the existing field and ignore the newer one.
    CONFLICT_IGNORE,
    // Replace the existing field with the newer one.
    CONFLICT_REPLACE,
    // Refuse the new field and error out.
    CONFLICT_ERROR
  }

  /**
   *  Base constructor that sets default conflict policy to APPEND.
   */
  protected AbstractStructVector(String name,
                                 BufferAllocator allocator,
                                 CallBack callBack,
                                 ConflictPolicy conflictPolicy,
                                 boolean allowConflictPolicyChanges) {
    super(name, allocator, callBack);
    this.conflictPolicy = conflictPolicy == null ? DEFAULT_CONFLICT_POLICY : conflictPolicy;
    this.vectors = new PromotableMultiMapWithOrdinal<>(allowConflictPolicyChanges, this.conflictPolicy);
    this.allowConflictPolicyChanges = allowConflictPolicyChanges;
  }

  /**
   * Set conflict policy and return last conflict policy state.
   */
  public ConflictPolicy setConflictPolicy(ConflictPolicy conflictPolicy) {
    ConflictPolicy tmp = this.conflictPolicy;
    this.conflictPolicy = conflictPolicy;
    this.vectors.setConflictPolicy(conflictPolicy);
    return tmp;
  }

  public ConflictPolicy getConflictPolicy() {
    return conflictPolicy;
  }

  @Override
  public void close() {
    for (final ValueVector valueVector : vectors.values()) {
      valueVector.close();
    }
    vectors.clear();

    super.close();
  }

  @Override
  public boolean allocateNewSafe() {
    /* boolean to keep track if all the memory allocation were successful
     * Used in the case of composite vectors when we need to allocate multiple
     * buffers for multiple vectors. If one of the allocations failed we need to
     * clear all the memory that we allocated
     */
    boolean success = false;
    try {
      for (final ValueVector v : vectors.values()) {
        if (!v.allocateNewSafe()) {
          return false;
        }
      }
      success = true;
    } finally {
      if (!success) {
        clear();
      }
    }
    return true;
  }

  @Override
  public void reAlloc() {
    for (final ValueVector v : vectors.values()) {
      v.reAlloc();
    }
  }

  /**
   * Adds a new field with the given parameters or replaces the existing one and consequently returns the resultant
   * {@link org.apache.arrow.vector.ValueVector}.
   *
   * <p>Execution takes place in the following order:
   * <ul>
   * <li>
   * if field is new, create and insert a new vector of desired type.
   * </li>
   * <li>
   * if field exists and existing vector is of desired vector type, return the vector.
   * </li>
   * <li>
   * if field exists and null filled, clear the existing vector; create and insert a new vector of desired type.
   * </li>
   * <li>
   * otherwise, throw an {@link java.lang.IllegalStateException}
   * </li>
   * </ul>
   *
   * @param childName the name of the field
   * @param fieldType the type for the vector
   * @param clazz     class of expected vector type
   * @param <T>       class type of expected vector type
   * @return resultant {@link org.apache.arrow.vector.ValueVector}
   * @throws java.lang.IllegalStateException raised if there is a hard schema change
   */
  public <T extends FieldVector> T addOrGet(String childName, FieldType fieldType, Class<T> clazz) {
    final ValueVector existing = getChild(childName);
    boolean create = false;
    if (existing == null) {
      create = true;
    } else if (clazz.isAssignableFrom(existing.getClass())) {
      return clazz.cast(existing);
    } else if (nullFilled(existing)) {
      existing.clear();
      create = true;
    }
    if (create) {
      final T vector = clazz.cast(fieldType.createNewSingleVector(childName, allocator, callBack));
      putChild(childName, vector);
      if (callBack != null) {
        callBack.doWork();
      }
      return vector;
    }
    final String message = "Arrow does not support schema change yet. Existing[%s] and desired[%s] vector types " +
        "mismatch";
    throw new IllegalStateException(String.format(message, existing.getClass().getSimpleName(), clazz.getSimpleName()));
  }

  private boolean nullFilled(ValueVector vector) {
    return BitVectorHelper.checkAllBitsEqualTo(vector.getValidityBuffer(), vector.getValueCount(), false);
  }

  /**
   * Returns a {@link org.apache.arrow.vector.ValueVector} corresponding to the given ordinal identifier.
   *
   * @param id the ordinal of the child to return
   * @return the corresponding child
   */
  public ValueVector getChildByOrdinal(int id) {
    return vectors.getByOrdinal(id);
  }

  /**
   * Returns a {@link org.apache.arrow.vector.ValueVector} instance of subtype of T corresponding to the given
   * field name if exists or null.
   *
   * If there is more than one element for name this will return the first inserted.
   *
   * @param name  the name of the child to return
   * @param clazz the expected type of the child
   * @return the child corresponding to this name
   */
  @Override
  public <T extends FieldVector> T getChild(String name, Class<T> clazz) {
    final FieldVector f = vectors.get(name);
    if (f == null) {
      return null;
    }
    return typeify(f, clazz);
  }

  protected ValueVector add(String childName, FieldType fieldType) {
    FieldVector vector = fieldType.createNewSingleVector(childName, allocator, callBack);
    putChild(childName, vector);
    if (callBack != null) {
      callBack.doWork();
    }
    return vector;
  }

  /**
   * Inserts the vector with the given name if it does not exist else replaces it with the new value.
   *
   * <p>Note that this method does not enforce any vector type check nor throws a schema change exception.
   *
   * @param name   the name of the child to add
   * @param vector the vector to add as a child
   */
  protected void putChild(String name, FieldVector vector) {
    putVector(name, vector);
  }

  private void put(String name, FieldVector vector, boolean overwrite) {
    final boolean old = vectors.put(
        Preconditions.checkNotNull(name, "field name cannot be null"),
        Preconditions.checkNotNull(vector, "vector cannot be null"),
        overwrite
    );
    if (old) {
      logger.debug("Field [{}] mutated to [{}] ", name,
          vector.getClass().getSimpleName());
    }
  }

  /**
   * Inserts the input vector into the map if it does not exist.
   *
   * <p>
   * If the field name already exists the conflict is handled according to the currently set ConflictPolicy
   * </p>
   *
   * @param name   field name
   * @param vector vector to be inserted
   */
  protected void putVector(String name, FieldVector vector) {
    switch (conflictPolicy) {
      case CONFLICT_APPEND:
        put(name, vector, false);
        break;
      case CONFLICT_IGNORE:
        if (!vectors.containsKey(name)) {
          put(name, vector, false);
        }
        break;
      case CONFLICT_REPLACE:
        if (vectors.containsKey(name)) {
          vectors.removeAll(name);
        }
        put(name, vector, true);
        break;
      case CONFLICT_ERROR:
        if (vectors.containsKey(name)) {
          throw new IllegalStateException(String.format("Vector already exists: Existing[%s], Requested[%s] ",
            vector.getClass().getSimpleName(), vector.getField().getFieldType()));
        }
        put(name, vector, false);
        break;
      default:
        throw new IllegalStateException(String.format("%s type not a valid conflict state", conflictPolicy));
    }

  }

  /**
   * Get child vectors.
   * @return a sequence of underlying child vectors.
   */
  protected List<FieldVector> getChildren() {
    int size = vectors.size();
    List<FieldVector> children = new ArrayList<>();
    for (int i = 0; i < size; i++) {
      children.add(vectors.getByOrdinal(i));
    }
    return children;
  }

  /**
   * Get child field names.
   */
  public List<String> getChildFieldNames() {
    return getChildren().stream()
        .map(child -> child.getField().getName())
        .collect(Collectors.toList());
  }

  /**
   * Get the number of child vectors.
   * @return the number of underlying child vectors.
   */
  @Override
  public int size() {
    return vectors.size();
  }

  @Override
  public Iterator<ValueVector> iterator() {
    return Collections.<ValueVector>unmodifiableCollection(vectors.values()).iterator();
  }

  /**
   * Get primitive child vectors.
   * @return a list of scalar child vectors recursing the entire vector hierarchy.
   */
  public List<ValueVector> getPrimitiveVectors() {
    final List<ValueVector> primitiveVectors = new ArrayList<>();
    for (final FieldVector v : vectors.values()) {
      primitiveVectors.addAll(getPrimitiveVectors(v));
    }
    return primitiveVectors;
  }

  private List<ValueVector> getPrimitiveVectors(FieldVector v) {
    final List<ValueVector> primitives = new ArrayList<>();
    if (v instanceof AbstractStructVector) {
      AbstractStructVector structVector = (AbstractStructVector) v;
      primitives.addAll(structVector.getPrimitiveVectors());
    } else if (v instanceof ListVector) {
      ListVector listVector = (ListVector) v;
      primitives.addAll(getPrimitiveVectors(listVector.getDataVector()));
    } else if (v instanceof FixedSizeListVector) {
      FixedSizeListVector listVector = (FixedSizeListVector) v;
      primitives.addAll(getPrimitiveVectors(listVector.getDataVector()));
    } else if (v instanceof UnionVector) {
      UnionVector unionVector = (UnionVector) v;
      for (final FieldVector vector : unionVector.getChildrenFromFields()) {
        primitives.addAll(getPrimitiveVectors(vector));
      }
    } else {
      primitives.add(v);
    }
    return primitives;
  }

  /**
   * Get a child vector by name. If duplicate names this returns the first inserted.
   * @param name the name of the child to return
   * @return a vector with its corresponding ordinal mapping if field exists or null.
   */
  @Override
  public VectorWithOrdinal getChildVectorWithOrdinal(String name) {
    final int ordinal = vectors.getOrdinal(name);
    if (ordinal < 0) {
      return null;
    }
    final ValueVector vector = vectors.getByOrdinal(ordinal);
    return new VectorWithOrdinal(vector, ordinal);
  }

  @Override
  public ArrowBuf[] getBuffers(boolean clear) {
    final List<ArrowBuf> buffers = new ArrayList<>();

    for (final ValueVector vector : vectors.values()) {
      for (final ArrowBuf buf : vector.getBuffers(false)) {
        buffers.add(buf);
        if (clear) {
          buf.getReferenceManager().retain(1);
        }
      }
      if (clear) {
        vector.clear();
      }
    }

    return buffers.toArray(new ArrowBuf[buffers.size()]);
  }

  @Override
  public int getBufferSize() {
    int actualBufSize = 0;

    for (final ValueVector v : vectors.values()) {
      for (final ArrowBuf buf : v.getBuffers(false)) {
        actualBufSize += buf.writerIndex();
      }
    }
    return actualBufSize;
  }

  @Override
  public String toString() {
    return ValueVectorUtility.getToString(this, 0 , getValueCount());
  }

}
