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

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.complex.reader.FieldReader;
import org.apache.arrow.vector.dictionary.Dictionary;
import org.apache.arrow.vector.dictionary.DictionaryEncoder;
import org.apache.arrow.vector.dictionary.DictionaryProvider;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.TransferPair;

/**
 * Abstract base class for Table.
 *
 * <p>This API is EXPERIMENTAL.
 */
public abstract class BaseTable implements AutoCloseable {

  /** The field vectors holding the data in this table. */
  protected final List<FieldVector> fieldVectors;

  /**
   * An optional DictionaryProvider. One must be present if any vector in the table is dictionary
   * encoded.
   */
  protected DictionaryProvider dictionaryProvider;

  /** A map of Fields to FieldVectors used to select Fields. */
  protected final Map<Field, FieldVector> fieldVectorsMap = new LinkedHashMap<>();

  /** The schema for the table. */
  protected Schema schema;

  /**
   * The number of rows of data in the table; not necessarily the same as the table row capacity.
   */
  protected int rowCount;

  /**
   * Constructs new instance with the given rowCount, and containing the schema and each of the
   * given vectors.
   *
   * @param fieldVectors the FieldVectors containing the table's data
   * @param rowCount the number of rows in the table
   * @param provider a dictionary provider, may be null if none of the vectors in the table are
   *     encoded
   */
  public BaseTable(List<FieldVector> fieldVectors, int rowCount, DictionaryProvider provider) {

    this.dictionaryProvider = provider;
    this.rowCount = rowCount;
    this.fieldVectors = new ArrayList<>();
    List<Field> fields = new ArrayList<>();
    for (FieldVector fv : fieldVectors) {
      TransferPair transferPair = fv.getTransferPair(fv.getAllocator());
      transferPair.transfer();
      FieldVector newVector = (FieldVector) transferPair.getTo();
      newVector.setValueCount(rowCount);

      Field newField = newVector.getField();
      this.fieldVectors.add(newVector);
      fields.add(newField);
      fieldVectorsMap.put(newField, newVector);
    }
    this.schema = new Schema(fields);
  }

  BaseTable() {
    this.fieldVectors = new ArrayList<>();
  }

  /**
   * Returns a FieldReader for the vector with the given name.
   *
   * @param name The name of a vector in this Table (case-sensitive)
   * @return A FieldReader for the named FieldVector
   */
  public FieldReader getReader(String name) {
    for (Map.Entry<Field, FieldVector> entry : fieldVectorsMap.entrySet()) {
      if (entry.getKey().getName().equals(name)) {
        return entry.getValue().getReader();
      }
    }
    return null;
  }

  /**
   * Returns a FieldReader for the given field.
   *
   * @param field The field to be read
   * @return A FieldReader for the given field
   */
  public FieldReader getReader(Field field) {
    return fieldVectorsMap.get(field).getReader();
  }

  /**
   * Returns a FieldReader for the field at the given vector index.
   *
   * @param index The 0-based index of the field desired.
   * @return A FieldReader for the requested field
   */
  public FieldReader getReader(int index) {
    Preconditions.checkArgument(index >= 0 && index < fieldVectors.size());
    return fieldVectors.get(index).getReader();
  }

  /** Returns the schema for this Table. */
  public Schema getSchema() {
    return schema;
  }

  /**
   * Returns the Field with the given name if one exists in this table.
   *
   * @param fieldName the name of the field to return
   * @return a field with the given name if one is present
   * @throws IllegalArgumentException – if the field was not found
   */
  public Field getField(String fieldName) {
    return getSchema().findField(fieldName);
  }

  /**
   * Returns a list of Field created by adding the given vector to the vectors in this Table.
   *
   * @param index field index
   * @param vector vector to be added.
   * @return out List of FieldVectors with vector added
   */
  List<FieldVector> insertVector(int index, FieldVector vector) {
    Preconditions.checkNotNull(vector);
    Preconditions.checkArgument(index >= 0 && index <= fieldVectors.size());
    List<FieldVector> newVectors = new ArrayList<>();
    if (index == fieldVectors.size()) {
      newVectors.addAll(fieldVectors);
      newVectors.add(vector);
    } else {
      for (int i = 0; i < fieldVectors.size(); i++) {
        if (i == index) {
          newVectors.add(vector);
        }
        newVectors.add(fieldVectors.get(i));
      }
    }
    return newVectors;
  }

