/*
 * 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.celeborn.plugin.flink;

import static org.apache.celeborn.plugin.flink.utils.Utils.checkState;
import static org.apache.flink.runtime.metrics.groups.ShuffleIOMetricGroup.createShuffleIOMetricGroup;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import org.apache.commons.lang3.tuple.Pair;
import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
import org.apache.flink.runtime.event.AbstractEvent;
import org.apache.flink.runtime.io.AvailabilityProvider;
import org.apache.flink.runtime.io.network.api.EndOfData;
import org.apache.flink.runtime.io.network.api.EndOfPartitionEvent;
import org.apache.flink.runtime.io.network.api.serialization.EventSerializer;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.buffer.BufferDecompressor;
import org.apache.flink.runtime.io.network.buffer.BufferPool;
import org.apache.flink.runtime.io.network.partition.PartitionNotFoundException;
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.io.network.partition.consumer.BufferOrEvent;
import org.apache.flink.runtime.io.network.partition.consumer.PartitionConnectionException;
import org.apache.flink.runtime.metrics.groups.ShuffleIOMetricGroup;
import org.apache.flink.runtime.shuffle.ShuffleDescriptor;
import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.function.SupplierWithException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.DriverChangedException;
import org.apache.celeborn.common.exception.PartitionUnRetryAbleException;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.plugin.flink.buffer.BufferPacker;
import org.apache.celeborn.plugin.flink.buffer.TransferBufferPool;
import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
import org.apache.celeborn.plugin.flink.utils.BufferUtils;

public class RemoteShuffleInputGateDelegation {
  private static final Logger LOG = LoggerFactory.getLogger(RemoteShuffleInputGateDelegation.class);
  /** Lock to protect {@link #receivedBuffers} and {@link #cause} and {@link #closed}. */
  private final Object lock = new Object();

  /** Name of the corresponding computing task. */
  private String taskName;

  /** Index of the gate of the corresponding computing task. */
  private int gateIndex;

  /** Deployment descriptor for a single input gate instance. */
  private InputGateDeploymentDescriptor gateDescriptor;

  /** Buffer pool provider. */
  private SupplierWithException<BufferPool, IOException> bufferPoolFactory;

  /** Flink buffer pools to allocate network memory. */
  private BufferPool bufferPool;

  /** Buffer pool used by the transfer layer. */
  private TransferBufferPool transferBufferPool = new TransferBufferPool(Collections.emptySet());

  private List<RemoteBufferStreamReader> bufferReaders = new ArrayList<>();

  private List<InputChannelInfo> channelsInfo;
  /** Map from channel index to reader index. */
  private int[] channelIndexToReaderIndex;

  /** Map from read index to channel index. */
  private int[] readerIndexToChannelIndex;

  /** The number of subpartitions that has not consumed per channel. */
  private int[] numSubPartitionsNotConsumed;

  /** The overall number of subpartitions that has not been consumed. */
  private long numUnconsumedSubpartitions;

  /** Received buffers from remote shuffle worker. It's consumed by upper computing task. */
  private Queue<Pair<Buffer, InputChannelInfo>> receivedBuffers = new LinkedList<>();

  private Map<InputChannelInfo, ResultPartitionID> channelToResultPartitionId;

  /** {@link Throwable} when reading failure. */
  private Throwable cause;

  /** Whether this remote input gate has been closed or not. */
  private boolean closed;

  /** Whether we have opened all initial channels or not. */
  private boolean initialChannelsOpened;

  /** Number of pending {@link EndOfData} events to be received. */
  private long pendingEndOfDataEvents;

  /** Max concurrent reader count */
  private int numConcurrentReading;

  /** Keep compatibility with streaming mode. */
  private boolean shouldDrainOnEndOfData = true;

  /** Data decompressor. */
  private BufferDecompressor bufferDecompressor;

  private FlinkShuffleClientImpl shuffleClient;

  private int numOpenedReaders = 0;
  private AvailabilityProvider.AvailabilityHelper availabilityHelper;
  private int startSubIndex;
  private int endSubIndex;
  private boolean partitionConnectionExceptionEnabled;

  private final ShuffleIOMetricGroup shuffleIOMetricGroup;

  public RemoteShuffleInputGateDelegation(
      CelebornConf celebornConf,
      ShuffleIOOwnerContext ownerContext,
      int gateIndex,
      InputGateDeploymentDescriptor gateDescriptor,
      SupplierWithException<BufferPool, IOException> bufferPoolFactory,
      BufferDecompressor bufferDecompressor,
      int numConcurrentReading,
      AvailabilityProvider.AvailabilityHelper availabilityHelper,
      int startSubIndex,
      int endSubIndex,
      Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
    this.taskName = ownerContext.getOwnerName();
    this.gateIndex = gateIndex;
    this.gateDescriptor = gateDescriptor;
    this.bufferPoolFactory = bufferPoolFactory;

    int numChannels = gateDescriptor.getShuffleDescriptors().length;
    this.channelIndexToReaderIndex = new int[numChannels];
    this.readerIndexToChannelIndex = new int[numChannels];
    this.numSubPartitionsNotConsumed = new int[numChannels];
    this.bufferDecompressor = bufferDecompressor;

    RemoteShuffleDescriptor remoteShuffleDescriptor =
        (RemoteShuffleDescriptor) gateDescriptor.getShuffleDescriptors()[0];
    RemoteShuffleResource shuffleResource = remoteShuffleDescriptor.getShuffleResource();
    int shuffleId =
        remoteShuffleDescriptor
            .getShuffleResource()
            .getMapPartitionShuffleDescriptor()
            .getShuffleId();
    this.shuffleIOMetricGroup =
        shuffleIOMetricGroups.computeIfAbsent(
            shuffleId, k -> createShuffleIOMetricGroup(ownerContext, shuffleId, celebornConf));

    try {
      String appUniqueId =
          ((RemoteShuffleDescriptor) gateDescriptor.getShuffleDescriptors()[0]).getCelebornAppId();
      this.shuffleClient =
          FlinkShuffleClientImpl.get(
              appUniqueId,
              shuffleResource.getLifecycleManagerHost(),
              shuffleResource.getLifecycleManagerPort(),
              shuffleResource.getLifecycleManagerTimestamp(),
              celebornConf,
              new UserIdentifier("default", "default"));
    } catch (DriverChangedException e) {
      throw new RuntimeException(e.getMessage());
    }

    this.startSubIndex = startSubIndex;
    this.endSubIndex = endSubIndex;

    initShuffleReadClients();
    channelToResultPartitionId = new HashMap<>();
    channelsInfo = createChannelInfos();
    this.numConcurrentReading = numConcurrentReading;
    this.availabilityHelper = availabilityHelper;
    this.partitionConnectionExceptionEnabled = celebornConf.partitionConnectionExceptionEnabled();
    LOG.debug("Initial input gate with numConcurrentReading {}", this.numConcurrentReading);
  }

  private void initShuffleReadClients() {
    int numSubpartitionsPerChannel = endSubIndex - startSubIndex + 1;
    long numUnconsumedSubpartitions = 0;

    // left element is index
    List<Pair<Integer, ShuffleDescriptor>> descriptors =
        IntStream.range(0, gateDescriptor.getShuffleDescriptors().length)
            .mapToObj(i -> Pair.of(i, gateDescriptor.getShuffleDescriptors()[i]))
            .collect(Collectors.toList());

    int readerIndex = 0;
    for (Pair<Integer, ShuffleDescriptor> descriptor : descriptors) {
      RemoteShuffleDescriptor remoteDescriptor = (RemoteShuffleDescriptor) descriptor.getRight();
      ShuffleResourceDescriptor shuffleDescriptor =
          remoteDescriptor.getShuffleResource().getMapPartitionShuffleDescriptor();

      LOG.debug("create shuffle reader for descriptor {}", shuffleDescriptor);

      RemoteBufferStreamReader reader =
          new RemoteBufferStreamReader(
              shuffleClient,
              shuffleDescriptor,
              startSubIndex,
              endSubIndex,
              transferBufferPool,
              getDataListener(descriptor.getLeft(), shuffleIOMetricGroup),
              getFailureListener(remoteDescriptor.getResultPartitionID()));

      bufferReaders.add(reader);
      numSubPartitionsNotConsumed[descriptor.getLeft()] = numSubpartitionsPerChannel;
      numUnconsumedSubpartitions += numSubpartitionsPerChannel;
      channelIndexToReaderIndex[descriptor.getLeft()] = readerIndex;
      readerIndexToChannelIndex[readerIndex] = descriptor.getLeft();
      ++readerIndex;
    }

    this.numUnconsumedSubpartitions = numUnconsumedSubpartitions;
    this.pendingEndOfDataEvents = numUnconsumedSubpartitions;
  }

  private List<InputChannelInfo> createChannelInfos() {
    return IntStream.range(0, gateDescriptor.getShuffleDescriptors().length)
        .mapToObj(
            i -> {
              InputChannelInfo channelInfo = new InputChannelInfo(gateIndex, i);
              channelToResultPartitionId.put(
                  channelInfo, gateDescriptor.getShuffleDescriptors()[i].getResultPartitionID());
              return channelInfo;
            })
        .collect(Collectors.toList());
  }

  private Consumer<ByteBuf> getDataListener(
      int channelIdx, ShuffleIOMetricGroup shuffleIOMetricGroup) {
    return byteBuf -> {
      Queue<Buffer> unpackedBuffers = null;
      try {
        unpackedBuffers = BufferPacker.unpack(byteBuf);
        while (!unpackedBuffers.isEmpty()) {
          onBuffer(unpackedBuffers.poll(), channelIdx, shuffleIOMetricGroup);
        }
      } catch (Throwable throwable) {
        synchronized (lock) {
          cause = cause == null ? throwable : cause;
          availabilityHelper.getUnavailableToResetAvailable().complete(null);
        }

        if (unpackedBuffers != null) {
          unpackedBuffers.forEach(Buffer::recycleBuffer);
        }
        LOG.error("Failed to process the received buffer.", throwable);
      }
    };
  }

  private Consumer<Throwable> getFailureListener(ResultPartitionID rpID) {
    return throwable -> {
      synchronized (lock) {
        if (cause != null) {
          return;
        }
        Class<?> clazz = PartitionUnRetryAbleException.class;
        String message = throwable.getMessage();
        if (null != message && message.contains(clazz.getName())) {
          cause = new PartitionNotFoundException(rpID);
        } else if (partitionConnectionExceptionEnabled
            && null != message
            && message.contains("Failed to connect to")) {
          cause = new PartitionConnectionException(rpID, throwable);
        } else {
          cause = throwable;
        }
        availabilityHelper.getUnavailableToResetAvailable().complete(null);
      }
    };
  }

  private void onBuffer(Buffer buffer, int channelIdx, ShuffleIOMetricGroup shuffleIOMetricGroup) {
    synchronized (lock) {
      if (closed || cause != null) {
        buffer.recycleBuffer();
        throw new IllegalStateException("Input gate already closed or failed.");
      }

      boolean needRecycle = true;
      try {
        boolean wasEmpty = receivedBuffers.isEmpty();
        InputChannelInfo channelInfo = channelsInfo.get(channelIdx);
        checkState(channelInfo.getInputChannelIdx() == channelIdx, "Illegal channel index.");
        LOG.debug("ReceivedBuffers is adding buffer {} on {}", buffer, channelInfo);
        receivedBuffers.add(Pair.of(buffer, channelInfo));
        shuffleIOMetricGroup.getNumBytesIn().inc(buffer.getSize());
        needRecycle = false;
        if (wasEmpty) {
          availabilityHelper.getUnavailableToResetAvailable().complete(null);
        }
      } catch (Throwable throwable) {
        if (needRecycle) {
          buffer.recycleBuffer();
        }
        throw throwable;
      }
    }
  }

  public void setup() throws IOException {
    long startTime = System.nanoTime();

    bufferPool = bufferPoolFactory.get();
    BufferUtils.reserveNumRequiredBuffers(bufferPool, 16);
    tryRequestBuffers();
    availabilityHelper.getUnavailableToResetAvailable().complete(null);
    // Complete availability future though handshake not fired yet, thus to allow fetcher to
    // 'pollNext' and fire handshake to remote. This mechanism is to avoid bookkeeping remote
    // reading resource before task start processing data from input gate.

    LOG.info(
        "Set up read gate by {} ms.", TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime));
  }

  private void tryRequestBuffers() {
    checkState(bufferPool != null, "Not initialized yet.");

    Buffer buffer;
    List<ByteBuf> buffers = new ArrayList<>();
    while ((buffer = bufferPool.requestBuffer()) != null) {
      buffers.add(buffer.asByteBuf());
    }

    if (!buffers.isEmpty()) {
      transferBufferPool.addBuffers(buffers);
    }
  }

  private Buffer decompressBufferIfNeeded(Buffer buffer, InputChannelInfo info) throws IOException {
    if (buffer.isCompressed()) {
      try {
        checkState(bufferDecompressor != null, "Buffer decompressor not set.");
        return bufferDecompressor.decompressToIntermediateBuffer(buffer);
      } catch (Throwable t) {
        LOG.error("Decompress buffer error: inputChannelInfo: {}", info, t);
        throw new PartitionNotFoundException(channelToResultPartitionId.get(info));
      } finally {
        buffer.recycleBuffer();
      }
    }
    return buffer;
  }

  public List<InputChannelInfo> getChannelsInfo() {
    return channelsInfo;
  }

  /** Each one corresponds to a reading channel. */
  public List<RemoteBufferStreamReader> getBufferReaders() {
    return bufferReaders;
  }

  /** Try to open more readers to {@link #numConcurrentReading}. */
  private void tryOpenSomeChannels() throws IOException {
    if (bufferReaders.size() == numOpenedReaders) {
      // all readers are already opened
      return;
    }

    List<RemoteBufferStreamReader> readersToOpen = new ArrayList<>();

    synchronized (lock) {
      if (closed) {
        throw new IOException("Input gate already closed.");
      }

      LOG.debug("Try open some partition readers.");
      int numOnGoing = 0;
      for (int i = 0; i < bufferReaders.size(); i++) {
        RemoteBufferStreamReader bufferStreamReader = bufferReaders.get(i);
        LOG.debug(
            "Trying reader: {}, isOpened={}, numSubPartitionsHasNotConsumed={}.",
            bufferStreamReader,
            bufferStreamReader.isOpened(),
            numSubPartitionsNotConsumed[readerIndexToChannelIndex[i]]);
        if (numOnGoing >= numConcurrentReading) {
          break;
        }

        if (bufferStreamReader.isOpened()
            && numSubPartitionsNotConsumed[readerIndexToChannelIndex[i]] > 0) {
          numOnGoing++;
          continue;
        }

        if (!bufferStreamReader.isOpened()) {
          readersToOpen.add(bufferStreamReader);
          numOnGoing++;
        }
      }
    }

    for (RemoteBufferStreamReader reader : readersToOpen) {
      reader.open(0);
      numOpenedReaders++;
    }
  }

  private Pair<Buffer, InputChannelInfo> getReceived() throws IOException {
    synchronized (lock) {
      healthCheck();
      if (!receivedBuffers.isEmpty()) {
        return receivedBuffers.poll();
      } else {
        if (!allReadersEOF()) {
          availabilityHelper.resetUnavailable();
        }
        return null;
      }
    }
  }

  private void healthCheck() throws IOException {
    if (closed) {
      throw new IOException("Input gate already closed.");
    }
    if (cause != null) {
      if (cause instanceof IOException) {
        throw (IOException) cause;
      } else {
        throw new IOException(cause);
      }
    }
  }

  private boolean allReadersEOF() {
    return numUnconsumedSubpartitions <= 0;
  }

  private Optional<BufferOrEvent> transformBuffer(Buffer buf, InputChannelInfo info)
      throws IOException {
    return Optional.of(
        new BufferOrEvent(decompressBufferIfNeeded(buf, info), info, !isFinished(), false));
  }

  private Optional<BufferOrEvent> transformEvent(Buffer buffer, InputChannelInfo channelInfo)
      throws IOException {
    final AbstractEvent event;
    try {
      event = EventSerializer.fromBuffer(buffer, getClass().getClassLoader());
    } catch (Throwable t) {
      LOG.error("Failed to deserialize event: inputChannelInfo {}.", channelInfo, t);
      throw new PartitionNotFoundException(channelToResultPartitionId.get(channelInfo));
    } finally {
      buffer.recycleBuffer();
    }
    if (event.getClass() == EndOfPartitionEvent.class) {
      checkState(
          numSubPartitionsNotConsumed[channelInfo.getInputChannelIdx()] > 0,
          "BUG -- EndOfPartitionEvent received repeatedly.");
      numSubPartitionsNotConsumed[channelInfo.getInputChannelIdx()]--;
      numUnconsumedSubpartitions--;
      // not the real end.
      if (numSubPartitionsNotConsumed[channelInfo.getInputChannelIdx()] != 0) {
        LOG.debug(
            "numSubPartitionsNotConsumed: {}",
            numSubPartitionsNotConsumed[channelInfo.getInputChannelIdx()]);
      } else {
        // the real end.
        bufferReaders.get(channelIndexToReaderIndex[channelInfo.getInputChannelIdx()]).close();
        tryOpenSomeChannels();
        if (allReadersEOF()) {
          availabilityHelper.getUnavailableToResetAvailable().complete(null);
        }
      }
    } else if (event.getClass() == EndOfData.class) {
      checkState(!hasReceivedEndOfData(), "Too many EndOfData event.");
      --pendingEndOfDataEvents;
    }

    return Optional.of(
        new BufferOrEvent(
            event,
            buffer.getDataType().hasPriority(),
            channelInfo,
            !isFinished(),
            buffer.getSize(),
            false));
  }

  public boolean isFinished() {
    synchronized (lock) {
      return allReadersEOF() && receivedBuffers.isEmpty();
    }
  }

  public boolean hasReceivedEndOfData() {
    return pendingEndOfDataEvents <= 0;
  }

  public int getGateIndex() {
    return gateIndex;
  }

  public Optional<BufferOrEvent> pollNext() throws IOException {
    if (!initialChannelsOpened) {
      tryOpenSomeChannels();
      initialChannelsOpened = true;
      // DO NOT return, method of 'getReceived' will manipulate 'availabilityHelper'.
    }

    Pair<Buffer, InputChannelInfo> pair = getReceived();
    Optional<BufferOrEvent> bufferOrEvent = Optional.empty();
    LOG.debug("pollNext called with pair null {}", pair == null);
    while (pair != null) {
      Buffer buffer = pair.getLeft();
      InputChannelInfo channelInfo = pair.getRight();
      if (buffer.isBuffer()) {
        bufferOrEvent = transformBuffer(buffer, channelInfo);
      } else {
        bufferOrEvent = transformEvent(buffer, channelInfo);
        LOG.debug(
            "received event: {}.",
            bufferOrEvent.isPresent()
                ? bufferOrEvent.get().getEvent().getClass().getName()
                : Optional.empty());
      }

      if (bufferOrEvent.isPresent()) {
        break;
      }
      pair = getReceived();
    }

    tryRequestBuffers();
    return bufferOrEvent;
  }

  public void close() throws Exception {
    List<Buffer> buffersToRecycle;
    Throwable closeException = null;
    // Do not check closed flag, thus to allow calling this method from both task thread and
    // cancel thread.
    for (RemoteBufferStreamReader reader : bufferReaders) {
      try {
        reader.close();
      } catch (Throwable throwable) {
        closeException = closeException == null ? throwable : closeException;
        LOG.error("Failed to close shuffle read client.", throwable);
      }
    }
    synchronized (lock) {
      buffersToRecycle = receivedBuffers.stream().map(Pair::getLeft).collect(Collectors.toList());
      receivedBuffers.clear();
      closed = true;
    }

    try {
      buffersToRecycle.forEach(Buffer::recycleBuffer);
    } catch (Throwable throwable) {
      closeException = closeException == null ? throwable : closeException;
      LOG.error("Failed to recycle buffers.", throwable);
    }

    try {
      transferBufferPool.destroy();
    } catch (Throwable throwable) {
      closeException = closeException == null ? throwable : closeException;
      LOG.error("Failed to close transfer buffer pool.", throwable);
    }

    try {
      if (bufferPool != null) {
        bufferPool.lazyDestroy();
      }
    } catch (Throwable throwable) {
      closeException = closeException == null ? throwable : closeException;
      LOG.error("Failed to close local buffer pool.", throwable);
    }

    if (closeException != null) {
      ExceptionUtils.rethrowException(closeException);
    }
  }

  public String getTaskName() {
    return taskName;
  }

  public InputGateDeploymentDescriptor getGateDescriptor() {
    return gateDescriptor;
  }

  public long getPendingEndOfDataEvents() {
    return pendingEndOfDataEvents;
  }
}
