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

import java.util.List;
import java.util.Map;

import scala.Function0;
import scala.Function1;
import scala.Option;
import scala.Product2;
import scala.collection.AbstractIterator;
import scala.collection.Iterator;
import scala.runtime.AbstractFunction0;
import scala.runtime.AbstractFunction1;
import scala.runtime.BoxedUnit;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Lists;
import org.apache.hadoop.conf.Configuration;
import org.apache.spark.Aggregator;
import org.apache.spark.InterruptibleIterator;
import org.apache.spark.ShuffleDependency;
import org.apache.spark.TaskContext;
import org.apache.spark.executor.ShuffleReadMetrics;
import org.apache.spark.serializer.Serializer;
import org.apache.spark.shuffle.RssShuffleHandle;
import org.apache.spark.shuffle.ShuffleReader;
import org.apache.spark.util.CompletionIterator;
import org.apache.spark.util.CompletionIterator$;
import org.apache.spark.util.collection.ExternalSorter;
import org.roaringbitmap.longlong.Roaring64NavigableMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.uniffle.client.api.ShuffleReadClient;
import org.apache.uniffle.client.factory.ShuffleClientFactory;
import org.apache.uniffle.client.util.RssClientConfig;
import org.apache.uniffle.common.ShuffleDataDistributionType;
import org.apache.uniffle.common.ShuffleServerInfo;
import org.apache.uniffle.common.config.RssClientConf;
import org.apache.uniffle.common.config.RssConf;

import static org.apache.uniffle.common.util.Constants.DRIVER_HOST;

public class RssShuffleReader<K, C> implements ShuffleReader<K, C> {
  private static final Logger LOG = LoggerFactory.getLogger(RssShuffleReader.class);
  private final Map<Integer, List<ShuffleServerInfo>> partitionToShuffleServers;

  private String appId;
  private int shuffleId;
  private int startPartition;
  private int endPartition;
  private TaskContext context;
  private ShuffleDependency<K, ?, C> shuffleDependency;
  private int numMaps;
  private Serializer serializer;
  private String taskId;
  private String basePath;
  private int partitionNum;
  private Map<Integer, Roaring64NavigableMap> partitionToExpectBlocks;
  private Roaring64NavigableMap taskIdBitmap;
  private Configuration hadoopConf;
  private int mapStartIndex;
  private int mapEndIndex;
  private ShuffleReadMetrics readMetrics;
  private RssConf rssConf;
  private ShuffleDataDistributionType dataDistributionType;

  public RssShuffleReader(
      int startPartition,
      int endPartition,
      int mapStartIndex,
      int mapEndIndex,
      TaskContext context,
      RssShuffleHandle<K, ?, C> rssShuffleHandle,
      String basePath,
      Configuration hadoopConf,
      int partitionNum,
      Map<Integer, Roaring64NavigableMap> partitionToExpectBlocks,
      Roaring64NavigableMap taskIdBitmap,
      ShuffleReadMetrics readMetrics,
      RssConf rssConf,
      ShuffleDataDistributionType dataDistributionType,
      Map<Integer, List<ShuffleServerInfo>> allPartitionToServers) {
    this.appId = rssShuffleHandle.getAppId();
    this.startPartition = startPartition;
    this.endPartition = endPartition;
    this.mapStartIndex = mapStartIndex;
    this.mapEndIndex = mapEndIndex;
    this.context = context;
    this.numMaps = rssShuffleHandle.getNumMaps();
    this.shuffleDependency = rssShuffleHandle.getDependency();
    this.shuffleId = shuffleDependency.shuffleId();
    this.serializer = rssShuffleHandle.getDependency().serializer();
    this.taskId = "" + context.taskAttemptId() + "_" + context.attemptNumber();
    this.basePath = basePath;
    this.partitionNum = partitionNum;
    this.partitionToExpectBlocks = partitionToExpectBlocks;
    this.taskIdBitmap = taskIdBitmap;
    this.hadoopConf = hadoopConf;
    this.readMetrics = readMetrics;
    this.partitionToShuffleServers = allPartitionToServers;
    this.rssConf = rssConf;
    this.dataDistributionType = dataDistributionType;
  }

