/*
 * 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.tez.runtime.library.common.writers;

import java.io.IOException;
import java.util.Iterator;
import java.util.concurrent.atomic.AtomicBoolean;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.RawLocalFileSystem;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.serializer.Serialization;
import org.apache.hadoop.io.serializer.SerializationFactory;
import org.apache.hadoop.io.serializer.Serializer;
import org.apache.tez.common.counters.TaskCounter;
import org.apache.tez.common.counters.TezCounter;
import org.apache.tez.runtime.api.OutputContext;
import org.apache.tez.runtime.library.api.KeyValuesWriter;
import org.apache.tez.runtime.library.api.Partitioner;
import org.apache.tez.runtime.library.api.TezRuntimeConfiguration;
import org.apache.tez.runtime.library.common.ConfigUtils;
import org.apache.tez.runtime.library.common.TezRuntimeUtils;
import org.apache.tez.runtime.library.sort.CelebornSortBasedPusher;
import org.apache.tez.runtime.library.utils.CodecUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.client.CelebornTezWriter;
import org.apache.celeborn.common.CelebornConf;

public class CelebornUnorderedPartitionedKVWriter extends KeyValuesWriter {
  private static final Logger Logger =
      LoggerFactory.getLogger(CelebornUnorderedPartitionedKVWriter.class);

  protected final OutputContext outputContext;
  protected final Configuration conf;
  protected final RawLocalFileSystem localFs;
  protected final Partitioner partitioner;
  protected final Class keyClass;
  protected final Class valClass;
  protected final Serializer keySerializer;
  protected final Serializer valSerializer;
  protected final SerializationFactory serializationFactory;
  protected final Serialization keySerialization;
  protected final Serialization valSerialization;
  protected final int numOutputs;
  protected final CompressionCodec codec;

  protected final TezCounter outputRecordBytesCounter;
  protected final TezCounter outputRecordsCounter;
  protected final TezCounter outputBytesWithOverheadCounter;

  private long availableMemory;
  private int[] numRecordsPerPartition;
  private long[] sizePerPartition;
  private AtomicBoolean isShutdown = new AtomicBoolean(false);

  final TezRuntimeConfiguration.ReportPartitionStats reportPartitionStats;

  private CelebornSortBasedPusher pusher;

  public CelebornUnorderedPartitionedKVWriter(
      OutputContext outputContext,
      Configuration conf,
      int numOutputs,
      long availableMemoryBytes,
      CelebornTezWriter celebornTezWriter,
      CelebornConf celebornConf) {
    this.outputContext = outputContext;
    this.conf = conf;
    try {
      this.localFs = (RawLocalFileSystem) FileSystem.getLocal(conf).getRaw();
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    this.numOutputs = numOutputs;

    // k/v serialization
    keyClass = ConfigUtils.getIntermediateOutputKeyClass(this.conf);
    valClass = ConfigUtils.getIntermediateOutputValueClass(this.conf);
    serializationFactory = new SerializationFactory(this.conf);
    keySerialization = serializationFactory.getSerialization(keyClass);
    valSerialization = serializationFactory.getSerialization(valClass);
    keySerializer = keySerialization.getSerializer(keyClass);
    valSerializer = valSerialization.getSerializer(valClass);

    outputRecordBytesCounter = outputContext.getCounters().findCounter(TaskCounter.OUTPUT_BYTES);
    outputRecordsCounter = outputContext.getCounters().findCounter(TaskCounter.OUTPUT_RECORDS);
    outputBytesWithOverheadCounter =
        outputContext.getCounters().findCounter(TaskCounter.OUTPUT_BYTES_WITH_OVERHEAD);

    // stats
    reportPartitionStats =
        TezRuntimeConfiguration.ReportPartitionStats.fromString(
            conf.get(
                TezRuntimeConfiguration.TEZ_RUNTIME_REPORT_PARTITION_STATS,
                TezRuntimeConfiguration.TEZ_RUNTIME_REPORT_PARTITION_STATS_DEFAULT));
    sizePerPartition = (reportPartitionStats.isEnabled()) ? new long[numOutputs] : null;
    numRecordsPerPartition = new int[numOutputs];

    // compression
    try {
      this.codec = CodecUtils.getCodec(conf);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }

    Logger.info(
        "Instantiating Partitioner: [{}]",
        conf.get(TezRuntimeConfiguration.TEZ_RUNTIME_PARTITIONER_CLASS));

    try {
      this.partitioner = TezRuntimeUtils.instantiatePartitioner(this.conf);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }

    availableMemory = availableMemoryBytes;
    // assume that there is 64MB memory to writer shuffle data
    if (availableMemory == 0) {
      availableMemory = 64 * 1024 * 1024;
    }
    pusher =
        new CelebornSortBasedPusher(
            keySerializer,
            valSerializer,
            (int) availableMemory,
            (int) (availableMemory * 0.8),
            null,
            outputRecordBytesCounter,
            outputRecordsCounter,
            celebornTezWriter,
            celebornConf,
            false);
  }

  @Override
  public void write(Object key, Iterable<Object> iterable) throws IOException {
    Iterator<Object> it = iterable.iterator();
    while (it.hasNext()) {
      write(key, it.next());
    }
  }

  @Override
  public void write(Object key, Object value) throws IOException {
    if (isShutdown.get()) {
      throw new RuntimeException("Writer already closed");
    }
    pusher.insert(key, value, partitioner.getPartition(key, value, numOutputs));
  }

  public void close() throws IOException {
    pusher.close();
    isShutdown.set(true);
    updateTezCountersAndNotify();
  }

  private void updateTezCountersAndNotify() {
    numRecordsPerPartition = pusher.getRecordsPerPartition();
    if (sizePerPartition != null) {
      sizePerPartition = pusher.getBytesPerPartition();
    }
    outputContext.notifyProgress();
  }

  public int[] getNumRecordsPerPartition() {
    return numRecordsPerPartition;
  }

  public boolean reportDetailedPartitionStats() {
    return reportPartitionStats.isPrecise();
  }

  public long[] getPartitionStats() {
    return sizePerPartition;
  }
}
