/**
 * 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.avro.hadoop.file;

import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStream;
import java.util.Iterator;
import java.util.Map;
import java.util.NavigableMap;
import java.util.TreeMap;

import org.apache.avro.Schema;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.specific.SpecificData;
import org.apache.avro.io.DatumReader;
import org.apache.avro.io.DatumWriter;
import org.apache.avro.file.CodecFactory;
import org.apache.avro.file.DataFileReader;
import org.apache.avro.file.DataFileWriter;
import org.apache.avro.hadoop.util.AvroCharSequenceComparator;
import org.apache.avro.hadoop.io.AvroKeyValue;
import org.apache.avro.mapred.FsInput;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A SortedKeyValueFile is an indexed Avro container file of KeyValue records
 * sorted by key.
 *
 * <p>The SortedKeyValueFile is a directory with two files, named 'data' and
 * 'index'. The 'data' file is an ordinary Avro container file with
 * records. Each record has exactly two fields, 'key' and 'value'. The keys are
 * sorted lexicographically. The 'index' file is a small Avro container file
 * mapping keys in the 'data' file to their byte positions. The index file is
 * intended to fit in memory, so it should remain small. There is one entry in
 * the index file for each data block in the Avro container file.</p>
 *
 * <p>SortedKeyValueFile is to Avro container file as MapFile is to
 * SequenceFile.</p>
 */
public class SortedKeyValueFile {
  private static final Logger LOG = LoggerFactory.getLogger(SortedKeyValueFile.class);

  /** The name of the data file within the SortedKeyValueFile directory. */
  public static final String DATA_FILENAME = "data";

  /** The name of the index file within the SortedKeyValueFile directory. */
  public static final String INDEX_FILENAME = "index";

  /**
   * Reads a SortedKeyValueFile by loading the key index into memory.
   *
   * <p>When doing a lookup, this reader finds the correct block in the data file using
   * the key index. It performs a single disk seek to the block and loads the entire block
   * into memory. The block is scanned until the key is found or is determined not to
   * exist.</p>
   *
   * @param <K> The key type.
   * @param <V> The value type.
   */
  public static class Reader<K, V> implements Closeable, Iterable<AvroKeyValue<K, V>> {
    /** The index from key to its byte offset into the data file. */
    private final NavigableMap<K, Long> mIndex;

    /** The reader for the data file. */
    private final DataFileReader<GenericRecord> mDataFileReader;

    /** The key schema for the data file. */
    private final Schema mKeySchema;

    /** The model for the data. */
    private GenericData model;

    /** A class to encapsulate the options of a Reader. */
    public static class Options {
      /** The configuration. */
      private Configuration mConf;

      /** The path to the SortedKeyValueFile to read. */
      private Path mPath;

      /** The reader schema for the key. */
      private Schema mKeySchema;

      /** The reader schema for the value. */
      private Schema mValueSchema;

      /** The model for the data. */
      private GenericData model = SpecificData.get();

      /**
       * Sets the configuration.
       *
       * @param conf The configuration.
       * @return This options instance.
       */
      public Options withConfiguration(Configuration conf) {
        mConf = conf;
        return this;
      }

      /**
       * Gets the configuration.
       *
       * @return The configuration.
       */
      public Configuration getConfiguration() {
        return mConf;
      }

      /**
       * Sets the input path.
       *
       * @param path The input path.
       * @return This options instance.
       */
      public Options withPath(Path path) {
        mPath = path;
        return this;
      }

      /**
       * Gets the input path.
       *
       * @return The input path.
       */
      public Path getPath() {
        return mPath;
      }

      /**
       * Sets the reader schema for the key.
       *
       * @param keySchema The reader schema for the key.
       * @return This options instance.
       */
      public Options withKeySchema(Schema keySchema) {
        mKeySchema = keySchema;
        return this;
      }

      /**
       * Gets the reader schema for the key.
       *
       * @return The reader schema for the key.
       */
      public Schema getKeySchema() {
        return mKeySchema;
      }

      /**
       * Sets the reader schema for the value.
       *
       * @param valueSchema The reader schema for the value.
       * @return This options instance.
       */
      public Options withValueSchema(Schema valueSchema) {
        mValueSchema = valueSchema;
        return this;
      }

      /**
       * Gets the reader schema for the value.
       *
       * @return The reader schema for the value.
       */
      public Schema getValueSchema() {
        return mValueSchema;
      }

