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

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.VectorLoader;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.compression.CompressionCodec;
import org.apache.arrow.vector.dictionary.Dictionary;
import org.apache.arrow.vector.dictionary.DictionaryProvider;
import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch;
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.DictionaryUtility;
import org.apache.arrow.vector.util.VectorBatchAppender;

/** Abstract class to read Schema and ArrowRecordBatches. */
public abstract class ArrowReader implements DictionaryProvider, AutoCloseable {

  protected final BufferAllocator allocator;
  private VectorLoader loader;
  private VectorSchemaRoot root;
  protected Map<Long, Dictionary> dictionaries;
  private boolean initialized = false;

  private final CompressionCodec.Factory compressionFactory;

  protected ArrowReader(BufferAllocator allocator) {
    this(allocator, CompressionCodec.Factory.INSTANCE);
  }

  protected ArrowReader(BufferAllocator allocator, CompressionCodec.Factory compressionFactory) {
    this.allocator = allocator;
    this.compressionFactory = compressionFactory;
  }

  /**
   * Returns the vector schema root. This will be loaded with new values on every call to
   * loadNextBatch.
   *
   * @return the vector schema root
   * @throws IOException if reading of schema fails
   */
  public VectorSchemaRoot getVectorSchemaRoot() throws IOException {
    ensureInitialized();
    return root;
  }

  /**
   * Returns any dictionaries that were loaded along with ArrowRecordBatches.
   *
   * @return Map of dictionaries to dictionary id, empty if no dictionaries loaded
   * @throws IOException if reading of schema fails
   */
  public Map<Long, Dictionary> getDictionaryVectors() throws IOException {
    ensureInitialized();
    return dictionaries;
  }

  /**
   * Lookup a dictionary that has been loaded using the dictionary id.
   *
   * @param id Unique identifier for a dictionary
   * @return the requested dictionary or null if not found
   */
  @Override
  public Dictionary lookup(long id) {
    if (!initialized) {
      throw new IllegalStateException("Unable to lookup until reader has been initialized");
    }

    return dictionaries.get(id);
  }

  @Override
  public Set<Long> getDictionaryIds() {
    return dictionaries.keySet();
  }

  /**
   * Load the next ArrowRecordBatch to the vector schema root if available.
   *
   * @return true if a batch was read, false on EOS
   * @throws IOException on error
   */
  public abstract boolean loadNextBatch() throws IOException;

  /**
   * Return the number of bytes read from the ReadChannel.
   *
   * @return number of bytes read
   */
  public abstract long bytesRead();

  /**
   * Close resources, including vector schema root and dictionary vectors, and the underlying read
   * source.
   *
   * @throws IOException on error
   */
  @Override
  public void close() throws IOException {
    close(true);
  }

  /**
   * Close resources, including vector schema root and dictionary vectors. If the flag
   * closeReadChannel is true then close the underlying read source, otherwise leave it open.
   *
   * @param closeReadSource Flag to control if closing the underlying read source
   * @throws IOException on error
   */
  public void close(boolean closeReadSource) throws IOException {
    if (initialized) {
      root.close();
      for (Dictionary dictionary : dictionaries.values()) {
        dictionary.getVector().close();
      }
    }

    if (closeReadSource) {
      closeReadSource();
    }
  }

  /**
   * Close the underlying read source.
   *
   * @throws IOException on error
   */
  protected abstract void closeReadSource() throws IOException;

  /**
   * Read the Schema from the source, will be invoked at the beginning the initialization.
   *
   * @return the read Schema
   * @throws IOException on error
   */
  protected abstract Schema readSchema() throws IOException;

  /**
   * Initialize if not done previously.
   *
   * @throws IOException on error
   */
  protected void ensureInitialized() throws IOException {
    if (!initialized) {
      initialize();
      initialized = true;
    }
  }

  /** Reads the schema and initializes the vectors. */
  protected void initialize() throws IOException {
    Schema originalSchema = readSchema();
    List<Field> fields = new ArrayList<>(originalSchema.getFields().size());
    List<FieldVector> vectors = new ArrayList<>(originalSchema.getFields().size());
    Map<Long, Dictionary> dictionaries = new HashMap<>();

    // Convert fields with dictionaries to have the index type
    for (Field field : originalSchema.getFields()) {
      Field updated = DictionaryUtility.toMemoryFormat(field, allocator, dictionaries);
      fields.add(updated);
      vectors.add(updated.createVector(allocator));
    }
    Schema schema = new Schema(fields, originalSchema.getCustomMetadata());

    this.root = new VectorSchemaRoot(schema, vectors, 0);
    this.loader = new VectorLoader(root, compressionFactory);
    this.dictionaries = Collections.unmodifiableMap(dictionaries);
  }

  /**
   * Ensure the reader has been initialized and reset the VectorSchemaRoot row count to 0.
   *
   * @throws IOException on error
   */
  protected void prepareLoadNextBatch() throws IOException {
    ensureInitialized();
    root.setRowCount(0);
  }

  /**
   * Load an ArrowRecordBatch to the readers VectorSchemaRoot.
   *
   * @param batch the record batch to load
   */
  protected void loadRecordBatch(ArrowRecordBatch batch) {
    try {
      loader.load(batch);
    } finally {
      batch.close();
    }
  }

  /**
   * Load an ArrowDictionaryBatch to the readers dictionary vectors.
   *
   * @param dictionaryBatch dictionary batch to load
   */
  protected void loadDictionary(ArrowDictionaryBatch dictionaryBatch) {
    long id = dictionaryBatch.getDictionaryId();
    Dictionary dictionary = dictionaries.get(id);
    if (dictionary == null) {
      throw new IllegalArgumentException("Dictionary ID " + id + " not defined in schema");
    }
    FieldVector vector = dictionary.getVector();
    // if is deltaVector, concat it with non-delta vector with the same ID.
    if (dictionaryBatch.isDelta()) {
      try (FieldVector deltaVector = vector.getField().createVector(allocator)) {
        load(dictionaryBatch, deltaVector);
        VectorBatchAppender.batchAppend(vector, deltaVector);
      }
      return;
    }

    load(dictionaryBatch, vector);
  }

  private void load(ArrowDictionaryBatch dictionaryBatch, FieldVector vector) {
    VectorSchemaRoot root =
        new VectorSchemaRoot(
            Collections.singletonList(vector.getField()), Collections.singletonList(vector), 0);
    VectorLoader loader = new VectorLoader(root, this.compressionFactory);
    try {
      loader.load(dictionaryBatch.getDictionary());
    } finally {
      dictionaryBatch.close();
    }
  }
}
