/*
    * 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.block_app.framework.api.local;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Iterator;
  import java.util.List;
  import java.util.concurrent.ThreadLocalRandom;
  
  import org.apache.giraph.comm.messages.InMemoryMessageStoreFactory;
  import org.apache.giraph.comm.messages.MessageStore;
  import org.apache.giraph.comm.messages.PartitionSplitInfo;
  import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
  import org.apache.giraph.conf.MessageClasses;
  import org.apache.giraph.factories.MessageValueFactory;
  import org.apache.giraph.utils.WritableUtils;
  import org.apache.hadoop.io.Writable;
  import org.apache.hadoop.io.WritableComparable;
  
  import com.google.common.collect.Iterators;
  
  /**
   * Interface for internal message store, used by LocalBlockRunner
   *
   * @param <I> Vertex id type
   * @param <M> Message type
   */
  @SuppressWarnings("rawtypes")
  interface InternalMessageStore
      <I extends WritableComparable, M extends Writable> {
    Iterator<I> targetVertexIds();
    boolean hasMessage(I id);
    Iterable<M> takeMessages(I id);
    void sendMessage(I id, M message);
    void sendMessageToMultipleEdges(Iterator<I> idIter, M message);
    void finalizeStore();
    Iterable<I> getPartitionDestinationVertices(int partitionId);
  
    /**
     * A wrapper that uses InMemoryMessageStoreFactory to
     * create MessageStore
     *
     * @param <I> Vertex id type
     * @param <M> Message type
     */
    class InternalWrappedMessageStore
    <I extends WritableComparable, M extends Writable>
    implements InternalMessageStore<I, M> {
      private final MessageStore<I, M> messageStore;
      private final PartitionSplitInfo<I> partitionInfo;
  
      private InternalWrappedMessageStore(
        ImmutableClassesGiraphConfiguration<I, ?, ?> conf,
        MessageStore<I, M> messageStore,
        PartitionSplitInfo<I> partitionInfo
      ) {
        this.messageStore = messageStore;
        this.partitionInfo = partitionInfo;
      }
  
      public static <I extends WritableComparable, M extends Writable>
      InternalMessageStore<I, M> create(
        ImmutableClassesGiraphConfiguration<I, ?, ?> conf,
        MessageClasses<I, M> messageClasses,
        PartitionSplitInfo<I> partitionInfo
      ) {
        InMemoryMessageStoreFactory<I, M> factory =
          new InMemoryMessageStoreFactory<>();
        factory.initialize(partitionInfo, conf);
        return new InternalWrappedMessageStore<>(
          conf,
          factory.newStore(messageClasses),
          partitionInfo
        );
      }
  
      @Override
      public void sendMessage(I id, M message) {
        try {
          messageStore.addMessage(id, message);
        } catch (IOException e) {
          throw new RuntimeException(e);
        }
     }
 
     @Override
     public void sendMessageToMultipleEdges(Iterator<I> idIter, M message) {
       while (idIter.hasNext()) {
         sendMessage(idIter.next(), message);
       }
     }
 
     @Override
     public Iterable<M> takeMessages(I id) {
       Iterable<M> result = messageStore.getVertexMessages(id);
       messageStore.clearVertexMessages(id);
       return result;
     }
 
     @Override
     public Iterable<I> getPartitionDestinationVertices(int partitionId) {
       return messageStore.getPartitionDestinationVertices(partitionId);
     }
 
     @Override
     public Iterator<I> targetVertexIds() {
       List<Iterator<I>> iterators = new ArrayList<>();
       for (int partition : partitionInfo.getPartitionIds()) {
         Iterable<I> vertices =
           messageStore.getPartitionDestinationVertices(partition);
         iterators.add(vertices.iterator());
       }
       return Iterators.concat(iterators.iterator());
     }
 
     @Override
     public boolean hasMessage(I id) {
       return messageStore.hasMessagesForVertex(id);
     }
 
     @Override
     public void finalizeStore() {
       messageStore.finalizeStore();
     }
   }
 
   /**
    * Message store that add checks for whether serialization seems to be
    * working fine
    */
   static class InternalChecksMessageStore
       <I extends WritableComparable, M extends Writable>
       implements InternalMessageStore<I, M> {
     private final InternalMessageStore<I, M> messageStore;
     private final ImmutableClassesGiraphConfiguration<I, ?, ?> conf;
     private final MessageValueFactory<M> messageFactory;
 
     public InternalChecksMessageStore(
       InternalMessageStore<I, M> messageStore,
       ImmutableClassesGiraphConfiguration<I, ?, ?> conf,
       MessageValueFactory<M> messageFactory
     ) {
       this.messageStore = messageStore;
       this.conf = conf;
       this.messageFactory = messageFactory;
     }
 
     // Use message copies probabilistically, to catch both not serializing some
     // fields, and storing references from message object itself
     // (which can be reusable).
     private M maybeMessageCopy(M message) {
       M messageCopy = WritableUtils.createCopy(
           message, messageFactory, conf);
       return ThreadLocalRandom.current().nextBoolean() ? messageCopy : message;
     }
 
     private void checkIdCopy(I id) {
       WritableUtils.createCopy(id, conf.getVertexIdFactory(), conf);
     }
 
     @Override
     public void sendMessage(I id, M message) {
       checkIdCopy(id);
       messageStore.sendMessage(id, maybeMessageCopy(message));
     }
 
     @Override
     public void sendMessageToMultipleEdges(
         final Iterator<I> idIter, M message) {
       messageStore.sendMessageToMultipleEdges(
           new Iterator<I>() {
             @Override
             public boolean hasNext() {
               return idIter.hasNext();
             }
 
             @Override
             public I next() {
               I id = idIter.next();
               checkIdCopy(id);
               return id;
             }
 
             @Override
             public void remove() {
               idIter.remove();
             }
           },
           maybeMessageCopy(message));
     }
 
     @Override
     public Iterable<M> takeMessages(I id) {
       checkIdCopy(id);
       return messageStore.takeMessages(id);
     }
 
     @Override
     public boolean hasMessage(I id) {
       return messageStore.hasMessage(id);
     }
 
     @Override
     public Iterator<I> targetVertexIds() {
       return messageStore.targetVertexIds();
     }
 
     @Override
     public Iterable<I> getPartitionDestinationVertices(int partitionId) {
       return messageStore.getPartitionDestinationVertices(partitionId);
     }
 
     @Override
     public void finalizeStore() {
       messageStore.finalizeStore();
     }
   }
 }