/*
 * 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.beam.runners.dataflow.worker;

import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkNotNull;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkState;

import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.NonWritableChannelException;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.SeekableByteChannel;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.NavigableMap;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.SortedMap;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import org.apache.beam.runners.dataflow.internal.IsmFormat;
import org.apache.beam.runners.dataflow.internal.IsmFormat.Footer;
import org.apache.beam.runners.dataflow.internal.IsmFormat.FooterCoder;
import org.apache.beam.runners.dataflow.internal.IsmFormat.IsmRecord;
import org.apache.beam.runners.dataflow.internal.IsmFormat.IsmRecordCoder;
import org.apache.beam.runners.dataflow.internal.IsmFormat.IsmShard;
import org.apache.beam.runners.dataflow.internal.IsmFormat.KeyPrefix;
import org.apache.beam.runners.dataflow.internal.IsmFormat.KeyPrefixCoder;
import org.apache.beam.runners.dataflow.util.RandomAccessData;
import org.apache.beam.runners.dataflow.worker.util.ScalableBloomFilter;
import org.apache.beam.runners.dataflow.worker.util.ScalableBloomFilter.ScalableBloomFilterCoder;
import org.apache.beam.runners.dataflow.worker.util.ValueInEmptyWindows;
import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader;
import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.util.VarInt;
import org.apache.beam.sdk.util.WeightedValue;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.guava.v20_0.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.MoreObjects;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Optional;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Throwables;
import org.apache.beam.vendor.guava.v20_0.com.google.common.cache.Cache;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableSortedMap;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Ordering;
import org.apache.beam.vendor.guava.v20_0.com.google.common.io.ByteStreams;
import org.apache.beam.vendor.guava.v20_0.com.google.common.primitives.Ints;
import org.apache.beam.vendor.guava.v20_0.com.google.common.primitives.Longs;

/**
 * A {@link NativeReader} that reads Ism files.
 *
 * @param <V> the type of the value written to the sink
 */
// Possible real inconsistency - https://issues.apache.org/jira/browse/BEAM-6560
@SuppressFBWarnings("IS2_INCONSISTENT_SYNC")
public class IsmReaderImpl<V> extends IsmReader<V> {
  /**
   * This constant represents the distance we would rather read and drop bytes for versus doing an
   * actual repositioning of the underlying stream. Tuned for operation within GCS.
   */
  private static final int SEEK_VS_READ = 6 * 1024 * 1024;

  static final int MAX_SHARD_INDEX_AND_FOOTER_SIZE = 1024 * 1024;

  private final ResourceId resourceId;
  private final IsmRecordCoder<V> coder;

  /** Lazily initialized on first read. */
  private long length;

  private Footer footer;

  /** A map indexed by shard id, storing the Ism shard descriptors. */
  private NavigableMap<Integer, IsmShard> shardIdToShardMap;

  /**
   * A map sorted and indexed by the block offset for a given shard. The sorting is important so the
   * {@link LazyIsmPrefixReaderIterator} can read through the file in increasing order.
   */
  private NavigableMap<Long, IsmShard> shardOffsetToShardMap;

  /** Values lazily initialized per shard on first keyed read of each shard. */
  private Map<Integer, ImmutableSortedMap<RandomAccessData, IsmShardKey>> indexPerShard;

  ScalableBloomFilter bloomFilter;

  /**
   * A cache instance which if set on this reader is used to cache blocks of data that are read.
   * Each value represents the decoded form of a block.
   */
  private final Cache<
          IsmShardKey, WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>>>>
      cache;

  /**
   * Produces a reader for the specified {@code resourceId} and {@code coder}. See {@link IsmFormat}
   * for encoded format details.
   */
  IsmReaderImpl(
      final ResourceId resourceId,
      IsmRecordCoder<V> coder,
      Cache<IsmShardKey, WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>>>>
          cache) {
    checkNotNull(cache);
    IsmFormat.validateCoderIsCompatible(coder);
    this.resourceId = resourceId;
    this.coder = coder;
    this.cache = cache;
  }

  @Override
  public IsmPrefixReaderIterator iterator() throws IOException {
    SideInputReadCounter readCounter = IsmReader.getCurrentSideInputCounter();
    return new LazyIsmPrefixReaderIterator(readCounter);
  }

  @Override
  public IsmPrefixReaderIterator overKeyComponents(List<?> keyComponents) throws IOException {
    if (keyComponents.isEmpty()) {
      return overKeyComponents(keyComponents, 0, new RandomAccessData(0));
    }
    RandomAccessData keyBytes = new RandomAccessData();
    int shardId = coder.encodeAndHash(keyComponents, keyBytes);
    return overKeyComponents(keyComponents, shardId, keyBytes);
  }