  /**
   * Returns a new List of FieldVectors created by removing the selected Vector from the list in
   * this Table.
   *
   * @param index field index
   * @return out List of FieldVectors like the list in this table, but with the argument removed
   */
  List<FieldVector> extractVector(int index) {
    Preconditions.checkArgument(index >= 0 && index < fieldVectors.size());
    List<FieldVector> newVectors = new ArrayList<>();
    for (int i = 0; i < fieldVectors.size(); i++) {
      if (i != index) {
        newVectors.add(fieldVectors.get(i));
      }
    }
    return newVectors;
  }

  /** Returns the number of vectors (columns) in this table. */
  public int getVectorCount() {
    return fieldVectors.size();
  }

  /**
   * Closes all the vectors holding data for this table and sets the rowcount to 0, preventing
   * enumeration.
   */
  void clear() {
    close();
    rowCount = 0;
  }

  /** Closes all the vectors holding data for this table. */
  @Override
  public void close() {
    try {
      AutoCloseables.close(fieldVectors);
    } catch (RuntimeException ex) {
      throw ex;
    } catch (Exception ex) {
      // should never happen since FieldVector.close() doesn't throw IOException
      throw new RuntimeException(ex);
    }
  }

  /** Returns the number of rows in this table. */
  public long getRowCount() {
    return rowCount;
  }

  /**
   * Returns a new VectorSchemaRoot with the data and schema from this table. Data is transferred to
   * the new VectorSchemaRoot, so this table is cleared and the rowCount is set to 0;
   *
   * @return a new VectorSchemaRoot
   */
  public VectorSchemaRoot toVectorSchemaRoot() {
    VectorSchemaRoot vsr =
        new VectorSchemaRoot(
            fieldVectors.stream()
                .map(
                    v -> {
                      TransferPair transferPair = v.getTransferPair(v.getAllocator());
                      transferPair.transfer();
                      return (FieldVector) transferPair.getTo();
                    })
                .collect(Collectors.toList()));
    clear();
    return vsr;
  }

  /**
   * Returns the vector with the given name, or throws IllegalArgumentException if the name is not
   * found. Names are case-sensitive.
   *
   * @param columnName The name of the vector
   * @return the Vector with the given name, or null
   * @throws IllegalArgumentException if the name is not the name of a vector in the table.
   */
  FieldVector getVector(String columnName) {
    for (Map.Entry<Field, FieldVector> entry : fieldVectorsMap.entrySet()) {
      if (entry.getKey().getName().equals(columnName)) {
        return entry.getValue();
      }
    }
    throw new IllegalArgumentException(
        String.format("No vector named '%s' is present in the table", columnName));
  }

  /**
   * Returns the vector at the given position.
   *
   * @param columnIndex The 0-based position of the vector
   */
  FieldVector getVector(int columnIndex) {
    return fieldVectors.get(columnIndex);
  }

  /**
   * Returns a copy of the vector with the given name, or throws IllegalArgumentException if the
   * name is not found. Names are case-sensitive.
   *
   * @param columnName The name of the vector to copy
   * @return A copy of the Vector with the given name
   * @throws IllegalArgumentException if the name is not the name of a vector in the table.
   */
  public FieldVector getVectorCopy(String columnName) {
    FieldVector source;
    for (Map.Entry<Field, FieldVector> entry : fieldVectorsMap.entrySet()) {
      if (entry.getKey().getName().equals(columnName)) {
        source = entry.getValue();
        FieldVector copy = source.getField().createVector(source.getAllocator());
        copy.allocateNew();
        for (int i = 0; i < source.getValueCount(); i++) {
          copy.copyFromSafe(i, i, source);
        }
        copy.setValueCount(source.getValueCount());
        return copy;
      }
    }
    throw new IllegalStateException(
        String.format("No vector named '%s' is present in the table", columnName));
  }

