/*
 * 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.hadoop.mapreduce.task.reduce;

import java.io.IOException;
import java.lang.reflect.Field;

import org.apache.hadoop.io.compress.CodecPool;
import org.apache.hadoop.io.compress.Decompressor;
import org.apache.hadoop.mapred.Counters;
import org.apache.hadoop.mapred.Reporter;
import org.apache.hadoop.mapred.TaskStatus;
import org.apache.hadoop.mapreduce.TaskAttemptID;
import org.apache.hadoop.mapreduce.TaskID;
import org.apache.hadoop.mapreduce.TaskType;
import org.apache.hadoop.util.Progress;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.client.read.CelebornInputStream;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.unsafe.Platform;

public class CelebornShuffleFetcher<K, V> {
  private static final Logger logger = LoggerFactory.getLogger(CelebornShuffleFetcher.class);
  private final TaskAttemptID reduceId;
  private final Reporter reporter;
  private final TaskStatus status;
  private final MergeManager<K, V> merger;
  private final Progress progress;
  private final ShuffleClientMetrics metrics;
  private final CelebornInputStream celebornInputStream;
  private volatile boolean stopped = false;
  private int uniqueMapId = 0;
  private final Counters.Counter ioErrs;
  private boolean hasPendingData = false;
  private long inputShuffleSize;
  private byte[] shuffleData;

  public CelebornShuffleFetcher(
      TaskAttemptID reduceId,
      TaskStatus status,
      MergeManager<K, V> merger,
      Progress progress,
      Reporter reporter,
      ShuffleClientMetrics metrics,
      CelebornInputStream input) {
    this.reduceId = reduceId;
    this.reporter = reporter;
    this.status = status;
    this.merger = merger;
    this.progress = progress;
    this.metrics = metrics;
    this.celebornInputStream = input;

    ioErrs = reporter.getCounter("Shuffle Errors", "IO_ERROR");
  }

  // fetch all push data and merge
  public void fetchAndMerge() {
    while (!stopped) {
      try {
        // If merge is on, block
        merger.waitForResource();
        // Do shuffle
        if (metrics != null) {
          metrics.threadBusy();
        }
        // read blocks
        fetchToLocalAndMerge();
      } catch (Exception e) {
        logger.error("Celeborn shuffle fetcher fetch data failed.", e);
      } finally {
        if (metrics != null) {
          metrics.threadFree();
        }
      }
    }
  }

  private byte[] getShuffleBlock() throws IOException {
    // get len
    byte[] header = new byte[4];
    int count = celebornInputStream.read(header);
    if (count == -1) {
      stopped = true;
      return null;
    }
    while (count != header.length) {
      count += celebornInputStream.read(header, count, 4 - count);
    }

    // get data
    int blockLen = Platform.getInt(header, Platform.BYTE_ARRAY_OFFSET);
    inputShuffleSize += blockLen;
    byte[] shuffleData = new byte[blockLen];
    count = celebornInputStream.read(shuffleData);
    while (count != shuffleData.length) {
      count += celebornInputStream.read(shuffleData, count, blockLen - count);
      if (count == -1) {
        // read shuffle is done.
        stopped = true;
        throw new CelebornIOException("Read mr shuffle failed.");
      }
    }
    return shuffleData;
  }

  private void fetchToLocalAndMerge() throws IOException {
    if (!hasPendingData) {
      shuffleData = getShuffleBlock();
    }

    if (shuffleData != null) {
      // start to merge
      if (wrapMapOutput(shuffleData)) {
        hasPendingData = false;
      } else {
        return;
      }

      updateStatus();
      reporter.progress();
    } else {
      celebornInputStream.close();
      if (metrics != null) {
        metrics.inputBytes(inputShuffleSize);
      }
      logger.info("reduce task {} read {} bytes", reduceId, inputShuffleSize);
      stopped = true;
    }
  }

  private boolean wrapMapOutput(byte[] shuffleData) throws IOException {
    // treat push data as mapoutput
    TaskAttemptID mapId =
        new TaskAttemptID(new TaskID(reduceId.getJobID(), TaskType.MAP, uniqueMapId++), 0);
    MapOutput<K, V> mapOutput = null;
    try {
      mapOutput = merger.reserve(mapId, shuffleData.length, 0);
    } catch (IOException ioe) {
      ioErrs.increment(1);
      throw ioe;
    }
    if (mapOutput == null) {
      logger.info(
          "Celeborn fetcher returned status wait because reserve buffer for shuffle get null");
      hasPendingData = true;
      return false;
    }

    // write data to mapOutput
    try {
      writeShuffle(mapOutput, shuffleData);
      // let the merger knows this block is ready for merging
      mapOutput.commit();
    } catch (Throwable t) {
      ioErrs.increment(1);
      mapOutput.abort();
      throw new CelebornIOException(
          "Reduce: {} "
              + reduceId
              + " fetch failed to {} "
              + mapOutput.getClass().getSimpleName()
              + " due to: {} "
              + t.getClass().getName());
    }
    return true;
  }

  private Decompressor getDecompressor(InMemoryMapOutput inMemoryMapOutput)
      throws CelebornIOException {
    try {
      Class clazz = Class.forName(InMemoryMapOutput.class.getName());
      Field deCompressorField = clazz.getDeclaredField("decompressor");
      deCompressorField.setAccessible(true);
      return (Decompressor) deCompressorField.get(inMemoryMapOutput);
    } catch (Exception e) {
      throw new CelebornIOException("Get Decompressor fail " + e.getMessage());
    }
  }

  private void writeShuffle(MapOutput mapOutput, byte[] shuffle) throws CelebornIOException {
    if (mapOutput instanceof InMemoryMapOutput) {
      InMemoryMapOutput inMemoryMapOutput = (InMemoryMapOutput) mapOutput;
      CodecPool.returnDecompressor(getDecompressor(inMemoryMapOutput));
      byte[] memory = inMemoryMapOutput.getMemory();
      System.arraycopy(shuffle, 0, memory, 0, shuffle.length);
    } else if (mapOutput instanceof OnDiskMapOutput) {
      throw new IllegalStateException(
          "Celeborn map reduce client do not support OnDiskMapOutput. Try to increase mapreduce.reduce.shuffle.memory.limit.percent");
    } else {
      throw new IllegalStateException(
          "Merger reserve unknown type of MapOutput: " + mapOutput.getClass().getCanonicalName());
    }
  }

  private void updateStatus() {
    progress.set(
        (float) celebornInputStream.partitionsRead() / celebornInputStream.totalPartitionsToRead());
    String statusString =
        celebornInputStream.partitionsRead()
            + " / "
            + celebornInputStream.totalPartitionsToRead()
            + " copied.";
    status.setStateString(statusString);

    progress.setStatus(
        "copy("
            + celebornInputStream.partitionsRead()
            + " of "
            + celebornInputStream.totalPartitionsToRead());
  }
}
