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

import java.io.IOException;

import com.google.common.annotations.VisibleForTesting;
import org.apache.tez.common.CallableWithNdc;
import org.apache.tez.runtime.library.common.InputAttemptIdentifier;
import org.apache.tez.runtime.library.common.shuffle.FetchResult;
import org.apache.tez.runtime.library.common.shuffle.FetchedInput;
import org.apache.tez.runtime.library.common.shuffle.FetchedInputAllocator;
import org.apache.tez.runtime.library.common.shuffle.FetcherCallback;
import org.apache.tez.runtime.library.common.shuffle.orderedgrouped.CelebornTezBypassWriter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.client.CelebornTezReader;
import org.apache.celeborn.common.exception.CelebornIOException;

public class CelebornFetcher extends CallableWithNdc<FetchResult> {
  private static final Logger LOG = LoggerFactory.getLogger(CelebornFetcher.class);
  private final FetcherCallback fetcherCallback;

  private final FetchedInputAllocator inputManager;

  private long copyBlockCount = 0;

  private volatile boolean stopped = false;

  private final CelebornTezReader celebornTezReader;
  private long readTime = 0;
  private long decompressTime = 0;
  private long serializeTime = 0;
  private long waitTime = 0;
  private long copyTime = 0; // the sum of readTime + decompressTime + serializeTime + waitTime
  private static int uniqueMapId = 0;

  private boolean hasPendingData = false;
  private long startWait;
  private int waitCount = 0;
  private byte[] shuffleData = null;

  CelebornFetcher(
      FetcherCallback fetcherCallback,
      FetchedInputAllocator inputManager,
      CelebornTezReader celebornTezReader) {
    this.fetcherCallback = fetcherCallback;
    this.inputManager = inputManager;
    this.celebornTezReader = celebornTezReader;
  }

  public void fetchAllRssBlocks() throws IOException {
    while (!stopped) {
      try {
        copyFromRssServer();
      } catch (Exception e) {
        LOG.error("Failed to fetchAllRssBlocks.", e);
        throw e;
      }
    }
  }

  @VisibleForTesting
  public void copyFromRssServer() throws IOException {
    long blockStartFetch = 0;
    // fetch a block
    if (!hasPendingData) {
      final long startFetch = System.currentTimeMillis();
      blockStartFetch = System.currentTimeMillis();
      shuffleData = celebornTezReader.getShuffleBlock();
      long fetchDuration = System.currentTimeMillis() - startFetch;
      readTime += fetchDuration;
    }

    if (shuffleData != null) {
      // start to merge
      final long startSerialization = System.currentTimeMillis();
      if (issueMapOutputMerge(blockStartFetch)) {
        long serializationDuration = System.currentTimeMillis() - startSerialization;
        serializeTime += serializationDuration;
        // if reserve successes, reset status for next fetch
        if (hasPendingData) {
          waitTime += System.currentTimeMillis() - startWait;
        }
        hasPendingData = false;
        shuffleData = null;
      } else {
        LOG.info("UncompressedData is null");
        // if reserve fail, return and wait
        waitCount++;
        startWait = System.currentTimeMillis();
        return;
      }

      // update some status
      copyBlockCount++;
      copyTime = readTime + decompressTime + serializeTime + waitTime;
    } else {
      LOG.info("UncompressedData is null");
      // finish reading data, close related reader and check data consistent
      celebornTezReader.close();
      LOG.info(
          "Reduce task partition:"
              + celebornTezReader.getPartitionId()
              + " read block cnt: "
              + copyBlockCount
              + " cost "
              + readTime
              + " ms to fetch and "
              + serializeTime
              + " ms to serialize and "
              + waitTime
              + " ms to wait resource, total copy time: "
              + copyTime);
      LOG.info("Stop fetcher");
      stopFetch();
    }
  }

  private boolean issueMapOutputMerge(long blockStartFetch) throws IOException {
    // Allocate a MapOutput (either in-memory or on-disk) to put uncompressed block
    // In Rss, a MapOutput is sent as multiple blocks, so the reducer needs to
    // treat each "block" as a faked "mapout".
    // To avoid name conflicts, we use getNextUniqueTaskAttemptID instead.
    // It will generate a unique TaskAttemptID(increased_seq++, 0).
    InputAttemptIdentifier uniqueInputAttemptIdentifier = getNextUniqueInputAttemptIdentifier();
    FetchedInput fetchedInput = null;

    try {
      fetchedInput =
          inputManager.allocate(
              shuffleData.length, shuffleData.length, uniqueInputAttemptIdentifier);
    } catch (IOException ioe) {
      // kill this reduce attempt
      throw ioe;
    }

    // Allocated space and then write data to mapOutput
    try {
      CelebornTezBypassWriter.write(fetchedInput, shuffleData);
      // let the merger knows this block is ready for merging
      fetcherCallback.fetchSucceeded(
          null,
          uniqueInputAttemptIdentifier,
          fetchedInput,
          shuffleData.length,
          shuffleData.length,
          System.currentTimeMillis() - blockStartFetch);
    } catch (Throwable t) {
      LOG.error("Failed to write fetchedInput.", t);
      throw new CelebornIOException(
          "Partition: "
              + celebornTezReader.getPartitionId()
              + " cannot write block to "
              + fetchedInput.getClass().getSimpleName()
              + " due to: "
              + t.getClass().getName());
    }
    return true;
  }

  private InputAttemptIdentifier getNextUniqueInputAttemptIdentifier() {
    return new InputAttemptIdentifier(uniqueMapId++, 0);
  }

  private void stopFetch() {
    stopped = true;
  }

  @VisibleForTesting
  public int getRetryCount() {
    return waitCount;
  }

  @Override
  protected FetchResult callInternal() throws Exception {
    celebornTezReader.init();
    fetchAllRssBlocks();
    return null;
  }

  public void shutdown() {}

  public Integer getPartitionId() {
    return celebornTezReader.getPartitionId();
  }
}