  /**
   * Returns a copy of the vector at the given position.
   *
   * @param columnIndex The 0-based position of the vector to be copied
   */
  public FieldVector getVectorCopy(int columnIndex) {
    FieldVector source = fieldVectors.get(columnIndex);
    FieldVector copy = source.getField().createVector(source.getAllocator());
    copy.allocateNew();
    for (int i = 0; i < source.getValueCount(); i++) {
      copy.copyFromSafe(i, i, source);
    }
    copy.setValueCount(source.getValueCount());
    return copy;
  }

  /**
   * Returns an immutable Row object holding a reference to this table. The default character
   * encoding used by the cursor to decode Strings will be StandardCharsets.UTF_8 as this is the
   * only charset supported in Arrow format.
   */
  public Row immutableRow() {
    return new Row(this);
  }

  /** Returns a tab separated value of vectors (based on their java object representation). */
  public String contentToTSVString() {
    StringBuilder sb = new StringBuilder();
    List<Object> row = new ArrayList<>(schema.getFields().size());
    for (Field field : schema.getFields()) {
      row.add(field.getName());
    }
    printRow(sb, row);
    for (int i = 0; i < rowCount; i++) {
      row.clear();
      for (FieldVector v : fieldVectors) {
        row.add(v.getObject(i));
      }
      printRow(sb, row);
    }
    return sb.toString();
  }

  /**
   * Prints a single row without a header to the given StringBuilder.
   *
   * @param sb the StringBuilder to write to
   * @param row the row to write
   */
  private void printRow(StringBuilder sb, List<Object> row) {
    boolean first = true;
    for (Object v : row) {
      if (first) {
        first = false;
      } else {
        sb.append("\t");
      }
      sb.append(v);
    }
    sb.append("\n");
  }

  /**
   * Returns true if the row at the given index has been deleted and false otherwise.
   *
   * <p>If the index is larger than the number of rows, the method returns true.
   *
   * @param rowNumber The 0-based index of the possibly deleted row
   * @return true if the row at the index was deleted; false otherwise
   */
  public boolean isRowDeleted(int rowNumber) {
    return false;
  }

  /** Returns the DictionaryProvider for this table. It can be used to decode an encoded values */
  public DictionaryProvider getDictionaryProvider() {
    return dictionaryProvider;
  }

  /**
   * Returns a ValueVector containing the decoded version of the vector with the given name.
   *
   * @param vectorName The name of the vector to decode
   * @param dictionaryId The identifier for the dictionary to use when decoding. Must match the id
   *     returned by the dictionary's getId() method.
   * @return A ValueVector
   */
  public ValueVector decode(String vectorName, long dictionaryId) {
    Dictionary dictionary = getDictionary(dictionaryId);

    FieldVector vector = getVector(vectorName);
    if (vector == null) {
      throw new IllegalArgumentException(
          String.format("No vector with name '%s' is present in table", vectorName));
    }

    DictionaryEncoder decoder = new DictionaryEncoder(dictionary, vector.getAllocator());
    return decoder.decode(vector);
  }

  /**
   * Returns a ValueVector containing the encoded version of the vector with the given name.
   *
   * @param vectorName The name of the vector to encode
   * @param dictionaryId The identifier for the dictionary to use when encoding. Must match the id
   *     returned by the dictionary's getId() method.
   * @return A ValueVector
   */
  public ValueVector encode(String vectorName, long dictionaryId) {
    Dictionary dictionary = getDictionary(dictionaryId);
    FieldVector vector = getVector(vectorName);
    if (vector == null) {
      throw new IllegalArgumentException(
          String.format("No vector with name '%s' is present in table", vectorName));
    }
    DictionaryEncoder decoder = new DictionaryEncoder(dictionary, vector.getAllocator());
    return decoder.encode(vector);
  }

  /**
   * Returns the dictionary with given id.
   *
   * @param dictionaryId A long integer that is the id returned by the dictionary's getId() method
   */
  private Dictionary getDictionary(long dictionaryId) {
    if (dictionaryProvider == null) {
      throw new IllegalStateException("No dictionary provider is present in table.");
    }

    Dictionary dictionary = dictionaryProvider.lookup(dictionaryId);
    if (dictionary == null) {
      throw new IllegalArgumentException("No dictionary with id '%n' exists in the table");
    }
    return dictionary;
  }
}