  @Override
  public IsmPrefixReaderIterator overKeyComponents(
      List<?> keyComponents, int shardId, RandomAccessData keyBytes) throws IOException {
    checkNotNull(keyComponents);
    checkNotNull(keyBytes);
    SideInputReadCounter readCounter = IsmReader.getCurrentSideInputCounter();
    if (keyComponents.isEmpty()) {
      checkArgument(
          shardId == 0 && keyBytes.size() == 0,
          "Expected shard id to be 0 and key bytes to be empty "
              + "but got shard id %s and key bytes of length %s",
          shardId,
          keyBytes.size());
    }
    checkArgument(
        keyComponents.size() <= coder.getKeyComponentCoders().size(),
        "Expected at most %s key component(s) but received %s.",
        coder.getKeyComponentCoders().size(),
        keyComponents);

    Optional<SeekableByteChannel> inChannel =
        initializeFooterAndShardIndex(Optional.<SeekableByteChannel>absent(), readCounter);

    // If this file is empty, we can return an empty iterator.
    if (footer.getNumberOfKeys() == 0) {
      return new EmptyIsmPrefixReaderIterator(keyComponents);
    }

    // If not enough key components to figure out which shard was requested, we return a reader
    // iterator over all the keys.
    if (keyComponents.size() < coder.getNumberOfShardKeyCoders(keyComponents)) {
      return new ShardAwareIsmPrefixReaderIterator(
          keyComponents, openIfNeeded(inChannel), readCounter);
    }

    // If this file does not contain the shard or Bloom filter does not contain the key prefix,
    // we know that we can return an empty reader iterator.
    if (!shardIdToShardMap.containsKey(shardId)) {
      return new EmptyIsmPrefixReaderIterator(keyComponents);
    }
    inChannel = initializeForKeyedRead(shardId, inChannel, readCounter);
    closeIfPresent(inChannel);
    if (!bloomFilterMightContain(keyBytes)) {
      return new EmptyIsmPrefixReaderIterator(keyComponents);
    }

    // Otherwise we may actually contain the key so construct a reader iterator
    // which will fetch the data blocks containing the requested key prefix.

    // We find the first key in the index which may contain our prefix
    RandomAccessData floorKey = indexPerShard.get(shardId).floorKey(keyBytes);

    // We compute an upper bound on the key prefix by incrementing the prefix
    RandomAccessData keyBytesUpperBound = keyBytes.increment();
    // Compute the sub-range of the index map that we want to iterate over since
    // any of these blocks may contain the key prefix.
    Iterator<IsmShardKey> blockEntries =
        indexPerShard.get(shardId).subMap(floorKey, keyBytesUpperBound).values().iterator();

    return new WithinShardIsmPrefixReaderIterator(
        keyComponents, keyBytes, keyBytesUpperBound, blockEntries, readCounter);
  }

  /** Returns whether this ISM reader has been initialized. */
  @Override
  public boolean isInitialized() {
    return footer != null;
  }

  @Override
  public boolean isEmpty() throws IOException {
    SideInputReadCounter readCounter = IsmReader.getCurrentSideInputCounter();
    closeIfPresent(
        initializeFooterAndShardIndex(Optional.<SeekableByteChannel>absent(), readCounter));
    return footer.getNumberOfKeys() == 0;
  }

  @Override
  public ResourceId getResourceId() {
    return resourceId;
  }

  @Override
  public IsmRecordCoder<V> getCoder() {
    return coder;
  }

  @Override
  public String toString() {
    return MoreObjects.toStringHelper(IsmReaderImpl.class)
        .add("resource id", resourceId)
        .add("coder", coder)
        .toString();
  }

  // Overridable by tests to get around the bloom filter not containing any values.
  @VisibleForTesting
  boolean bloomFilterMightContain(RandomAccessData keyBytes) {
    return bloomFilter.mightContain(keyBytes.array(), 0, keyBytes.size());
  }

