/*
 * 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;
  }
}
