/*
 * 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.lucene.codecs.uniformsplit;

import java.io.IOException;

import org.apache.lucene.codecs.BlockTermState;
import org.apache.lucene.codecs.PostingsReaderBase;
import org.apache.lucene.index.BaseTermsEnum;
import org.apache.lucene.index.CorruptIndexException;
import org.apache.lucene.index.ImpactsEnum;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.TermState;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.store.ByteArrayDataInput;
import org.apache.lucene.store.IndexInput;
import org.apache.lucene.util.Accountable;
import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.BytesRefBuilder;
import org.apache.lucene.util.RamUsageEstimator;

/**
 * Seeks the block corresponding to a given term, read the block bytes, and
 * scans the block terms.
 * <p>
 * Reads fully the block in {@link #blockReadBuffer}. Then scans the block
 * terms in memory. The details region is lazily decoded with {@link #termStatesReadBuffer}
 * which shares the same byte array with {@link #blockReadBuffer}.
 * See {@link BlockWriter} and {@link BlockLine} for the block format.
 *
 * @lucene.experimental
 */
public class BlockReader extends BaseTermsEnum implements Accountable {

  private static final long BASE_RAM_USAGE = RamUsageEstimator.shallowSizeOfInstance(BlockReader.class)
      + RamUsageEstimator.shallowSizeOfInstance(IndexInput .class)
      +RamUsageEstimator.shallowSizeOfInstance(ByteArrayDataInput .class)*2;

  /**
   * {@link IndexInput} on the {@link UniformSplitPostingsFormat#TERMS_BLOCKS_EXTENSION block file}.
   */
  protected IndexInput blockInput;
  protected final PostingsReaderBase postingsReader;
  protected final FieldMetadata fieldMetadata;
  protected final BlockDecoder blockDecoder;

  protected BlockHeader.Serializer blockHeaderReader;
  protected BlockLine.Serializer blockLineReader;
  /**
   * In-memory read buffer for the current block.
   */
  protected ByteArrayDataInput blockReadBuffer;
  /**
   * In-memory read buffer for the details region of the current block.
   * It shares the same byte array as {@link #blockReadBuffer}, with a
   * different position.
   */
  protected ByteArrayDataInput termStatesReadBuffer;
  protected DeltaBaseTermStateSerializer termStateSerializer;

  /**
   * {@link IndexDictionary.Browser} supplier for lazy loading.
   */
  protected final IndexDictionary.BrowserSupplier dictionaryBrowserSupplier;
  /**
   * Holds the {@link IndexDictionary.Browser} once loaded.
   */
  protected IndexDictionary.Browser dictionaryBrowser;

  /**
   * Current block start file pointer, absolute in the
   * {@link UniformSplitPostingsFormat#TERMS_BLOCKS_EXTENSION block file}.
   */
  protected long blockStartFP;
  /**
   * Current block header.
   */
  protected BlockHeader blockHeader;
  /**
   * Current block line.
   */
  protected BlockLine blockLine;
  /**
   * Current block line details.
   */
  protected BlockTermState termState;
  /**
   * Offset of the start of the first line of the current block (just after the header), relative to the block start.
   */
  protected int blockFirstLineStart;
  /**
   * Current line index in the block.
   */
  protected int lineIndexInBlock;
  /**
   * Whether the current {@link TermState} has been forced with a call to
   * {@link #seekExact(BytesRef, TermState)}.
   *
   * @see #forcedTerm
   */
  protected boolean termStateForced;
  /**
   * Set when {@link #seekExact(BytesRef, TermState)} is called.
   * <p>
   * This optimizes the use-case when the caller calls first {@link #seekExact(BytesRef, TermState)}
   * and then {@link #postings(PostingsEnum, int)}. In this case we don't access
   * the terms block file (we don't seek) but directly the postings file because
   * we already have the {@link TermState} with the file pointers to the postings
   * file.
   */
  protected BytesRefBuilder forcedTerm;

  // Scratch objects to avoid object reallocation.
  protected BytesRef scratchBlockBytes;
  protected final BlockTermState scratchTermState;
  protected BlockLine scratchBlockLine;