  /**
   * Initialize this Ism reader by reading the footer and shard index. Returns a channel for re-use
   * if this method was required to open one.
   */
  private synchronized Optional<SeekableByteChannel> initializeFooterAndShardIndex(
      Optional<SeekableByteChannel> inChannel, SideInputReadCounter readCounter)
      throws IOException {
    if (footer != null) {
      checkState(
          shardIdToShardMap != null, "Expected shard id to shard map to have been initialized.");
      checkState(
          shardOffsetToShardMap != null,
          "Expected shard offset to shard map to have been initialized.");
      return inChannel;
    }
    checkState(
        shardIdToShardMap == null, "Expected shard id to shard map to not have been initialized.");
    checkState(
        shardOffsetToShardMap == null,
        "Expected shard offset to shard map to not have been initialized.");

    SeekableByteChannel rawChannel;
    RandomAccessData data;
    long startPosition;
    try (Closeable closeReadCounter = readCounter.enter()) {
      rawChannel = openIfNeeded(inChannel);
      this.length = rawChannel.size();

      // We read the last chunk of data, for small files we will capture the entire file.
      // We may capture the Bloom filter, shard index, and footer for slightly larger files.
      // Otherwise we are guaranteed to capture the footer and the shard index.
      startPosition = Math.max(length - MAX_SHARD_INDEX_AND_FOOTER_SIZE, 0);
      position(rawChannel, startPosition);
      data = new RandomAccessData(ByteStreams.toByteArray(Channels.newInputStream(rawChannel)));
    }
    readCounter.addBytesRead(data.size());
    // Read the fixed length footer.
    this.footer =
        FooterCoder.of()
            .decode(data.asInputStream(data.size() - Footer.FIXED_LENGTH, Footer.FIXED_LENGTH));

    checkState(
        startPosition < footer.getIndexPosition(),
        "Malformed file, expected to have been able to read entire shard index.");
    int offsetWithinReadData = (int) (footer.getIndexPosition() - startPosition);

    // Decode the list of Ism shard descriptors
    List<IsmShard> ismShards =
        IsmFormat.ISM_SHARD_INDEX_CODER.decode(
            data.asInputStream(offsetWithinReadData, data.size() - offsetWithinReadData));

    // Build the shard id to shard descriptor map
    ImmutableSortedMap.Builder<Integer, IsmShard> shardIdToShardMapBuilder =
        ImmutableSortedMap.orderedBy(Ordering.<Integer>natural());
    for (IsmShard ismShard : ismShards) {
      shardIdToShardMapBuilder.put(ismShard.getId(), ismShard);
    }
    shardIdToShardMap = shardIdToShardMapBuilder.build();

    // Build the shard block offset to shard descriptor map
    ImmutableSortedMap.Builder<Long, IsmShard> shardOffsetToShardMapBuilder =
        ImmutableSortedMap.orderedBy(Ordering.<Long>natural());
    for (IsmShard ismShard : ismShards) {
      shardOffsetToShardMapBuilder.put(ismShard.getBlockOffset(), ismShard);
    }
    shardOffsetToShardMap = shardOffsetToShardMapBuilder.build();

    // We may have gotten the Bloom filter, if so lets store it.
    if (startPosition < footer.getBloomFilterPosition()) {
      Optional<SeekableByteChannel> cachedDataChannel =
          Optional.<SeekableByteChannel>of(
              new CachedTailSeekableByteChannel(startPosition, data.array()));
      initializeBloomFilterAndIndexPerShard(cachedDataChannel);

      // For small files, we may have read the whole thing during initialization so
      // lets cache all this information. Note that this is important for the many small files
      // case since the IsmSideInputReader does initialization in parallel.
      if (cache != null && startPosition == 0) {
        for (IsmShard ismShard : ismShards) {
          initializeForKeyedRead(ismShard.getId(), cachedDataChannel, readCounter);
        }
        for (SortedMap<RandomAccessData, IsmShardKey> shards : indexPerShard.values()) {
          for (Map.Entry<RandomAccessData, IsmShardKey> block : shards.entrySet()) {
            cache.put(
                block.getValue(),
                new IsmCacheLoader(block.getValue()).call(cachedDataChannel.get()));
          }
        }
      }
    }

    return Optional.of(rawChannel);
  }

  /**
   * Initializes the Bloom filter and index per shard. We prepopulate empty indices for shards where
   * the index offset matches the following shard block offset. Re-uses the provided channel,
   * returning it or a new one if this method was required to open one.
   */
  private synchronized Optional<SeekableByteChannel> initializeBloomFilterAndIndexPerShard(
      Optional<SeekableByteChannel> inChannel) throws IOException {
    if (indexPerShard != null) {
      checkState(bloomFilter != null, "Expected Bloom filter to have been initialized.");
      return inChannel;
    }

    SeekableByteChannel rawChannel = openIfNeeded(inChannel);

    // Set the position to where the bloom filter is and read it in.
    position(rawChannel, footer.getBloomFilterPosition());
    bloomFilter = ScalableBloomFilterCoder.of().decode(Channels.newInputStream(rawChannel));

    indexPerShard = new HashMap<>();
    // If a shard is small, it may not contain an index and we can detect this and
    // prepopulate the shard index map with an empty entry if the start of the index
    // and start of the next block are equal
    Iterator<IsmShard> shardIterator = shardOffsetToShardMap.values().iterator();

    // If file is empty we just return here.
    if (!shardIterator.hasNext()) {
      return Optional.of(rawChannel);
    }

    // If the current shard's index position is equal to the next shards block offset
    // then we know that the index contains no data and we can pre-populate it with
    // the empty map.
    IsmShard currentShard = shardIterator.next();
    while (shardIterator.hasNext()) {
      IsmShard nextShard = shardIterator.next();
      if (currentShard.getIndexOffset() == nextShard.getBlockOffset()) {
        indexPerShard.put(
            currentShard.getId(),
            ImmutableSortedMap.<RandomAccessData, IsmShardKey>orderedBy(
                    RandomAccessData.UNSIGNED_LEXICOGRAPHICAL_COMPARATOR)
                .put(
                    new RandomAccessData(0),
                    new IsmShardKey(
                        IsmReaderImpl.this.resourceId.toString(),
                        new RandomAccessData(0),
                        currentShard.getBlockOffset(),
                        currentShard.getIndexOffset()))
                .build());
      }
      currentShard = nextShard;
    }

    // Add an entry for the last shard if its index offset is equivalent to the
    // start of the Bloom filter, then we know that the index is empty.
    if (currentShard.getIndexOffset() == footer.getBloomFilterPosition()) {
      indexPerShard.put(
          currentShard.getId(),
          ImmutableSortedMap.<RandomAccessData, IsmShardKey>orderedBy(
                  RandomAccessData.UNSIGNED_LEXICOGRAPHICAL_COMPARATOR)
              .put(
                  new RandomAccessData(0),
                  new IsmShardKey(
                      IsmReaderImpl.this.resourceId.toString(),
                      new RandomAccessData(0),
                      currentShard.getBlockOffset(),
                      currentShard.getIndexOffset()))
              .build());
    }

    return Optional.of(rawChannel);
  }

