/*
 * 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 java.io.IOException;
import java.util.Optional;

import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.buffer.BufferPool;
import org.apache.flink.runtime.metrics.groups.ShuffleIOMetricGroup;
import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
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.identity.UserIdentifier;
import org.apache.celeborn.common.protocol.PartitionLocation;
import org.apache.celeborn.plugin.flink.buffer.BufferHeader;
import org.apache.celeborn.plugin.flink.buffer.BufferPacker;
import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
import org.apache.celeborn.plugin.flink.utils.BufferUtils;
import org.apache.celeborn.plugin.flink.utils.Utils;

/**
 * A transportation gate used to spill buffers from {@link ResultPartitionWriter} to remote shuffle
 * worker. The whole process of communication between outputGate and shuffle worker could be
 * described as below:
 *
 * <ul>
 *   <li>1. Client registers shuffle to get partitionLocation which stores the data in remote
 *       shuffle;
 *   <li>2. Client sends PushDataHandShake to transfer handshake message;
 *   <li>3. Client sends RegionStart which announces the start of a writing region and maybe get a
 *       new partitionLocation;
 *   <li>4. Client write data;
 *   <li>5. Client sends RegionFinish to indicate writing finish of a region;
 *   <li>6. Repeat from step-2 to step-5;
 *   <li>7. Client sends mapend to indicate writing finish;
 * </ul>
 */
public class RemoteShuffleOutputGate {
  private static final Logger LOG = LoggerFactory.getLogger(RemoteShuffleOutputGate.class);
  private final RemoteShuffleDescriptor shuffleDesc;
  protected final int numSubs;
  protected FlinkShuffleClientImpl flinkShuffleClient;
  protected final SupplierWithException<BufferPool, IOException> bufferPoolFactory;
  protected BufferPool bufferPool;
  private CelebornConf celebornConf;
  private final int numMappers;
  private PartitionLocation partitionLocation;

  private int currentRegionIndex = 0;

  private int bufferSize;
  private BufferPacker bufferPacker;
  private String applicationId;
  private int shuffleId;
  private int mapId;
  private int attemptId;
  private int partitionId;
  private String lifecycleManagerHost;
  private int lifecycleManagerPort;
  private long lifecycleManagerTimestamp;
  private UserIdentifier userIdentifier;
  private boolean isRegisterShuffle = false;
  private int maxReviveTimes;
  private boolean hasSentHandshake = false;
  protected final ShuffleIOMetricGroup shuffleIOMetricGroup;

  /**
   * @param shuffleDesc Describes shuffle meta and shuffle worker address.
   * @param numSubs Number of subpartitions of the corresponding {@link ResultPartitionWriter}.
   * @param bufferPoolFactory {@link BufferPool} provider.
   */
  public RemoteShuffleOutputGate(
      RemoteShuffleDescriptor shuffleDesc,
      int numSubs,
      int bufferSize,
      SupplierWithException<BufferPool, IOException> bufferPoolFactory,
      CelebornConf celebornConf,
      int numMappers,
      ShuffleIOMetricGroup shuffleIOMetricGroup) {
    this(
        shuffleDesc,
        numSubs,
        bufferSize,
        bufferPoolFactory,
        celebornConf,
        numMappers,
        shuffleIOMetricGroup,
        null);
  }

  public RemoteShuffleOutputGate(
      RemoteShuffleDescriptor shuffleDesc,
      int numSubs,
      int bufferSize,
      SupplierWithException<BufferPool, IOException> bufferPoolFactory,
      CelebornConf celebornConf,
      int numMappers,
      ShuffleIOMetricGroup shuffleIOMetricGroup,
      FlinkShuffleClientImpl flinkShuffleClient) {
    this.shuffleDesc = shuffleDesc;
    this.numSubs = numSubs;
    this.bufferPoolFactory = bufferPoolFactory;
    this.bufferPacker = new BufferPacker(this::write);
    this.celebornConf = celebornConf;
    this.numMappers = numMappers;
    this.bufferSize = bufferSize;
    this.applicationId = shuffleDesc.getCelebornAppId();
    this.shuffleId =
        shuffleDesc.getShuffleResource().getMapPartitionShuffleDescriptor().getShuffleId();
    this.mapId = shuffleDesc.getShuffleResource().getMapPartitionShuffleDescriptor().getMapId();
    this.attemptId =
        shuffleDesc.getShuffleResource().getMapPartitionShuffleDescriptor().getAttemptId();
    this.partitionId =
        shuffleDesc.getShuffleResource().getMapPartitionShuffleDescriptor().getPartitionId();
    this.lifecycleManagerHost = shuffleDesc.getShuffleResource().getLifecycleManagerHost();
    this.lifecycleManagerPort = shuffleDesc.getShuffleResource().getLifecycleManagerPort();
    this.lifecycleManagerTimestamp =
        shuffleDesc.getShuffleResource().getLifecycleManagerTimestamp();
    this.flinkShuffleClient = flinkShuffleClient == null ? getShuffleClient() : flinkShuffleClient;
    this.maxReviveTimes = celebornConf.clientPushMaxReviveTimes();
    this.shuffleIOMetricGroup = shuffleIOMetricGroup;
  }

  /** Initialize transportation gate. */
  public void setup() throws IOException, InterruptedException {
    bufferPool = Utils.checkNotNull(bufferPoolFactory.get());
    Utils.checkArgument(
        bufferPool.getNumberOfRequiredMemorySegments() >= 2,
        "Too few buffers for transfer, the minimum valid required size is 2.");

    // guarantee that we have at least one buffer
    BufferUtils.reserveNumRequiredBuffers(bufferPool, 1);
  }