      /** Set the data model. */
      public Options withDataModel(GenericData model) {
        this.model = model;
        return this;
      }

      /** Return the data model. */
      public GenericData getDataModel() {
        return model;
      }

    }

    /**
     * Constructs a reader.
     *
     * @param options The options.
     * @throws IOException If there is an error.
     */
    public Reader(Options options) throws IOException {
      mKeySchema = options.getKeySchema();
      this.model = options.getDataModel();

      // Load the whole index file into memory.
      Path indexFilePath = new Path(options.getPath(), INDEX_FILENAME);
      LOG.debug("Loading the index from " + indexFilePath);
      mIndex = loadIndexFile(options.getConfiguration(), indexFilePath, mKeySchema);

      // Open the data file.
      Path dataFilePath = new Path(options.getPath(), DATA_FILENAME);
      LOG.debug("Loading the data file " + dataFilePath);
      Schema recordSchema = AvroKeyValue.getSchema(mKeySchema, options.getValueSchema());
      DatumReader<GenericRecord> datumReader =
        model.createDatumReader(recordSchema);
      mDataFileReader =
        new DataFileReader<GenericRecord>
        (new FsInput(dataFilePath, options.getConfiguration()), datumReader);

    }

    /**
     * Gets the first value associated with a given key, or null if it is not found.
     *
     * <p>This method will move the current position in the file to the record immediately
     * following the requested key.</p>
     *
     * @param key The key to look up.
     * @return The value associated with the key, or null if not found.
     * @throws IOException If there is an error.
     */
    public V get(K key) throws IOException {
      // Look up the entry in the index.
      LOG.debug("Looking up key " + key + " in the index.");
      Map.Entry<K, Long> indexEntry = mIndex.floorEntry(key);
      if (null == indexEntry) {
        LOG.debug("Key " + key + " was not found in the index (it is before the first entry)");
        return null;
      }
      LOG.debug("Key was found in the index, seeking to syncpoint " + indexEntry.getValue());

      // Seek to the data block that would contain the entry.
      mDataFileReader.seek(indexEntry.getValue());

      // Scan from this position of the file until we find it or pass it.
      Iterator<AvroKeyValue<K, V>> iter = iterator();
      while (iter.hasNext()) {
        AvroKeyValue<K, V> record = iter.next();
        int comparison = model.compare(record.getKey(), key, mKeySchema);
        if (0 == comparison) {
          // We've found it!
          LOG.debug("Found record for key " + key);
          return record.getValue();
        }
        if (comparison > 0) {
          // We've passed it.
          LOG.debug("Searched beyond the point where key " + key + " would appear in the file.");
          return null;
        }
      }

      // We've reached the end of the file.
      LOG.debug("Searched to the end of the file but did not find key " + key);
      return null;
    }

    /**
     * Returns an iterator starting at the current position in the file.
     *
     * <p>Use the get() method to move the current position.</p>
     *
     * <p>Note that this iterator is shared with other clients of the file; it does not
     * contain a separate pointer into the file.</p>
     *
     * @return An iterator.
     */
    public Iterator<AvroKeyValue<K, V>> iterator() {
      return new AvroKeyValue.Iterator<K, V>(mDataFileReader.iterator());
    }

    /** {@inheritDoc} */
    @Override
    public void close() throws IOException {
      mDataFileReader.close();
    }

    /**
     * Loads an index file into an in-memory map, from key to file offset in bytes.
     *
     * @param conf The configuration.
     * @param path The path to the index file.
     * @param keySchema The reader schema for the key.
     * @throws IOException If there is an error.
     */
    private <K> NavigableMap<K, Long> loadIndexFile(
        Configuration conf, Path path, Schema keySchema) throws IOException {
      DatumReader<GenericRecord> datumReader = model.createDatumReader(
          AvroKeyValue.getSchema(keySchema, Schema.create(Schema.Type.LONG)));
      DataFileReader<GenericRecord> fileReader = new DataFileReader<GenericRecord>(
          new FsInput(path, conf), datumReader);

      NavigableMap<K, Long> index;
      if (Schema.create(Schema.Type.STRING).equals(keySchema)) {
        // Because Avro STRING types are mapped to the Java CharSequence class that does not
        // mandate the implementation of Comparable, we need to specify a special
        // CharSequence comparator if the key type is a string.  This hack only fixes the
        // problem for primitive string types.  If, for example, you tried to use a record
        // type as the key, any string fields inside of it would not be compared correctly
        // against java.lang.Strings.
        index = new TreeMap<K, Long>(new AvroCharSequenceComparator<K>());
      } else {
        index = new TreeMap<K, Long>();
      }
      try {
        for (GenericRecord genericRecord : fileReader) {
          AvroKeyValue<K, Long> indexRecord = new AvroKeyValue<K, Long>(genericRecord);
          index.put(indexRecord.getKey(), indexRecord.getValue());
        }
      } finally {
        fileReader.close();
      }
      return index;
    }
  }