  /** A unique key used to fully describe an Ism shard. */
  static final class IsmShardKey {
    private final String resourceId;
    private final RandomAccessData firstKey;
    private final long startOffset;
    private final long endOffset;

    IsmShardKey(String resourceId, RandomAccessData firstKey, long startOffset, long endOffset) {
      this.resourceId = resourceId;
      this.firstKey = firstKey;
      this.startOffset = startOffset;
      this.endOffset = endOffset;
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper(IsmShardKey.class)
          .add("resource id", resourceId)
          .add("firstKey", firstKey)
          .add("startOffset", startOffset)
          .add("endOffset", endOffset)
          .toString();
    }

    @Override
    public boolean equals(Object obj) {
      if (!(obj instanceof IsmShardKey)) {
        return false;
      }
      IsmShardKey key = (IsmShardKey) obj;
      return startOffset == key.startOffset
          && endOffset == key.endOffset
          && Objects.equals(resourceId, key.resourceId)
          && Objects.equals(firstKey, key.firstKey);
    }

    @Override
    public int hashCode() {
      int result = resourceId.hashCode();
      result = result * 31 + firstKey.hashCode();
      result = result * 31 + Longs.hashCode(startOffset);
      result = result * 31 + Longs.hashCode(endOffset);
      return result;
    }
  }

  /**
   * Initializes the footer, shard index, Bloom filter and index for the requested shard id if they
   * have not been initialized yet. Re-uses the provided channel, returning it or a new one if this
   * method was required to open one.
   */
  // Real bug - https://issues.apache.org/jira/browse/BEAM-6559
  @SuppressFBWarnings("NP_NULL_ON_SOME_PATH")
  private Optional<SeekableByteChannel> initializeForKeyedRead(
      int shardId, Optional<SeekableByteChannel> inChannel, SideInputReadCounter readCounter)
      throws IOException {
    inChannel = initializeFooterAndShardIndex(inChannel, readCounter);

    IsmShard shardWithIndex = shardIdToShardMap.get(shardId);

    // If this shard id is not within this file, we can return immediately.
    if (shardWithIndex == null) {
      return inChannel;
    }

    inChannel = initializeBloomFilterAndIndexPerShard(inChannel);

    // If the index has been populated and contains the shard id, we can return.
    if (indexPerShard != null && indexPerShard.containsKey(shardId)) {
      checkState(bloomFilter != null, "Bloom filter expected to have been initialized.");
      return inChannel;
    }

    checkState(
        indexPerShard.get(shardId) == null,
        "Expected to not have initialized index for shard %s",
        shardId);

    Long startOfNextBlock = shardOffsetToShardMap.higherKey(shardWithIndex.getBlockOffset());
    // If this is the last block, then we need to grab the position of the Bloom filter
    // as the upper bound.
    if (startOfNextBlock == null) {
      startOfNextBlock = footer.getBloomFilterPosition();
    }

    // Open the channel if needed and seek to the start of the index.
    SeekableByteChannel rawChannel = openIfNeeded(inChannel);
    rawChannel.position(shardWithIndex.getIndexOffset());
    InputStream inStream = Channels.newInputStream(rawChannel);

    ImmutableSortedMap.Builder<RandomAccessData, IsmShardKey> builder =
        ImmutableSortedMap.orderedBy(RandomAccessData.UNSIGNED_LEXICOGRAPHICAL_COMPARATOR);

    // Read the first key
    RandomAccessData currentKeyBytes = new RandomAccessData();
    readKey(inStream, currentKeyBytes);
    long currentOffset = VarInt.decodeLong(inStream);

    // Insert the entry that happens at the beginning limiting the shard block by the
    // first keys block offset.
    builder.put(
        new RandomAccessData(0),
        new IsmShardKey(
            IsmReaderImpl.this.resourceId.toString(),
            new RandomAccessData(0),
            shardWithIndex.getBlockOffset(),
            currentOffset));

    // While another index entry exists, insert an index entry with the key, and offsets
    // that limit the range of the shard block.
    while (rawChannel.position() < startOfNextBlock) {
      RandomAccessData nextKeyBytes = currentKeyBytes.copy();
      readKey(inStream, nextKeyBytes);
      long nextOffset = VarInt.decodeLong(inStream);

      builder.put(
          currentKeyBytes,
          new IsmShardKey(
              IsmReaderImpl.this.resourceId.toString(),
              currentKeyBytes,
              currentOffset,
              nextOffset));

      currentKeyBytes = nextKeyBytes;
      currentOffset = nextOffset;
    }

    // Upper bound the last entry with the index offset.
    builder.put(
        currentKeyBytes,
        new IsmShardKey(
            IsmReaderImpl.this.resourceId.toString(),
            currentKeyBytes,
            currentOffset,
            shardWithIndex.getIndexOffset()));
    indexPerShard.put(shardId, builder.build());

    return Optional.of(rawChannel);
  }