  /**
   * @param dictionaryBrowserSupplier to load the {@link IndexDictionary.Browser}
   *                                  lazily in {@link #seekCeil(BytesRef)}.
   * @param blockDecoder              Optional block decoder, may be null if none.
   *                                  It can be used for decompression or decryption.
   */
  protected BlockReader(IndexDictionary.BrowserSupplier dictionaryBrowserSupplier, IndexInput blockInput,
                        PostingsReaderBase postingsReader, FieldMetadata fieldMetadata,
                        BlockDecoder blockDecoder) throws IOException {
    this.dictionaryBrowserSupplier = dictionaryBrowserSupplier;
    this.blockInput = blockInput;
    this.postingsReader = postingsReader;
    this.fieldMetadata = fieldMetadata;
    this.blockDecoder = blockDecoder;
    this.blockStartFP = -1;
    scratchTermState = postingsReader.newTermState();
  }

  @Override
  public SeekStatus seekCeil(BytesRef searchedTerm) throws IOException {
    if (isCurrentTerm(searchedTerm)) {
      return SeekStatus.FOUND;
    }
    clearTermState();

    long blockStartFP = getOrCreateDictionaryBrowser().seekBlock(searchedTerm);
    blockStartFP = Math.max(blockStartFP, fieldMetadata.getFirstBlockStartFP());
    if (isBeyondLastTerm(searchedTerm, blockStartFP)) {
      return SeekStatus.END;
    }
    SeekStatus seekStatus = seekInBlock(searchedTerm, blockStartFP);
    if (seekStatus != SeekStatus.END) {
      return seekStatus;
    }
    // Go to next block.
    return nextTerm() == null ? SeekStatus.END : SeekStatus.NOT_FOUND;
  }

  @Override
  public boolean seekExact(BytesRef searchedTerm) throws IOException {
    if (isCurrentTerm(searchedTerm)) {
      return true;
    }
    clearTermState();

    long blockStartFP = getOrCreateDictionaryBrowser().seekBlock(searchedTerm);
    if (blockStartFP < fieldMetadata.getFirstBlockStartFP() || isBeyondLastTerm(searchedTerm, blockStartFP)) {
      return false;
    }
    return seekInBlock(searchedTerm, blockStartFP) == SeekStatus.FOUND;
  }

  protected boolean isCurrentTerm(BytesRef searchedTerm) {
    // Optimization and also required to not search with the same BytesRef
    // instance as the BytesRef used to read the block line (BlockLine.Serializer).
    // Indeed getCurrentTerm() is allowed to return the same BytesRef instance.
    return searchedTerm.equals(term());
  }

  /**
   * Indicates whether the searched term is beyond the last term of the field.
   *
   * @param blockStartFP The current block start file pointer.
   */
  protected boolean isBeyondLastTerm(BytesRef searchedTerm, long blockStartFP) {
    return blockStartFP == fieldMetadata.getLastBlockStartFP()
        && searchedTerm.compareTo(fieldMetadata.getLastTerm()) > 0;
  }

  /**
   * Seeks to the provided term in the block starting at the provided file pointer.
   * Does not exceed the block.
   */
  protected SeekStatus seekInBlock(BytesRef searchedTerm, long blockStartFP) throws IOException {
    initializeHeader(searchedTerm, blockStartFP);
    if (blockHeader == null) {
      throw newCorruptIndexException("Illegal absence of block", blockStartFP);
    }
    return seekInBlock(searchedTerm);
  }

