/*
    * 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.giraph.ooc.data;
  
  import java.io.DataInput;
  import java.io.DataOutput;
  import java.io.IOException;
  
  import org.apache.giraph.comm.messages.MessageStore;
  import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
  import org.apache.giraph.factories.MessageValueFactory;
  import org.apache.giraph.ooc.OutOfCoreEngine;
  import org.apache.giraph.ooc.persistence.DataIndex;
  import org.apache.giraph.ooc.persistence.DataIndex.NumericIndexEntry;
  import org.apache.giraph.ooc.persistence.OutOfCoreDataAccessor;
  import org.apache.giraph.utils.ByteArrayOneMessageToManyIds;
  import org.apache.giraph.utils.ByteArrayVertexIdMessages;
  import org.apache.giraph.utils.VertexIdMessages;
  import org.apache.hadoop.io.Writable;
  import org.apache.hadoop.io.WritableComparable;
  import org.apache.log4j.Logger;
  
  
  /**
   * Implementation of a message store used for out-of-core mechanism.
   *
   * @param <I> Vertex id
   * @param <M> Message data
   */
  public class DiskBackedMessageStore<I extends WritableComparable,
      M extends Writable> extends DiskBackedDataStore<VertexIdMessages<I, M>>
      implements MessageStore<I, M> {
    /** Class logger. */
    private static final Logger LOG =
        Logger.getLogger(DiskBackedMessageStore.class);
    /** Configuration */
    private final ImmutableClassesGiraphConfiguration<I, ?, ?> config;
    /** In-memory message store */
    private final MessageStore<I, M> messageStore;
    /** Whether the message store uses message combiner or not */
    private final boolean useMessageCombiner;
    /** Which superstep this message store is used for */
    private final long superstep;
    /** Message value class */
    private final MessageValueFactory<M> messageValueFactory;
  
    /**
     * Type of VertexIdMessage class (container for serialized messages) received
     * for a particular message. If we write the received messages to disk before
     * adding them to message store, we need this type when we want to read the
     * messages back from disk (so that we know how to read the messages from
     * disk).
     */
    private enum SerializedMessageClass {
      /** ByteArrayVertexIdMessages */
      BYTE_ARRAY_VERTEX_ID_MESSAGES,
      /** ByteArrayOneMEssageToManyIds */
      BYTE_ARRAY_ONE_MESSAGE_TO_MANY_IDS
    }
  
    /**
     * Constructor
     *
     * @param config Configuration
     * @param oocEngine Out-of-core engine
     * @param messageStore In-memory message store for which out-of-core message
     *                     store would be wrapper
     * @param useMessageCombiner Whether message combiner is used for this message
     *                           store
     * @param superstep superstep number this messages store is used for
     */
    public DiskBackedMessageStore(ImmutableClassesGiraphConfiguration<I, ?, ?>
                                      config,
                                  OutOfCoreEngine oocEngine,
                                  MessageStore<I, M> messageStore,
                                  boolean useMessageCombiner, long superstep) {
      super(config, oocEngine);
      this.config = config;
      this.messageStore = messageStore;
      this.useMessageCombiner = useMessageCombiner;
      this.superstep = superstep;
      this.messageValueFactory = config.createOutgoingMessageValueFactory();
    }
 
   @Override
   public boolean isPointerListEncoding() {
     return messageStore.isPointerListEncoding();
   }
 
   @Override
   public Iterable<M> getVertexMessages(I vertexId) {
     return messageStore.getVertexMessages(vertexId);
   }
 
   @Override
   public void clearVertexMessages(I vertexId) {
     messageStore.clearVertexMessages(vertexId);
   }
 
   @Override
   public void clearAll() {
     messageStore.clearAll();
   }
 
   @Override
   public boolean hasMessagesForVertex(I vertexId) {
     return messageStore.hasMessagesForVertex(vertexId);
   }
 
   @Override
   public boolean hasMessagesForPartition(int partitionId) {
     return messageStore.hasMessagesForPartition(partitionId);
   }
 
   @Override
   public void addPartitionMessages(
       int partitionId, VertexIdMessages<I, M> messages) {
     if (useMessageCombiner) {
       messageStore.addPartitionMessages(partitionId, messages);
     } else {
       addEntry(partitionId, messages);
     }
   }
 
   @Override
   public void addMessage(I vertexId, M message) throws IOException {
     if (useMessageCombiner) {
       messageStore.addMessage(vertexId, message);
     } else {
       // TODO: implement if LocalBlockRunner needs this message store
       throw new UnsupportedOperationException();
     }
   }
 
   /**
    * Gets the path that should be used specifically for message data.
    *
    * @param basePath path prefix to build the actual path from
    * @param superstep superstep for which message data should be stored
    * @return path to files specific for message data
    */
   private static String getPath(String basePath, long superstep) {
     return basePath + "_messages-S" + superstep;
   }
 
   @Override
   public long loadPartitionData(int partitionId)
       throws IOException {
     if (!useMessageCombiner) {
       return loadPartitionDataProxy(partitionId,
           new DataIndex().addIndex(DataIndex.TypeIndexEntry.MESSAGE)
               .addIndex(NumericIndexEntry.createSuperstepEntry(superstep)));
     } else {
       return 0;
     }
   }
 
   @Override
   public long offloadPartitionData(int partitionId)
       throws IOException {
     if (!useMessageCombiner) {
       return offloadPartitionDataProxy(partitionId,
           new DataIndex().addIndex(DataIndex.TypeIndexEntry.MESSAGE)
               .addIndex(NumericIndexEntry.createSuperstepEntry(superstep)));
     } else {
       return 0;
     }
   }
 
   @Override
   public long offloadBuffers(int partitionId)
       throws IOException {
     if (!useMessageCombiner) {
       return offloadBuffersProxy(partitionId,
           new DataIndex().addIndex(DataIndex.TypeIndexEntry.MESSAGE)
               .addIndex(NumericIndexEntry.createSuperstepEntry(superstep)));
     } else {
       return 0;
     }
   }
 
   @Override
   public void finalizeStore() {
     messageStore.finalizeStore();
   }
 
   @Override
   public Iterable<I> getPartitionDestinationVertices(int partitionId) {
     return messageStore.getPartitionDestinationVertices(partitionId);
   }
 
   @Override
   public void clearPartition(int partitionId) {
     messageStore.clearPartition(partitionId);
   }
 
   @Override
   public void writePartition(DataOutput out, int partitionId)
       throws IOException {
     messageStore.writePartition(out, partitionId);
   }
 
   @Override
   public void readFieldsForPartition(DataInput in, int partitionId)
       throws IOException {
     messageStore.readFieldsForPartition(in, partitionId);
   }
 
   @Override
   protected void writeEntry(VertexIdMessages<I, M> messages, DataOutput out)
       throws IOException {
     SerializedMessageClass messageClass;
     if (messages instanceof ByteArrayVertexIdMessages) {
       messageClass = SerializedMessageClass.BYTE_ARRAY_VERTEX_ID_MESSAGES;
     } else if (messages instanceof ByteArrayOneMessageToManyIds) {
       messageClass = SerializedMessageClass.BYTE_ARRAY_ONE_MESSAGE_TO_MANY_IDS;
     } else {
       throw new IllegalStateException("writeEntry: serialized message " +
           "type is not supported");
     }
     out.writeByte(messageClass.ordinal());
     messages.write(out);
   }
 
   @Override
   protected VertexIdMessages<I, M> readNextEntry(DataInput in)
       throws IOException {
     byte messageType = in.readByte();
     SerializedMessageClass messageClass =
         SerializedMessageClass.values()[messageType];
     VertexIdMessages<I, M> vim;
     switch (messageClass) {
     case BYTE_ARRAY_VERTEX_ID_MESSAGES:
       vim = new ByteArrayVertexIdMessages<>(messageValueFactory);
       vim.setConf(config);
       break;
     case BYTE_ARRAY_ONE_MESSAGE_TO_MANY_IDS:
       vim = new ByteArrayOneMessageToManyIds<>(messageValueFactory);
       vim.setConf(config);
       break;
     default:
       throw new IllegalStateException("readNextEntry: unsupported " +
           "serialized message type!");
     }
     vim.readFields(in);
     return vim;
   }
 
   @Override
   protected long loadInMemoryPartitionData(int partitionId, int ioThreadId,
                                            DataIndex index) throws IOException {
     long numBytes = 0;
     if (hasPartitionDataOnFile.remove(partitionId)) {
       OutOfCoreDataAccessor.DataInputWrapper inputWrapper =
           oocEngine.getDataAccessor().prepareInput(ioThreadId, index.copy());
       messageStore.readFieldsForPartition(inputWrapper.getDataInput(),
           partitionId);
       numBytes = inputWrapper.finalizeInput(true);
     }
     return numBytes;
   }
 
   @Override
   protected long offloadInMemoryPartitionData(
       int partitionId, int ioThreadId, DataIndex index) throws IOException {
     long numBytes = 0;
     if (messageStore.hasMessagesForPartition(partitionId)) {
       OutOfCoreDataAccessor.DataOutputWrapper outputWrapper =
           oocEngine.getDataAccessor().prepareOutput(ioThreadId, index.copy(),
               false);
       messageStore.writePartition(outputWrapper.getDataOutput(), partitionId);
       messageStore.clearPartition(partitionId);
       numBytes = outputWrapper.finalizeOutput();
       hasPartitionDataOnFile.add(partitionId);
     }
     return numBytes;
   }
 
   @Override
   protected int entrySerializedSize(VertexIdMessages<I, M> messages) {
     return messages.getSerializedSize();
   }
 
   @Override
   protected void addEntryToInMemoryPartitionData(int partitionId,
                                                  VertexIdMessages<I, M>
                                                      messages) {
     messageStore.addPartitionMessages(partitionId, messages);
   }
 }