  /** A function which takes an IsmShardKey fully describing a data block to read and return. */
  private class IsmCacheLoader
      implements Callable<
          WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>>>> {

    private final IsmShardKey key;

    private IsmCacheLoader(IsmShardKey key) {
      this.key = key;
    }

    @Override
    public WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>>> call()
        throws IOException {
      // Open a channel and build a sorted map from key to value for each key value
      // pair found within the data block.
      try (SeekableByteChannel rawChannel = open()) {
        return call(rawChannel);
      }
    }

    public WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>>> call(
        SeekableByteChannel rawChannel) throws IOException {
      SideInputReadCounter readCounter = IsmReader.getCurrentSideInputCounter();
      try (WithinShardIsmReaderIterator readerIterator =
          new WithinShardIsmReaderIterator(
              rawChannel, key.firstKey, key.startOffset, key.endOffset, readCounter)) {

        ImmutableSortedMap.Builder<RandomAccessData, WindowedValue<IsmRecord<V>>> mapBuilder =
            ImmutableSortedMap.orderedBy(RandomAccessData.UNSIGNED_LEXICOGRAPHICAL_COMPARATOR);
        for (boolean more = readerIterator.start(); more; more = readerIterator.advance()) {
          RandomAccessData nextKey = readerIterator.getCurrentKeyBytes().copy();
          WindowedValue<IsmRecord<V>> next = readerIterator.getCurrent();
          mapBuilder.put(nextKey, next);
        }
        // We return the size of the data block as the weight of this data block.
        return WeightedValue.of(
            (NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>>) mapBuilder.build(),
            key.endOffset - key.startOffset);
      }
    }
  }

  /**
   * Fetches the data block requested.
   *
   * <p>If the cache is available, we will load and cache the requested block. Otherwise, we will
   * load and return the block.
   */
  private NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>> fetch(IsmShardKey key)
      throws IOException {
    try {
      if (cache == null) {
        return new IsmCacheLoader(key).call().getValue();
      } else {
        return cache.get(key, new IsmCacheLoader(key)).getValue();
      }
    } catch (ExecutionException e) {
      // Try and re-throw the root cause if its an IOException
      Throwables.propagateIfPossible(e.getCause(), IOException.class);
      throw new IOException(e.getCause());
    }
  }

  /** The base class of Ism reader iterators which operate over a given key prefix. */
  abstract class IsmPrefixReaderIteratorImpl extends IsmPrefixReaderIterator {
    private final List<?> keyComponents;
    final SideInputReadCounter readCounter;

    private IsmPrefixReaderIteratorImpl(List<?> keyComponents, SideInputReadCounter readCounter) {
      this.keyComponents = keyComponents;
      this.readCounter = readCounter;
    }

    /** Returns the list of key components representing this iterators key prefix. */
    protected List<?> getKeyComponents() {
      return keyComponents;
    }

    /**
     * Concatenates this reader iterators key components with the additionally supplied key
     * components and encodes them into their byte representations producing a key. Returns the
     * exact record represented by the key generated above.
     *
     * <p>Null is returned if no key has all the key components as a prefix within this file.
     */
    @Override
    public final WindowedValue<IsmRecord<V>> get(List<?> additionalKeyComponents)
        throws IOException {
      RandomAccessData keyBytes = new RandomAccessData();
      try (Closeable readerCloser = readCounter.enter()) {
        int shardId =
            coder.encodeAndHash(
                ImmutableList.builder()
                    .addAll(keyComponents)
                    .addAll(additionalKeyComponents)
                    .build(),
                keyBytes);
        return getBlock(keyBytes, shardId, readCounter).get(keyBytes);
      }
    }

    /**
     * Returns the record for the last key having this iterators key prefix. Last is defined as the
     * largest key with the same key prefix when comparing key's byte representations using an
     * unsigned lexicographical byte order.
     *
     * <p>Null is returned if the prefix is not present within this file.
     */
    @Override
    public WindowedValue<IsmRecord<V>> getLast() throws IOException {
      RandomAccessData keyBytes = new RandomAccessData();
      int shardId = coder.encodeAndHash(keyComponents, keyBytes);

      Optional<SeekableByteChannel> inChannel =
          initializeFooterAndShardIndex(Optional.<SeekableByteChannel>absent(), readCounter);

      // Key is not stored here
      if (!shardIdToShardMap.containsKey(shardId) || !bloomFilterMightContain(keyBytes)) {
        return null;
      }

      inChannel = initializeForKeyedRead(shardId, inChannel, readCounter);
      closeIfPresent(inChannel);

      final NavigableMap<RandomAccessData, IsmShardKey> indexInShard = indexPerShard.get(shardId);
      RandomAccessData end = keyBytes.increment();
      final IsmShardKey cacheEntry = indexInShard.floorEntry(end).getValue();

      NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>> block;
      try (Closeable readerCloser = IsmReader.setSideInputReadContext(readCounter)) {
        block = fetch(cacheEntry);
      }

      RandomAccessData lastKey = block.lastKey();

      // If the requested key is greater than the last key within the block, then it
      // does not exist.
      if (RandomAccessData.UNSIGNED_LEXICOGRAPHICAL_COMPARATOR.compare(keyBytes, lastKey) > 0) {
        return null;
      }

      Entry<RandomAccessData, WindowedValue<IsmRecord<V>>> rval = block.floorEntry(end);

      // If the prefix matches completely then we can return
      if (RandomAccessData.UNSIGNED_LEXICOGRAPHICAL_COMPARATOR.commonPrefixLength(
              keyBytes, rval.getKey())
          == keyBytes.size()) {
        return rval.getValue();
      }
      return null;
    }
  }