  /**
   * Seeks to the provided term in this block.
   * <p>
   * Does not exceed this block; {@link org.apache.lucene.index.TermsEnum.SeekStatus#END}
   * is returned if it follows the block.
   * <p>
   * Compares the line terms with the <code>searchedTerm</code>, taking
   * advantage of the incremental encoding properties.
   * <p>
   * Scans linearly the terms. Updates the current block line with the current
   * term.
   */
  protected SeekStatus seekInBlock(BytesRef searchedTerm) throws IOException {
    if (compareToMiddleAndJump(searchedTerm) == 0) {
      return SeekStatus.FOUND;
    }
    int comparisonOffset = 0;
    while (true) {
      if (readLineInBlock() == null) {
        // No more terms for the block.
        return SeekStatus.END;
      }
      TermBytes lineTermBytes = blockLine.getTermBytes();
      BytesRef lineTerm = lineTermBytes.getTerm();
      assert lineTerm.offset == 0;

      // Equivalent to comparing with BytesRef.compareTo(),
      // but faster since we start comparing from min(comparisonOffset, suffixOffset).
      int suffixOffset = lineTermBytes.getSuffixOffset();
      int start = Math.min(comparisonOffset, suffixOffset);
      int end = Math.min(searchedTerm.length, lineTerm.length);
      int comparison = searchedTerm.length - lineTerm.length;
      for (int i = start; i < end; i++) {
        // Compare unsigned bytes.
        int byteDiff = (searchedTerm.bytes[i + searchedTerm.offset] & 0xFF) - (lineTerm.bytes[i] & 0xFF);
        if (byteDiff != 0) {
          comparison = byteDiff;
          break;
        }
        comparisonOffset = i + 1;
      }
      if (comparison == 0) {
        return SeekStatus.FOUND;
      } else if (comparison < 0) {
        return SeekStatus.NOT_FOUND;
      }
    }
  }

  /**
   * Compares the searched term to the middle term of the block.
   * If the searched term is lexicographically equal or after the middle term
   * then jumps to the second half of the block directly.
   *
   * @return The comparison between the searched term and the middle term.
   */
  protected int compareToMiddleAndJump(BytesRef searchedTerm) throws IOException {
    if (lineIndexInBlock != 0) {
      // Don't try to compare and jump if we are not positioned at the first line.
      // This can happen if we seek in the same current block and we continue
      // scanning from the current line (see initializeHeader()).
      return -1;
    }
    blockReadBuffer.skipBytes(blockHeader.getMiddleLineOffset());
    lineIndexInBlock = blockHeader.getMiddleLineIndex();
    readLineInBlock();
    if (blockLine == null) {
      throw newCorruptIndexException("Illegal absence of line at the middle of the block", null);
    }
    int compare = searchedTerm.compareTo(term());
    if (compare < 0) {
      blockReadBuffer.setPosition(blockFirstLineStart);
      lineIndexInBlock = 0;
    }
    return compare;
  }

  /**
   * Reads the current block line.
   * Sets {@link #blockLine} and increments {@link #lineIndexInBlock}.
   *
   * @return The {@link BlockLine}; or null if there no more line in the block.
   */
  protected BlockLine readLineInBlock() throws IOException {
    if (lineIndexInBlock >= blockHeader.getLinesCount()) {
      return blockLine = null;
    }
    boolean isIncrementalEncodingSeed = lineIndexInBlock == 0 || lineIndexInBlock == blockHeader.getMiddleLineIndex();
    lineIndexInBlock++;
    return blockLine = blockLineReader.readLine(blockReadBuffer, isIncrementalEncodingSeed, scratchBlockLine);
  }

  /**
   * Positions this {@link BlockReader} without re-seeking the term dictionary.
   * <p>
   * The block containing the term is not read by this method. It will be read
   * lazily only if needed, for example if {@link #next()} is called.
   * Calling {@link #postings} after this method does require the block to be read.
   */
  @Override
  public void seekExact(BytesRef term, TermState state) {
    termStateForced = true;
    termState = scratchTermState;
    termState.copyFrom(state);
    if (forcedTerm == null) {
      forcedTerm = new BytesRefBuilder();
    }
    forcedTerm.copyBytes(term);
  }

  /**
   * Not supported.
   */
  @Override
  public void seekExact(long ord) {
    throw new UnsupportedOperationException();
  }

  @Override
  public BytesRef next() throws IOException {
    if (termStateForced) {
      initializeHeader(forcedTerm.get(), termState.blockFilePointer);
      if (blockHeader == null) {
        throw newCorruptIndexException("Illegal absence of block for TermState", termState.blockFilePointer);
      }
      for (int i = lineIndexInBlock; i < termState.termBlockOrd; i++) {
        readLineInBlock();
      }
      assert blockLine.getTermBytes().getTerm().equals(forcedTerm.get());
    }
    clearTermState();
    return nextTerm();
  }