  /** Get transportation buffer pool. */
  public BufferPool getBufferPool() {
    return bufferPool;
  }

  /** Writes a {@link Buffer} to a subpartition. */
  public void write(Buffer buffer, int subIdx) throws InterruptedException {
    bufferPacker.process(buffer, subIdx);
  }

  /**
   * Indicates the start of a region. A region of buffers guarantees the records inside are
   * completed.
   *
   * @param isBroadcast Whether it's a broadcast region.
   */
  public void regionStart(boolean isBroadcast) {
    try {
      registerShuffle();
      handshake();
      regionStartWithRevive(isBroadcast);
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
  }

  /**
   * Indicates the finish of a region. A region is always bounded by a pair of region-start and
   * region-finish.
   */
  public void regionFinish() throws InterruptedException {
    bufferPacker.drain();
    try {
      flinkShuffleClient.regionFinish(shuffleId, mapId, attemptId, partitionLocation);
      currentRegionIndex++;
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
  }

  /** Indicates the writing/spilling is finished. */
  public void finish() throws InterruptedException, IOException {
    flinkShuffleClient.mapPartitionMapperEnd(
        shuffleId, mapId, attemptId, numMappers, numSubs, partitionLocation.getId());
  }

  /** Close the transportation gate. */
  public void close() throws IOException {
    if (bufferPool != null) {
      bufferPool.lazyDestroy();
    }
    bufferPacker.close();
    flinkShuffleClient.cleanup(shuffleId, mapId, attemptId);
  }

  /** Returns shuffle descriptor. */
  public RemoteShuffleDescriptor getShuffleDesc() {
    return shuffleDesc;
  }

  @VisibleForTesting
  FlinkShuffleClientImpl getShuffleClient() {
    try {
      return FlinkShuffleClientImpl.get(
          applicationId,
          lifecycleManagerHost,
          lifecycleManagerPort,
          lifecycleManagerTimestamp,
          celebornConf,
          userIdentifier);
    } catch (DriverChangedException e) {
      // would generate a new attempt to retry output gate
      throw new RuntimeException(e.getMessage());
    }
  }

  /** Writes a piece of data to a subpartition. */
  public void write(ByteBuf byteBuf, BufferHeader bufferHeader) {
    try {
      int bytesWritten =
          flinkShuffleClient.pushDataToLocation(
              shuffleId,
              mapId,
              attemptId,
              bufferHeader.getSubPartitionId(),
              io.netty.buffer.Unpooled.wrappedBuffer(byteBuf.nioBuffer()),
              partitionLocation,
              byteBuf::release);
      shuffleIOMetricGroup.getNumBytesOut().inc(bytesWritten);
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
  }

  public void registerShuffle() throws IOException {
    if (!isRegisterShuffle) {
      partitionLocation =
          flinkShuffleClient.registerMapPartitionTask(
              shuffleId, numMappers, mapId, attemptId, partitionId);
      Utils.checkNotNull(partitionLocation);

      currentRegionIndex = 0;
      isRegisterShuffle = true;
    }
  }

  public void regionStartWithRevive(boolean isBroadcast) {
    try {
      int remainingReviveTimes = maxReviveTimes;
      boolean hasSentRegionStart = false;
      while (remainingReviveTimes-- > 0 && !hasSentRegionStart) {
        Optional<PartitionLocation> revivePartition =
            flinkShuffleClient.regionStart(
                shuffleId, mapId, attemptId, partitionLocation, currentRegionIndex, isBroadcast);
        if (revivePartition.isPresent()) {
          LOG.info(
              "Revive at regionStart, currentTimes:{}, totalTimes:{} for shuffleId:{}, mapId:{}, attempId:{}, currentRegionIndex:{}, isBroadcast:{}, newPartition:{}, oldPartition:{}",
              remainingReviveTimes,
              maxReviveTimes,
              shuffleId,
              mapId,
              attemptId,
              currentRegionIndex,
              isBroadcast,
              revivePartition,
              partitionLocation);
          partitionLocation = revivePartition.get();
          hasSentRegionStart = false;
          // For every revive partition, handshake should be sent firstly
          hasSentHandshake = false;
          handshake();
        } else {
          hasSentRegionStart = true;
        }
      }
      if (remainingReviveTimes == 0 && !hasSentRegionStart) {
        throw new RuntimeException(
            "After retry " + maxReviveTimes + " times, still failed to send regionStart");
      }
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
  }

  public void handshake() {
    try {
      int remainingReviveTimes = maxReviveTimes;
      while (remainingReviveTimes-- > 0 && !hasSentHandshake) {
        Optional<PartitionLocation> revivePartition =
            flinkShuffleClient.pushDataHandShake(
                shuffleId, mapId, attemptId, numSubs, bufferSize, partitionLocation);
        // if remainingReviveTimes == 0 and revivePartition.isPresent(), there is no need to send
        // handshake again
        if (revivePartition.isPresent() && remainingReviveTimes > 0) {
          LOG.info(
              "Revive at handshake, currentTimes:{}, totalTimes:{} for shuffleId:{}, mapId:{}, attempId:{}, currentRegionIndex:{}, newPartition:{}, oldPartition:{}",
              remainingReviveTimes,
              maxReviveTimes,
              shuffleId,
              mapId,
              attemptId,
              currentRegionIndex,
              revivePartition,
              partitionLocation);
          partitionLocation = revivePartition.get();
          hasSentHandshake = false;
        } else {
          hasSentHandshake = true;
        }
      }
      if (remainingReviveTimes == 0 && !hasSentHandshake) {
        throw new RuntimeException(
            "After retry " + maxReviveTimes + " times, still failed to send handshake");
      }
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
  }
}
