/*
 * 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.output;

import java.util.HashMap;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentLinkedQueue;

import org.apache.giraph.block_app.framework.api.BlockOutputApi;
import org.apache.giraph.conf.GiraphConstants;
import org.apache.giraph.utils.CallableFactory;
import org.apache.giraph.utils.ProgressableUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.util.Progressable;

/**
 * Handler for blocks output - keeps track of outputs and writers created
 */
@SuppressWarnings("unchecked")
public class BlockOutputHandle implements BlockOutputApi {
  private transient Configuration conf;
  private transient Progressable progressable;
  private final Map<String, BlockOutputDesc> outputDescMap;
  private final Map<String, Queue<BlockOutputWriter>> freeWriters =
      new HashMap<>();
  private final Map<String, Queue<BlockOutputWriter>> occupiedWriters =
      new HashMap<>();

  public BlockOutputHandle() {
    outputDescMap = null;
  }

  public BlockOutputHandle(String jobIdentifier, Configuration conf,
      Progressable hadoopProgressable) {
    outputDescMap = BlockOutputFormat.createInitAndCheckOutputDescsMap(
        conf, jobIdentifier);
    for (Map.Entry<String, BlockOutputDesc> entry : outputDescMap.entrySet()) {
      entry.getValue().preWriting();
      freeWriters.put(entry.getKey(),
          new ConcurrentLinkedQueue<BlockOutputWriter>());
      occupiedWriters.put(entry.getKey(),
          new ConcurrentLinkedQueue<BlockOutputWriter>());
    }
    initialize(conf, hadoopProgressable);
  }

  public void initialize(Configuration conf, Progressable progressable) {
    this.conf = conf;
    this.progressable = progressable;
  }


  @Override
  public <OW extends BlockOutputWriter, OD extends BlockOutputDesc<OW>>
  OD getOutputDesc(String confOption) {
    if (outputDescMap == null) {
      throw new IllegalArgumentException(
          "Output cannot be used with checkpointing");
    }
    return (OD) outputDescMap.get(confOption);
  }

  @Override
  public <OW extends BlockOutputWriter> OW getWriter(String confOption) {
    if (outputDescMap == null) {
      throw new IllegalArgumentException(
          "Output cannot be used with checkpointing");
    }
    OW outputWriter = (OW) freeWriters.get(confOption).poll();
    if (outputWriter == null) {
      outputWriter = (OW) outputDescMap.get(confOption).createOutputWriter(
          conf, progressable);
    }
    occupiedWriters.get(confOption).add(outputWriter);
    return outputWriter;
  }

  public void returnAllWriters() {
    for (Map.Entry<String, Queue<BlockOutputWriter>> entry :
        occupiedWriters.entrySet()) {
      freeWriters.get(entry.getKey()).addAll(entry.getValue());
      entry.getValue().clear();
    }
  }

  public void closeAllWriters() {
    final Queue<BlockOutputWriter> allWriters = new ConcurrentLinkedQueue<>();
    for (Queue<BlockOutputWriter> blockOutputWriters : freeWriters.values()) {
      allWriters.addAll(blockOutputWriters);
    }
    if (allWriters.isEmpty()) {
      return;
    }
    // Closing writers can take time - use multiple threads and call progress
    CallableFactory<Void> callableFactory = new CallableFactory<Void>() {
      @Override
      public Callable<Void> newCallable(int callableId) {
        return new Callable<Void>() {
          @Override
          public Void call() throws Exception {
            BlockOutputWriter writer = allWriters.poll();
            while (writer != null) {
              writer.close();
              writer = allWriters.poll();
            }
            return null;
          }
        };
      }
    };
    ProgressableUtils.getResultsWithNCallables(callableFactory,
        Math.min(GiraphConstants.NUM_OUTPUT_THREADS.get(conf),
            allWriters.size()), "close-writers-%d", progressable);
    // Close all output formats
    for (BlockOutputDesc outputDesc : outputDescMap.values()) {
      outputDesc.postWriting();
    }
  }
}