  @Override
  public Iterator<Product2<K, C>> read() {
    LOG.info("Shuffle read started:" + getReadInfo());

    Iterator<Product2<K, C>> resultIter;
    MultiPartitionIterator rssShuffleDataIterator = new MultiPartitionIterator<K, C>();

    if (shuffleDependency.keyOrdering().isDefined()) {
      // Create an ExternalSorter to sort the data
      Option<Aggregator<K, Object, C>> aggregator = Option.empty();
      if (shuffleDependency.aggregator().isDefined()) {
        if (shuffleDependency.mapSideCombine()) {
          aggregator =
              Option.apply(
                  (Aggregator<K, Object, C>)
                      new Aggregator<K, C, C>(
                          x -> x,
                          shuffleDependency.aggregator().get().mergeCombiners(),
                          shuffleDependency.aggregator().get().mergeCombiners()));
        } else {
          aggregator =
              Option.apply((Aggregator<K, Object, C>) shuffleDependency.aggregator().get());
        }
      }
      ExternalSorter<K, Object, C> sorter =
          new ExternalSorter<>(
              context, aggregator, Option.empty(), shuffleDependency.keyOrdering(), serializer);
      LOG.info("Inserting aggregated records to sorter");
      long startTime = System.currentTimeMillis();
      sorter.insertAll(rssShuffleDataIterator);
      LOG.info(
          "Inserted aggregated records to sorter: millis:"
              + (System.currentTimeMillis() - startTime));
      context.taskMetrics().incMemoryBytesSpilled(sorter.memoryBytesSpilled());
      context.taskMetrics().incPeakExecutionMemory(sorter.peakMemoryUsedBytes());
      context.taskMetrics().incDiskBytesSpilled(sorter.diskBytesSpilled());
      Function0<BoxedUnit> fn0 =
          new AbstractFunction0<BoxedUnit>() {
            @Override
            public BoxedUnit apply() {
              sorter.stop();
              return BoxedUnit.UNIT;
            }
          };
      Function1<TaskContext, Void> fn1 =
          new AbstractFunction1<TaskContext, Void>() {
            public Void apply(TaskContext context) {
              sorter.stop();
              return null;
            }
          };
      context.addTaskCompletionListener(fn1);
      resultIter = CompletionIterator$.MODULE$.apply(sorter.iterator(), fn0);
    } else if (shuffleDependency.aggregator().isDefined()) {
      Aggregator<K, Object, C> aggregator =
          (Aggregator<K, Object, C>) shuffleDependency.aggregator().get();
      if (shuffleDependency.mapSideCombine()) {
        resultIter = aggregator.combineCombinersByKey(rssShuffleDataIterator, context);
      } else {
        resultIter = aggregator.combineValuesByKey(rssShuffleDataIterator, context);
      }
    } else {
      resultIter = rssShuffleDataIterator;
    }

    if (!(resultIter instanceof InterruptibleIterator)) {
      resultIter = new InterruptibleIterator<>(context, resultIter);
    }
    // resubmit stage and shuffle manager server port are both set
    if (rssConf.getBoolean(RssClientConfig.RSS_RESUBMIT_STAGE, false)
        && rssConf.getInteger(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT, 0) > 0) {
      String driver = rssConf.getString(DRIVER_HOST, "");
      int port = rssConf.get(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT);
      resultIter =
          RssFetchFailedIterator.newBuilder()
              .appId(appId)
              .shuffleId(shuffleId)
              .partitionId(startPartition)
              .stageAttemptId(context.stageAttemptNumber())
              .reportServerHost(driver)
              .port(port)
              .build(resultIter);
    }
    return resultIter;
  }

  private String getReadInfo() {
    return "appId="
        + appId
        + ", shuffleId="
        + shuffleId
        + ",taskId="
        + taskId
        + ", partitions: ["
        + startPartition
        + ", "
        + endPartition
        + ")"
        + ", maps: ["
        + mapStartIndex
        + ", "
        + mapEndIndex
        + ")";
  }