  /**
   * Moves to the next term line and reads it, it may be in the next block.
   * The term details are not read yet. They will be read only when needed
   * with {@link #readTermStateIfNotRead()}.
   *
   * @return The read term bytes; or null if there is no more term for the field.
   */
  protected BytesRef nextTerm() throws IOException {
    if (blockHeader == null) {
      // Read the first block for the field.
      initializeHeader(null, fieldMetadata.getFirstBlockStartFP());
      if (blockHeader == null) {
        throw newCorruptIndexException("Illegal absence of first block", fieldMetadata.getFirstBlockStartFP());
      }
    }
    if (readLineInBlock() == null) {
      // No more line in the current block.
      // Read the next block starting at the current file pointer in the block file.
      initializeHeader(null, blockInput.getFilePointer());
      if (blockHeader == null) {
        // No more block for the field.
        return null;
      }
      readLineInBlock();
    }
    return term();
  }

  /**
   * Reads and sets {@link #blockHeader}. Sets null if there is no block for the field anymore.
   *
   * @param searchedTerm       The searched term; or null if none.
   * @param targetBlockStartFP The file pointer of the block to read.
   */
  protected void initializeHeader(BytesRef searchedTerm, long targetBlockStartFP) throws IOException {
    initializeBlockReadLazily();
    if (blockStartFP == targetBlockStartFP) {
      // Optimization: If the block to read is already the current block, then
      // reuse it directly without reading nor decoding the block bytes.
      if (blockHeader == null) {
        throw newCorruptIndexException("Illegal absence of block", blockStartFP);
      }
      if (searchedTerm == null || blockLine == null || searchedTerm.compareTo(blockLine.getTermBytes().getTerm()) <= 0) {
        // If the searched term precedes lexicographically the current term,
        // then reset the position to the first term line of the block.
        // If the searched term equals the current term, we also need to reset
        // to scan again the current line.
        blockReadBuffer.setPosition(blockFirstLineStart);
        lineIndexInBlock = 0;
      }
    } else {
      blockInput.seek(targetBlockStartFP);
      blockStartFP = targetBlockStartFP;
      readHeader();
      blockFirstLineStart = blockReadBuffer.getPosition();
      lineIndexInBlock = 0;
    }
  }

  protected void initializeBlockReadLazily() throws IOException {
    if (blockStartFP == -1) {
      blockInput = blockInput.clone();
      blockHeaderReader = createBlockHeaderSerializer();
      blockLineReader = createBlockLineSerializer();
      blockReadBuffer = new ByteArrayDataInput();
      termStatesReadBuffer = new ByteArrayDataInput();
      termStateSerializer = createDeltaBaseTermStateSerializer();
      scratchBlockBytes = new BytesRef();
      scratchBlockLine = new BlockLine(new TermBytes(0, scratchBlockBytes), 0);
    }
  }

  protected BlockHeader.Serializer createBlockHeaderSerializer() {
    return new BlockHeader.Serializer();
  }

  protected BlockLine.Serializer createBlockLineSerializer() {
    return new BlockLine.Serializer();
  }

  protected DeltaBaseTermStateSerializer createDeltaBaseTermStateSerializer() {
    return new DeltaBaseTermStateSerializer();
  }

  /**
   * Reads the block header.
   * Sets {@link #blockHeader}.
   *
   * @return The block header; or null if there is no block for the field anymore.
   */
  protected BlockHeader readHeader() throws IOException {
    if (blockInput.getFilePointer() > fieldMetadata.getLastBlockStartFP()) {
      return blockHeader = null;
    }
    int numBlockBytes = blockInput.readVInt();
    BytesRef blockBytesRef = decodeBlockBytesIfNeeded(numBlockBytes);
    blockReadBuffer.reset(blockBytesRef.bytes, blockBytesRef.offset, blockBytesRef.length);
    termStatesReadBuffer.reset(blockBytesRef.bytes, blockBytesRef.offset, blockBytesRef.length);
    return blockHeader = blockHeaderReader.read(blockReadBuffer, blockHeader);
  }