  /**
   * Writes a SortedKeyValueFile.
   *
   * @param <K> The key type.
   * @param <V> The value type.
   */
  public static class Writer<K, V> implements Closeable {
    /** The key schema. */
    private final Schema mKeySchema;

    /** The value schema. */
    private final Schema mValueSchema;

    /** The schema of the data file records. */
    private final Schema mRecordSchema;

    /** The schema of the index file records. */
    private final Schema mIndexSchema;

    /** The model for the data. */
    private GenericData model;

    /** The writer for the data file. */
    private final DataFileWriter<GenericRecord> mDataFileWriter;

    /** The writer for the index file. */
    private final DataFileWriter<GenericRecord> mIndexFileWriter;

    /** We store an indexed key for every mIndexInterval records written to the data file. */
    private final int mIndexInterval;

    /** The number of records written to the file so far. */
    private long mRecordsWritten;

    /** The most recent key that was appended to the file, or null. */
    private K mPreviousKey;

    /** A class to encapsulate the various options of a SortedKeyValueFile.Writer. */
    public static class Options {
      /** The key schema. */
      private Schema mKeySchema;

      /** The value schema. */
      private Schema mValueSchema;

      /** The configuration. */
      private Configuration mConf;

      /** The path to the output file. */
      private Path mPath;

      /** The number of records between indexed entries. */
      private int mIndexInterval = 128;

      /** The model for the data. */
      private GenericData model = SpecificData.get();

      /** The compression codec for the data. */
      private CodecFactory codec = CodecFactory.nullCodec();

      /**
       * Sets the key schema.
       *
       * @param keySchema The key schema.
       * @return This options instance.
       */
      public Options withKeySchema(Schema keySchema) {
        mKeySchema = keySchema;
        return this;
      }

      /**
       * Gets the key schema.
       *
       * @return The key schema.
       */
      public Schema getKeySchema() {
        return mKeySchema;
      }

      /**
       * Sets the value schema.
       *
       * @param valueSchema The value schema.
       * @return This options instance.
       */
      public Options withValueSchema(Schema valueSchema) {
        mValueSchema = valueSchema;
        return this;
      }

      /**
       * Gets the value schema.
       *
       * @return The value schema.
       */
      public Schema getValueSchema() {
        return mValueSchema;
      }

      /**
       * Sets the configuration.
       *
       * @param conf The configuration.
       * @return This options instance.
       */
      public Options withConfiguration(Configuration conf) {
        mConf = conf;
        return this;
      }

      /**
       * Gets the configuration.
       *
       * @return The configuration.
       */
      public Configuration getConfiguration() {
        return mConf;
      }

      /**
       * Sets the output path.
       *
       * @param path The output path.
       * @return This options instance.
       */
      public Options withPath(Path path) {
        mPath = path;
        return this;
      }

      /**
       * Gets the output path.
       *
       * @return The output path.
       */
      public Path getPath() {
        return mPath;
      }

      /**
       * Sets the index interval.
       *
       * <p>If the index inverval is N, then every N records will be indexed into the
       * index file.</p>
       *
       * @param indexInterval The index interval.
       * @return This options instance.
       */
      public Options withIndexInterval(int indexInterval) {
        mIndexInterval = indexInterval;
        return this;
      }

      /**
       * Gets the index interval.
       *
       * @return The index interval.
       */
      public int getIndexInterval() {
        return mIndexInterval;
      }

      /** Set the data model. */
      public Options withDataModel(GenericData model) {
        this.model = model;
        return this;
      }

      /** Return the data model. */
      public GenericData getDataModel() {
        return model;
      }

      /** Set the compression codec. */
      public Options withCodec(String codec) {
          this.codec = CodecFactory.fromString(codec);
          return this;
      }