  /** An empty reader iterator. */
  class EmptyIsmPrefixReaderIterator extends IsmPrefixReaderIteratorImpl {
    private EmptyIsmPrefixReaderIterator(List<?> keyComponents) {
      super(keyComponents, NoopSideInputReadCounter.INSTANCE);
    }

    @Override
    public boolean start() throws IOException {
      return false;
    }

    @Override
    public boolean advance() throws IOException {
      return false;
    }

    @Override
    public WindowedValue<IsmRecord<V>> getCurrent() throws NoSuchElementException {
      throw new NoSuchElementException();
    }
  }

  /** A reader iterator which initializes its input stream lazily. */
  class LazyIsmPrefixReaderIterator extends IsmPrefixReaderIteratorImpl {
    private IsmPrefixReaderIterator delegate;

    public LazyIsmPrefixReaderIterator(SideInputReadCounter readCounter) {
      super(ImmutableList.of(), readCounter);
    }

    @Override
    public boolean start() throws IOException {
      checkState(delegate == null, "Already started");
      try (Closeable counterCloser = IsmReader.setSideInputReadContext(readCounter)) {
        delegate = overKeyComponents(getKeyComponents());
      }
      return delegate.start();
    }

    @Override
    public boolean advance() throws IOException {
      checkState(delegate != null, "Not started");
      return delegate.advance();
    }

    @Override
    public WindowedValue<IsmRecord<V>> getCurrent() {
      checkState(delegate != null, "Not started");
      return delegate.getCurrent();
    }

    @Override
    public void close() throws IOException {
      if (delegate != null) {
        delegate.close();
      }
    }
  }

  /**
   * Returns a map from key to value, where the keys are in increasing lexicographical order. If the
   * requested key is not contained within this file, an empty map is returned.
   */
  private NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>> getBlock(
      RandomAccessData keyBytes, int shardId, SideInputReadCounter readCounter) throws IOException {
    Optional<SeekableByteChannel> inChannel =
        initializeFooterAndShardIndex(Optional.<SeekableByteChannel>absent(), readCounter);

    // Key is not stored here so return an empty map.
    if (!shardIdToShardMap.containsKey(shardId) || !bloomFilterMightContain(keyBytes)) {
      return ImmutableSortedMap.<RandomAccessData, WindowedValue<IsmRecord<V>>>orderedBy(
              RandomAccessData.UNSIGNED_LEXICOGRAPHICAL_COMPARATOR)
          .build();
    }

    inChannel = initializeForKeyedRead(shardId, inChannel, readCounter);
    closeIfPresent(inChannel);

    final NavigableMap<RandomAccessData, IsmShardKey> indexInShard = indexPerShard.get(shardId);
    final IsmShardKey cacheEntry = indexInShard.floorEntry(keyBytes).getValue();
    try (Closeable readerCloseable = IsmReader.setSideInputReadContext(readCounter)) {
      return fetch(cacheEntry);
    }
  }

  /**
   * A reader iterator that returns all elements from prefix (inclusive) to prefixUpperBound
   * (exclusive) within the set of block entries provided.
   */
  private class WithinShardIsmPrefixReaderIterator extends IsmPrefixReaderIteratorImpl {
    private final Iterator<IsmShardKey> blockEntriesIterator;
    Iterator<WindowedValue<IsmRecord<V>>> iterator;
    private final RandomAccessData prefix;
    private final RandomAccessData prefixUpperBound;
    private Optional<WindowedValue<IsmRecord<V>>> current;

    private WithinShardIsmPrefixReaderIterator(
        List<?> keyComponents,
        RandomAccessData prefix,
        RandomAccessData prefixUpperBound,
        Iterator<IsmShardKey> blockEntriesIterator,
        SideInputReadCounter readCounter) {
      super(keyComponents, readCounter);
      checkNotNull(blockEntriesIterator);
      this.prefix = prefix;
      this.prefixUpperBound = prefixUpperBound;
      this.blockEntriesIterator = blockEntriesIterator;
    }

    @Override
    public boolean start() throws IOException {
      return advance();
    }

    @Override
    public boolean advance() throws IOException {
      // This is in a while loop because the blocks that we are asked to look into may
      // not contain the key prefix.
      while (iterator == null || !iterator.hasNext()) {
        // If there are no blocks to iterate over we can return false
        if (!blockEntriesIterator.hasNext()) {
          return false;
        }

        NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>> map;
        try (Closeable counterCloseable = IsmReader.setSideInputReadContext(readCounter)) {
          IsmShardKey nextBlock = blockEntriesIterator.next();
          map = fetch(nextBlock);
        }
        SortedMap<RandomAccessData, WindowedValue<IsmRecord<V>>> submap =
            map.subMap(prefix, prefixUpperBound);
        Collection<WindowedValue<IsmRecord<V>>> values = submap.values();
        iterator = values.iterator();
      }
      current = Optional.of(iterator.next());
      return true;
    }