  protected BytesRef decodeBlockBytesIfNeeded(int numBlockBytes) throws IOException {
    scratchBlockBytes.bytes = ArrayUtil.grow(scratchBlockBytes.bytes, numBlockBytes);
    blockInput.readBytes(scratchBlockBytes.bytes, 0, numBlockBytes);
    scratchBlockBytes.length = numBlockBytes;
    if (blockDecoder == null) {
      return scratchBlockBytes;
    }
    blockReadBuffer.reset(scratchBlockBytes.bytes, 0, numBlockBytes);
    return blockDecoder.decode(blockReadBuffer, numBlockBytes);
  }

  /**
   * Reads the {@link BlockTermState} if it is not already set.
   * Sets {@link #termState}.
   */
  protected BlockTermState readTermStateIfNotRead() throws IOException {
    if (termState == null) {
      termState = readTermState();
      if (termState != null) {
        termState.termBlockOrd = lineIndexInBlock;
        termState.blockFilePointer = blockStartFP;
      }
    }
    return termState;
  }

  /**
   * Reads the {@link BlockTermState} on the current line.
   * Sets {@link #termState}.
   * <p>
   * Overriding method may return null if there is no {@link BlockTermState}
   * (in this case the extending class must support a null {@link #termState}).
   *
   * @return The {@link BlockTermState}; or null if none.
   */
  protected BlockTermState readTermState() throws IOException {
    // We reuse scratchTermState safely as the read TermState is cloned in the termState() method.
    termStatesReadBuffer.setPosition(blockFirstLineStart + blockHeader.getTermStatesBaseOffset() + blockLine.getTermStateRelativeOffset());
    return termState = termStateSerializer.readTermState(
        blockHeader.getBaseDocsFP(), blockHeader.getBasePositionsFP(), blockHeader.getBasePayloadsFP(),
        termStatesReadBuffer, fieldMetadata.getFieldInfo(), scratchTermState
    );
  }

  @Override
  public BytesRef term() {
    if (termStateForced) {
      return forcedTerm.get();
    }
    return blockLine == null ? null : blockLine.getTermBytes().getTerm();
  }

  @Override
  public long ord() {
    throw new UnsupportedOperationException();
  }

  @Override
  public int docFreq() throws IOException {
    readTermStateIfNotRead();
    return termState.docFreq;
  }

  @Override
  public long totalTermFreq() throws IOException {
    readTermStateIfNotRead();
    return termState.totalTermFreq;
  }

  @Override
  public TermState termState() throws IOException {
    readTermStateIfNotRead();
    return termState.clone();
  }

  @Override
  public PostingsEnum postings(PostingsEnum reuse, int flags) throws IOException {
    readTermStateIfNotRead();
    return postingsReader.postings(fieldMetadata.getFieldInfo(), termState, reuse, flags);
  }

  @Override
  public ImpactsEnum impacts(int flags) throws IOException {
    readTermStateIfNotRead();
    return postingsReader.impacts(fieldMetadata.getFieldInfo(), termState, flags);
  }

  @Override
  public long ramBytesUsed() {
    return BASE_RAM_USAGE
        + (blockLineReader == null ? 0 : blockLineReader.ramBytesUsed())
        + (blockReadBuffer == null ? 0 : RamUsageUtil.ramBytesUsedByByteArrayOfLength(blockReadBuffer.length()))
        + (termStateSerializer == null ? 0 : termStateSerializer.ramBytesUsed())
        + (forcedTerm == null ? 0 : RamUsageUtil.ramBytesUsed(forcedTerm))
        + (blockHeader == null ? 0 : blockHeader.ramBytesUsed())
        + (blockLine == null ? 0 : blockLine.ramBytesUsed())
        + (termState == null ? 0 : RamUsageUtil.ramBytesUsed(termState));
  }

  protected IndexDictionary.Browser getOrCreateDictionaryBrowser() throws IOException {
    if (dictionaryBrowser == null) {
      dictionaryBrowser = dictionaryBrowserSupplier.get();
    }
    return dictionaryBrowser;
  }

  /**
   * Called by the primary {@link TermsEnum} methods to clear the previous {@link TermState}.
   */
  protected void clearTermState() {
    termState = null;
    termStateForced = false;
  }

  protected CorruptIndexException newCorruptIndexException(String msg, Long fp) {
    return new CorruptIndexException(msg
        + (fp == null ? "" : " at FP " + fp)
        + " for field \"" + fieldMetadata.getFieldInfo().name + "\"", blockInput);
  }
}