      /** Set the compression codec. */
      public Options withCodec(CodecFactory codec) {
          this.codec = codec;
          return this;
      }

      /** Return the compression codec. */
      public CodecFactory getCodec() {
          return this.codec;
      }
    }

    /**
     * Creates a writer for a new file.
     *
     * @param options The options.
     * @throws IOException If there is an error.
     */
    public Writer(Options options) throws IOException {
      this.model = options.getDataModel();

      if (null == options.getConfiguration()) {
        throw new IllegalArgumentException("Configuration may not be null");
      }

      FileSystem fileSystem = options.getPath().getFileSystem(options.getConfiguration());

      // Save the key and value schemas.
      mKeySchema = options.getKeySchema();
      if (null == mKeySchema) {
        throw new IllegalArgumentException("Key schema may not be null");
      }
      mValueSchema = options.getValueSchema();
      if (null == mValueSchema) {
        throw new IllegalArgumentException("Value schema may not be null");
      }

      // Save the index interval.
      mIndexInterval = options.getIndexInterval();

      // Create the directory.
      if (!fileSystem.mkdirs(options.getPath())) {
        throw new IOException(
            "Unable to create directory for SortedKeyValueFile: " + options.getPath());
      }
      LOG.debug("Created directory " + options.getPath());

      // Open a writer for the data file.
      Path dataFilePath = new Path(options.getPath(), DATA_FILENAME);
      LOG.debug("Creating writer for avro data file: " + dataFilePath);
      mRecordSchema = AvroKeyValue.getSchema(mKeySchema, mValueSchema);
      DatumWriter<GenericRecord> datumWriter =
        model.createDatumWriter(mRecordSchema);
      OutputStream dataOutputStream = fileSystem.create(dataFilePath);
      mDataFileWriter = new DataFileWriter<GenericRecord>(datumWriter)
          .setSyncInterval(1 << 20)  // Set the auto-sync interval sufficiently large, since
                                     // we will manually sync every mIndexInterval records.
          .setCodec(options.getCodec())
          .create(mRecordSchema, dataOutputStream);

      // Open a writer for the index file.
      Path indexFilePath = new Path(options.getPath(), INDEX_FILENAME);
      LOG.debug("Creating writer for avro index file: " + indexFilePath);
      mIndexSchema = AvroKeyValue.getSchema(mKeySchema, Schema.create(Schema.Type.LONG));
      DatumWriter<GenericRecord> indexWriter =
        model.createDatumWriter(mIndexSchema);
      OutputStream indexOutputStream = fileSystem.create(indexFilePath);
      mIndexFileWriter = new DataFileWriter<GenericRecord>(indexWriter)
          .create(mIndexSchema, indexOutputStream);
    }

    /**
     * Appends a record to the SortedKeyValueFile.
     *
     * @param key The key.
     * @param value The value.
     * @throws IOException If there is an error.
     */
    public void append(K key, V value) throws IOException {
      // Make sure the keys are inserted in sorted order.
      if (null != mPreviousKey && model.compare(key, mPreviousKey, mKeySchema) < 0) {
        throw new IllegalArgumentException("Records must be inserted in sorted key order."
            + " Attempted to insert key " + key + " after " + mPreviousKey + ".");
      }
      mPreviousKey = model.deepCopy(mKeySchema, key);

      // Construct the data record.
      AvroKeyValue<K, V> dataRecord
          = new AvroKeyValue<K, V>(new GenericData.Record(mRecordSchema));
      dataRecord.setKey(key);
      dataRecord.setValue(value);

      // Index it if necessary.
      if (0 == mRecordsWritten++ % mIndexInterval) {
        // Force a sync to the data file writer, which closes the current data block (if
        // nonempty) and reports the current position in the file.
        long position = mDataFileWriter.sync();

        // Construct the record to put in the index.
        AvroKeyValue<K, Long> indexRecord
            = new AvroKeyValue<K, Long>(new GenericData.Record(mIndexSchema));
        indexRecord.setKey(key);
        indexRecord.setValue(position);
        mIndexFileWriter.append(indexRecord.get());
      }

      // Write it to the data file.
      mDataFileWriter.append(dataRecord.get());
    }

    /** {@inheritDoc} */
    @Override
    public void close() throws IOException {
      mIndexFileWriter.close();
      mDataFileWriter.close();
    }
  }
}