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

import java.io.IOException;

import org.apache.parquet.column.ColumnDescriptor;
import org.apache.parquet.column.page.PageReader;
import org.apache.spark.sql.types.DataType;

import org.apache.comet.CometSchemaImporter;
import org.apache.comet.vector.CometLazyVector;
import org.apache.comet.vector.CometVector;

public class LazyColumnReader extends ColumnReader {

  // Remember the largest skipped index for sanity checking.
  private int lastSkippedRowId = Integer.MAX_VALUE;

  // Track whether the underlying page is drained.
  private boolean isPageDrained = true;

  // Leftover number of rows that did not skip in the previous batch.
  private int numRowsToSkipFromPrevBatch;

  // The lazy vector being updated.
  private final CometLazyVector vector;

  LazyColumnReader(
      DataType sparkReadType,
      ColumnDescriptor descriptor,
      CometSchemaImporter importer,
      int batchSize,
      boolean useDecimal128,
      boolean useLegacyDateTimestamp) {
    super(sparkReadType, descriptor, importer, batchSize, useDecimal128, useLegacyDateTimestamp);
    this.batchSize = 0; // the batch size is set later in `readBatch`
    this.vector = new CometLazyVector(sparkReadType, this, useDecimal128);
  }

  @Override
  public void setPageReader(PageReader pageReader) throws IOException {
    super.setPageReader(pageReader);
    lastSkippedRowId = Integer.MAX_VALUE;
    isPageDrained = true;
    numRowsToSkipFromPrevBatch = 0;
    currentNumValues = batchSize;
  }

  /**
   * Lazily read a batch of 'total' rows for this column. The includes: 1) Skip any unused rows from
   * the previous batch 2) Reset the native columnar batch 3) Reset tracking variables
   *
   * @param total the number of rows in the batch. MUST be <= the number of rows available in this
   *     column chunk.
   */
  @Override
  public void readBatch(int total) {
    // Before starting a new batch, take care of the remaining rows to skip from the previous batch.
    tryPageSkip(batchSize);
    numRowsToSkipFromPrevBatch += batchSize - currentNumValues;

    // Now first reset the current columnar batch so that it can be used to fill in a new batch
    // of values. Then, keep reading more data pages (via 'readBatch') until the current batch is
    // full, or we have read 'total' number of values.
    Native.resetBatch(nativeHandle);

    batchSize = total;
    currentNumValues = 0;
    lastSkippedRowId = -1;
  }

  @Override
  public CometVector currentBatch() {
    return vector;
  }

  /** Read all rows up to the `batchSize`. Expects no rows are skipped so far. */
  public void readAllBatch() {
    // All rows should be read without any skips so far
    assert (lastSkippedRowId == -1);

    readBatch(batchSize - 1, 0);
  }

  /**
   * Read at least up to `rowId`. It may read beyond `rowId` if enough rows available in the page.
   * It may skip reading rows before `rowId`. In case `rowId` is already read, return immediately.
   *
   * @param rowId the row index in the batch to read.
   * @return true if `rowId` is newly materialized, or false if `rowId` is already materialized.
   */
  public boolean materializeUpToIfNecessary(int rowId) {
    // Not allowed reading rowId if it may have skipped previously.
    assert (rowId > lastSkippedRowId);

    // If `rowId` is already materialized, return immediately.
    if (rowId < currentNumValues) return false;

    int numRowsWholePageSkipped = tryPageSkip(rowId);
    readBatch(rowId, numRowsWholePageSkipped);
    return true;
  }

  /**
   * Read up to `rowId` (inclusive). If the whole pages are skipped previously in `tryPageSkip()`,
   * pad the number of whole page skipped rows with nulls to the underlying vector before reading.
   *
   * @param rowId the row index in the batch to read.
   * @param numNullRowsToPad the number of nulls to pad before reading.
   */
  private void readBatch(int rowId, int numNullRowsToPad) {
    if (numRowsToSkipFromPrevBatch > 0) {
      // Reaches here only when starting a new batch and the page is previously drained
      readPage();
      isPageDrained = false;
      Native.skipBatch(nativeHandle, numRowsToSkipFromPrevBatch, true);
      numRowsToSkipFromPrevBatch = 0;
    }
    while (rowId >= currentNumValues) {
      int numRowsToRead = batchSize - currentNumValues;
      if (isPageDrained) {
        readPage();
      }
      int[] array = Native.readBatch(nativeHandle, numRowsToRead, numNullRowsToPad);
      int read = array[0];
      isPageDrained = read < numRowsToRead;
      currentNumValues += read;
      currentNumNulls += array[1];
      // No need to update numNullRowsToPad. numNullRowsToPad > 0 means there were whole page skips.
      // That guarantees that the Native.readBatch can read up to rowId in the current page.
    }
  }

  /**
   * Try to skip until `rowId` (exclusive). If possible, it skips whole underlying pages without
   * decompressing. In that case, it returns early at the page end, so that the next iteration can
   * lazily decide to `readPage()` or `tryPageSkip()` again.
   *
   * @param rowId the row index in the batch that it tries to skip up until (exclusive).
   * @return the number of rows that the whole page skips were applied.
   */
  private int tryPageSkip(int rowId) {
    int total = rowId - currentNumValues;
    int wholePageSkipped = 0;
    if (total > 0) {
      // First try to skip from the non-drained underlying page.
      int skipped = isPageDrained ? 0 : Native.skipBatch(nativeHandle, total);
      total -= skipped;
      isPageDrained = total > 0;
      if (isPageDrained) {
        ColumnPageReader columnPageReader = (ColumnPageReader) pageReader;
        // It is always `columnPageReader.getPageValueCount() > numRowsToSkipFromPriorBatch`
        int pageValueCount = columnPageReader.getPageValueCount() - numRowsToSkipFromPrevBatch;
        while (pageValueCount <= total) {
          // skip the entire page if the next page is small enough
          columnPageReader.skipPage();
          numRowsToSkipFromPrevBatch = 0;
          total -= pageValueCount;
          wholePageSkipped += pageValueCount;
          pageValueCount = columnPageReader.getPageValueCount();
        }
      }

      currentNumValues += skipped + wholePageSkipped;
      currentNumNulls += skipped;
      lastSkippedRowId = currentNumValues - 1;
    }
    return wholePageSkipped;
  }
}