  @VisibleForTesting
  public Configuration getHadoopConf() {
    return hadoopConf;
  }

  class MultiPartitionIterator<K, C> extends AbstractIterator<Product2<K, C>> {
    java.util.Iterator<CompletionIterator<Product2<K, C>, RssShuffleDataIterator<K, C>>> iterator;
    CompletionIterator<Product2<K, C>, RssShuffleDataIterator<K, C>> dataIterator;

    MultiPartitionIterator() {
      List<CompletionIterator<Product2<K, C>, RssShuffleDataIterator<K, C>>> iterators =
          Lists.newArrayList();
      if (numMaps <= 0) {
        return;
      }
      for (int partition = startPartition; partition < endPartition; partition++) {
        if (partitionToExpectBlocks.get(partition).isEmpty()) {
          LOG.info("{} partition is empty partition", partition);
          continue;
        }
        List<ShuffleServerInfo> shuffleServerInfoList = partitionToShuffleServers.get(partition);
        // This mechanism of expectedTaskIdsBitmap filter is to filter out the most of data.
        // especially for AQE skew optimization
        boolean expectedTaskIdsBitmapFilterEnable =
            !(mapStartIndex == 0 && mapEndIndex == Integer.MAX_VALUE)
                || shuffleServerInfoList.size() > 1;
        int retryMax =
            rssConf.getInteger(
                RssClientConfig.RSS_CLIENT_RETRY_MAX,
                RssClientConfig.RSS_CLIENT_RETRY_MAX_DEFAULT_VALUE);
        long retryIntervalMax =
            rssConf.getLong(
                RssClientConfig.RSS_CLIENT_RETRY_INTERVAL_MAX,
                RssClientConfig.RSS_CLIENT_RETRY_INTERVAL_MAX_DEFAULT_VALUE);
        ShuffleReadClient shuffleReadClient =
            ShuffleClientFactory.getInstance()
                .createShuffleReadClient(
                    ShuffleClientFactory.newReadBuilder()
                        .appId(appId)
                        .shuffleId(shuffleId)
                        .partitionId(partition)
                        .basePath(basePath)
                        .partitionNumPerRange(1)
                        .partitionNum(partitionNum)
                        .blockIdBitmap(partitionToExpectBlocks.get(partition))
                        .taskIdBitmap(taskIdBitmap)
                        .shuffleServerInfoList(shuffleServerInfoList)
                        .hadoopConf(hadoopConf)
                        .shuffleDataDistributionType(dataDistributionType)
                        .expectedTaskIdsBitmapFilterEnable(expectedTaskIdsBitmapFilterEnable)
                        .retryMax(retryMax)
                        .retryIntervalMax(retryIntervalMax)
                        .rssConf(rssConf));
        RssShuffleDataIterator<K, C> iterator =
            new RssShuffleDataIterator<>(
                shuffleDependency.serializer(), shuffleReadClient, readMetrics, rssConf);
        CompletionIterator<Product2<K, C>, RssShuffleDataIterator<K, C>> completionIterator =
            CompletionIterator$.MODULE$.apply(
                iterator,
                () -> {
                  context.taskMetrics().mergeShuffleReadMetrics();
                  return iterator.cleanup();
                });
        iterators.add(completionIterator);
      }
      iterator = iterators.iterator();
      if (iterator.hasNext()) {
        dataIterator = iterator.next();
        iterator.remove();
      }
      context.addTaskCompletionListener(
          (taskContext) -> {
            if (dataIterator != null) {
              dataIterator.completion();
            }
            iterator.forEachRemaining(CompletionIterator::completion);
          });
    }

    @Override
    public boolean hasNext() {
      if (dataIterator == null) {
        return false;
      }
      while (!dataIterator.hasNext()) {
        if (!iterator.hasNext()) {
          return false;
        }
        dataIterator = iterator.next();
        iterator.remove();
      }
      return dataIterator.hasNext();
    }

    @Override
    public Product2<K, C> next() {
      Product2<K, C> result = dataIterator.next();
      return result;
    }
  }
}