    @Override
    public WindowedValue<IsmRecord<V>> getCurrent() throws NoSuchElementException {
      if (!current.isPresent()) {
        throw new NoSuchElementException();
      }
      return current.get();
    }
  }

  /** A reader iterator that returns all records across all shards contained within this file. */
  private class ShardAwareIsmPrefixReaderIterator extends IsmPrefixReaderIteratorImpl {
    private final SeekableByteChannel rawChannel;
    private WithinShardIsmReaderIterator delegate;
    private Iterator<IsmShard> shardEntries;

    private ShardAwareIsmPrefixReaderIterator(
        List<?> keyComponents, SeekableByteChannel rawChannel, SideInputReadCounter readCounter)
        throws IOException {
      super(keyComponents, readCounter);
      checkState(
          shardOffsetToShardMap.size() > 0,
          "Expected that shard offset to shard map has been initialized and is not empty.");

      this.rawChannel = rawChannel;
      this.shardEntries = shardOffsetToShardMap.values().iterator();
      IsmShard firstShard = shardEntries.next();
      delegate =
          new WithinShardIsmReaderIterator(
              rawChannel,
              new RandomAccessData(),
              firstShard.getBlockOffset(),
              firstShard.getIndexOffset(),
              readCounter);
    }

    @Override
    public boolean start() throws IOException {
      checkState(delegate.start(), "Expected each shard to contain at least one entry");
      return true;
    }

    @Override
    public boolean advance() throws IOException {
      if (delegate.advance()) {
        return true;
      }
      // If our current shard index is empty, we need to move to the next one.
      if (!shardEntries.hasNext()) {
        return false;
      }
      IsmShard nextIsmShard = shardEntries.next();
      checkState(
          nextIsmShard.getBlockOffset() >= rawChannel.position(),
          "Expected channel read order to be sequential.");

      delegate =
          new WithinShardIsmReaderIterator(
              rawChannel,
              new RandomAccessData(),
              nextIsmShard.getBlockOffset(),
              nextIsmShard.getIndexOffset(),
              readCounter);

      checkState(delegate.start(), "Expected each shard to contain at least one entry.");
      return true;
    }

    @Override
    public WindowedValue<IsmRecord<V>> getCurrent() throws NoSuchElementException {
      return delegate.getCurrent();
    }

    @Override
    public void close() throws IOException {
      rawChannel.close();
    }

    @Override
    public WindowedValue<IsmRecord<V>> getLast() throws IOException {
      // Since this is an iterator over all of the file, we return the last record within the file.
      try (Closeable readerCloser = readCounter.enter()) {
        int lastShardId = shardOffsetToShardMap.lastEntry().getValue().getId();
        initializeForKeyedRead(lastShardId, Optional.<SeekableByteChannel>absent(), readCounter);
        NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>> lastBlock =
            getBlock(indexPerShard.get(lastShardId).lastKey(), lastShardId, readCounter);
        return lastBlock.lastEntry().getValue();
      }
    }
  }

  /**
   * A reader iterator for Ism formatted files which returns a sequence of Ism records from the
   * underlying channel from where it is currently positioned till the supplied limit.
   *
   * <p>This reader iterator will use the supplied bytes as the bytes for the previous key which is
   * required to do the KeyPrefix based key decoding.
   */
  private class WithinShardIsmReaderIterator
      extends NativeReaderIterator<WindowedValue<IsmRecord<V>>> {
    private final SeekableByteChannel rawChannel;
    private final InputStream inStream;
    private final SideInputReadCounter readCounter;
    private long readLimit;
    private RandomAccessData keyBytes;
    private long position;
    private Optional<WindowedValue<IsmRecord<V>>> current;

    private WithinShardIsmReaderIterator(
        SeekableByteChannel rawChannel,
        RandomAccessData currentKeyBytes,
        long newPosition,
        long newLimit,
        SideInputReadCounter readCounter)
        throws IOException {
      checkNotNull(rawChannel);
      checkNotNull(currentKeyBytes);
      checkArgument(newLimit >= 0L);
      checkArgument(newPosition <= newLimit);
      this.rawChannel = rawChannel;
      this.inStream = Channels.newInputStream(rawChannel);
      this.keyBytes = currentKeyBytes.copy();
      this.position = newPosition;
      this.readLimit = newLimit;
      this.readCounter = readCounter;
      IsmReaderImpl.position(rawChannel, newPosition);
    }

    @Override
    public boolean start() throws IOException {
      return advance();
    }

    @Override
    public boolean advance() throws IOException, NoSuchElementException {
      checkState(position <= readLimit, "Read past end of stream");
      if (position == readLimit) {
        current = Optional.absent();
        return false;
      }

      final IsmRecord<V> ismRecord;
      readKey(inStream, keyBytes);
      InputStream keyInputStream = keyBytes.asInputStream(0, keyBytes.size());
      List<Object> keyComponents = new ArrayList<>(coder.getKeyComponentCoders().size());
      for (int i = 0; i < coder.getKeyComponentCoders().size(); ++i) {
        keyComponents.add(coder.getKeyComponentCoder(i).decode(keyInputStream));
      }

      if (IsmFormat.isMetadataKey(keyComponents)) {
        byte[] metadata = ByteArrayCoder.of().decode(inStream);
        ismRecord = IsmRecord.<V>meta(keyComponents, metadata);
      } else {
        V value = coder.getValueCoder().decode(inStream);
        ismRecord = IsmRecord.<V>of(keyComponents, value);
      }

      long newPosition = rawChannel.position();
      IsmReaderImpl.this.notifyElementRead(newPosition - position);
      readCounter.addBytesRead(newPosition - position);
      position = newPosition;
      current = Optional.of(new ValueInEmptyWindows<>(ismRecord));
      return true;
    }

    @Override
    public WindowedValue<IsmRecord<V>> getCurrent() {
      if (!current.isPresent()) {
        throw new NoSuchElementException();
      }
      return current.get();
    }

    public RandomAccessData getCurrentKeyBytes() {
      if (!current.isPresent()) {
        throw new NoSuchElementException();
      }
      return keyBytes;
    }
  }

