/*
    * 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.master.input;
  
  import org.apache.giraph.comm.MasterClient;
  import org.apache.giraph.comm.requests.ReplyWithInputSplitRequest;
  import org.apache.giraph.conf.StrConfOption;
  import org.apache.giraph.io.GiraphInputFormat;
  import org.apache.giraph.io.InputType;
  import org.apache.giraph.worker.WorkerInfo;
  import org.apache.hadoop.mapreduce.Counter;
  import org.apache.hadoop.mapreduce.InputSplit;
  import org.apache.hadoop.mapreduce.Mapper;
  
  import java.io.ByteArrayOutputStream;
  import java.io.DataOutput;
  import java.io.DataOutputStream;
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.EnumMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.atomic.AtomicInteger;
  
  /**
   * Handler for input splits on master
   *
   * Since currently Giraph fails if worker fails while reading input, we
   * didn't complicate this part with retries yet, later it could be added by
   * keeping track of which worker got which split and then if worker dies put
   * these splits back to queues.
   */
  public class MasterInputSplitsHandler {
    /**
     * Store in counters timestamps when we finished reading
     * these fractions of input
     */
    public static final StrConfOption DONE_FRACTIONS_TO_STORE_IN_COUNTERS =
        new StrConfOption("giraph.master.input.doneFractionsToStoreInCounters",
            "0.99,1", "Store in counters timestamps when we finished reading " +
            "these fractions of input");
    /** Map of counter group and names */
    private static Map<String, Set<String>> COUNTER_GROUP_AND_NAMES =
            new ConcurrentHashMap<>();
  
    /** Whether to use locality information */
    private final boolean useLocality;
    /** Master client */
    private MasterClient masterClient;
    /** Master client */
    private List<WorkerInfo> workers;
    /** Map of splits organizers for each split type */
    private Map<InputType, InputSplitsMasterOrganizer> splitsMap =
        new EnumMap<>(InputType.class);
    /** Latches to say when one input splits type is ready to be accessed */
    private Map<InputType, CountDownLatch> latchesMap =
        new EnumMap<>(InputType.class);
    /** Context for accessing counters */
    private final Mapper.Context context;
    /** How many splits per type are there total */
    private final Map<InputType, Integer> numSplitsPerType =
        new EnumMap<>(InputType.class);
    /** How many splits per type have been read so far */
    private final Map<InputType, AtomicInteger> numSplitsReadPerType =
        new EnumMap<>(InputType.class);
    /** Timestamps when various splits were created */
    private final Map<InputType, Long> splitsCreatedTimestamp =
        new EnumMap<>(InputType.class);
    /**
     * Store in counters timestamps when we finished reading
     * these fractions of input
     */
    private final double[] doneFractionsToStoreInCounters;
  
    /**
     * Constructor
     *
     * @param useLocality Whether to use locality information or not
     * @param context Context for accessing counters
    */
   public MasterInputSplitsHandler(boolean useLocality, Mapper.Context context) {
     this.useLocality = useLocality;
     this.context = context;
     for (InputType inputType : InputType.values()) {
       latchesMap.put(inputType, new CountDownLatch(1));
       numSplitsReadPerType.put(inputType, new AtomicInteger(0));
     }
 
     String[] tmp = DONE_FRACTIONS_TO_STORE_IN_COUNTERS.get(
         context.getConfiguration()).split(",");
     doneFractionsToStoreInCounters = new double[tmp.length];
     for (int i = 0; i < tmp.length; i++) {
       doneFractionsToStoreInCounters[i] = Double.parseDouble(tmp[i].trim());
     }
   }
 
   /**
    * Initialize
    *
    * @param masterClient Master client
    * @param workers List of workers
    */
   public void initialize(MasterClient masterClient, List<WorkerInfo> workers) {
     this.masterClient = masterClient;
     this.workers = workers;
   }
 
   /**
    * Add splits
    *
    * @param splitsType Type of splits
    * @param inputSplits Splits
    * @param inputFormat Format
    */
   public void addSplits(InputType splitsType, List<InputSplit> inputSplits,
       GiraphInputFormat inputFormat) {
     splitsCreatedTimestamp.put(splitsType, System.currentTimeMillis());
     List<byte[]> serializedSplits = new ArrayList<>();
     for (InputSplit inputSplit : inputSplits) {
       try {
         ByteArrayOutputStream byteArrayOutputStream =
             new ByteArrayOutputStream();
         DataOutput outputStream =
             new DataOutputStream(byteArrayOutputStream);
         inputFormat.writeInputSplit(inputSplit, outputStream);
         serializedSplits.add(byteArrayOutputStream.toByteArray());
       } catch (IOException e) {
         throw new IllegalStateException("IOException occurred", e);
       }
     }
     InputSplitsMasterOrganizer inputSplitsOrganizer;
     if (splitsType == InputType.MAPPING) {
       inputSplitsOrganizer = new MappingInputSplitsMasterOrganizer(
           serializedSplits, workers);
     } else {
       inputSplitsOrganizer = useLocality ?
           new LocalityAwareInputSplitsMasterOrganizer(serializedSplits,
               inputSplits, workers) :
           new BasicInputSplitsMasterOrganizer(serializedSplits);
     }
     splitsMap.put(splitsType, inputSplitsOrganizer);
     latchesMap.get(splitsType).countDown();
     numSplitsPerType.put(splitsType, serializedSplits.size());
   }
 
   /**
    * Called after we receive a split request from some worker, should send
    * split back to it if available, or send it information that there is no
    * more available
    *
    * @param splitType Type of split requested
    * @param workerTaskId Id of worker who requested split
    * @param isFirstSplit Whether this is the first split a thread is requesting,
    *   or this request indicates that previously requested input split was done
    */
   public void sendSplitTo(InputType splitType, int workerTaskId,
       boolean isFirstSplit) {
     try {
       // Make sure we don't try to retrieve splits before they were added
       latchesMap.get(splitType).await();
     } catch (InterruptedException e) {
       throw new IllegalStateException("Interrupted", e);
     }
     byte[] serializedInputSplit =
         splitsMap.get(splitType).getSerializedSplitFor(workerTaskId);
     masterClient.sendWritableRequest(workerTaskId,
         new ReplyWithInputSplitRequest(splitType,
             serializedInputSplit == null ? new byte[0] : serializedInputSplit));
     if (!isFirstSplit) {
       incrementSplitsRead(splitType);
     }
   }
 
   /**
    * Increment splits read
    *
    * @param splitType Type of split which was read
    */
   private void incrementSplitsRead(InputType splitType) {
     int splitsRead = numSplitsReadPerType.get(splitType).incrementAndGet();
     int splits = numSplitsPerType.get(splitType);
     for (int i = 0; i < doneFractionsToStoreInCounters.length; i++) {
       if (splitsRead == (int) (splits * doneFractionsToStoreInCounters[i])) {
         splitFractionReached(
             splitType, doneFractionsToStoreInCounters[i], context);
       }
     }
   }
 
   /**
    * Call when we reached some fraction of split type done to set the
    * timestamp counter
    *
    * @param inputType Type of input
    * @param fraction Which fraction of input type was done reading
    * @param context Context for accessing counters
    */
   private void splitFractionReached(
       InputType inputType, double fraction, Mapper.Context context) {
     getSplitFractionDoneTimestampCounter(inputType, fraction, context).setValue(
         System.currentTimeMillis() - splitsCreatedTimestamp.get(inputType));
   }
 
   /**
    * Get counter
    *
    * @param inputType Type of input for counter
    * @param fraction Fraction for counter
    * @param context Context to get counter from
    * @return Counter
    */
   public static Counter getSplitFractionDoneTimestampCounter(
       InputType inputType, double fraction, Mapper.Context context) {
     String groupName = inputType.name() + " input";
     String counterName = String.format("%.2f%% done time (ms)", fraction * 100);
     Set<String> counters = COUNTER_GROUP_AND_NAMES.getOrDefault(
             groupName, new HashSet<>());
     counters.add(counterName);
     COUNTER_GROUP_AND_NAMES.put(groupName, counters);
     return context.getCounter(groupName, counterName);
   }
 
   public static Map<String, Set<String>> getCounterGroupAndNames() {
     return COUNTER_GROUP_AND_NAMES;
   }
 }