  /**
   * Decodes a KeyPrefix from the stream and then reads unshared key bytes from the stream placing
   * them into the supplied keyBytes at position keyBytes[shared key bytes : shared key bytes +
   * unshared key bytes].
   */
  private static void readKey(InputStream inStream, RandomAccessData keyBytes) throws IOException {
    KeyPrefix keyPrefix = KeyPrefixCoder.of().decode(inStream);
    // currentKey = prevKey[0 : sharedKeySize] + read(unsharedKeySize)
    keyBytes.readFrom(
        inStream,
        keyPrefix.getSharedKeySize() /* start to overwrite the previous key at sharedKeySize */,
        keyPrefix.getUnsharedKeySize() /* read unsharedKeySize bytes from the stream */);
    // Reset the length incase the next key was shorter.
    keyBytes.resetTo(keyPrefix.getSharedKeySize() + keyPrefix.getUnsharedKeySize());
  }

  /** Closes the underlying channel if present. */
  private void closeIfPresent(Optional<SeekableByteChannel> inChannel) throws IOException {
    if (inChannel.isPresent()) {
      inChannel.get().close();
    }
  }

  /** Returns a channel, opening a new one if required. */
  private SeekableByteChannel openIfNeeded(Optional<SeekableByteChannel> inChannel)
      throws IOException {
    if (inChannel.isPresent()) {
      return inChannel.get();
    }
    return open();
  }

  /** Opens a new channel. */
  private SeekableByteChannel open() throws IOException {
    ReadableByteChannel channel = FileSystems.open(resourceId);
    Preconditions.checkArgument(
        channel instanceof SeekableByteChannel,
        "IsmReaderImpl requires a SeekableByteChannel for path %s but received %s.",
        resourceId,
        channel);
    return (SeekableByteChannel) channel;
  }

  /**
   * Seeks into the channel intelligently by either resetting the position or reading and discarding
   * bytes.
   */
  private static void position(SeekableByteChannel inChannel, long newPosition) throws IOException {
    long currentPosition = inChannel.position();
    // If just doing a read is cheaper discarding the bytes lets just do the read
    if (currentPosition < newPosition && newPosition - currentPosition <= SEEK_VS_READ) {
      ByteStreams.skipFully(Channels.newInputStream(inChannel), newPosition - currentPosition);
    } else {
      // Otherwise we will perform a seek
      inChannel.position(newPosition);
    }
  }

  /**
   * A {@link SeekableByteChannel} which uses a cached data segment representing the tail of a
   * {@link ReadableByteChannel}. Note that this channel only supports read operations. Closing this
   * channel is a no-op.
   */
  static class CachedTailSeekableByteChannel implements SeekableByteChannel {
    final long offset;
    final byte[] data;
    // Represents the relative position inside of data
    int localPosition;

    CachedTailSeekableByteChannel(long offset, byte[] data) {
      checkArgument(offset >= 0, "Offset must be non-negative.");
      checkNotNull(data, "Cached data must not be null.");
      this.offset = offset;
      this.data = data;
    }

    @Override
    public boolean isOpen() {
      return true;
    }

    @Override
    public void close() throws IOException {}

    @Override
    public int read(ByteBuffer dst) throws IOException {
      if (localPosition >= data.length) {
        return -1;
      }
      int length = Math.min(dst.remaining(), data.length - localPosition);
      dst.put(data, localPosition, length);
      localPosition += length;
      return length;
    }

    @Override
    public int write(ByteBuffer src) throws IOException {
      throw new NonWritableChannelException();
    }

    @Override
    public long position() throws IOException {
      // The external position is our relative position inside of data plus the offset.
      return offset + localPosition;
    }

    @Override
    public SeekableByteChannel position(long newPosition) throws IOException {
      checkArgument(
          newPosition >= offset && newPosition <= offset + data.length,
          "Cannot seek to position %s which is outside of cached data range [%s, %s].",
          newPosition,
          offset,
          offset + data.length);
      localPosition = Ints.checkedCast(newPosition - offset);
      return this;
    }

    @Override
    public long size() throws IOException {
      // The external size is the offset plus the amount of data that we are caching.
      return offset + data.length;
    }

    @Override
    public SeekableByteChannel truncate(long size) throws IOException {
      throw new NonWritableChannelException();
    }
  }
}