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

import static org.apache.celeborn.plugin.flink.utils.Utils.checkArgument;
import static org.apache.celeborn.plugin.flink.utils.Utils.checkState;
import static org.apache.flink.runtime.io.network.buffer.Buffer.DataType.END_OF_SEGMENT;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;

import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.core.memory.MemorySegmentFactory;
import org.apache.flink.runtime.io.network.api.EndOfSegmentEvent;
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.FreeingBufferRecycler;
import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageMemoryManager;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageResourceRegistry;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierProducerAgent;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleDescriptor;
import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
import org.apache.flink.shaded.netty4.io.netty.buffer.CompositeByteBuf;
import org.apache.flink.shaded.netty4.io.netty.buffer.Unpooled;
import org.apache.flink.util.ExceptionUtils;
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.protocol.PartitionLocation;
import org.apache.celeborn.plugin.flink.buffer.BufferHeader;
import org.apache.celeborn.plugin.flink.buffer.BufferPacker;
import org.apache.celeborn.plugin.flink.buffer.ReceivedNoHeaderBufferPacker;
import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
import org.apache.celeborn.plugin.flink.utils.BufferUtils;
import org.apache.celeborn.plugin.flink.utils.Utils;

public class CelebornTierProducerAgent implements TierProducerAgent {

  private static final Logger LOG = LoggerFactory.getLogger(CelebornTierProducerAgent.class);

  private final int numBuffersPerSegment;

  // The flink buffer size in bytes.
  private final int bufferSizeBytes;

  private final int numPartitions;

  private final int numSubPartitions;

  private final CelebornConf celebornConf;

  private final TieredStorageMemoryManager memoryManager;

  private final String applicationId;

  private final int shuffleId;

  private final int mapId;

  private final int attemptId;

  private final int partitionId;

  private final String lifecycleManagerHost;

  private final int lifecycleManagerPort;

  private final long lifecycleManagerTimestamp;

  private FlinkShuffleClientImpl flinkShuffleClient;

  private BufferPacker bufferPacker;

  private final int[] subPartitionSegmentIds;

  private final int[] subPartitionSegmentBuffers;

  private final int maxReviveTimes;

  private PartitionLocation partitionLocation;

  private boolean hasRegisteredShuffle;

  private int currentRegionIndex = 0;

  private int currentSubpartition = 0;

  private boolean hasSentHandshake = false;

  private boolean hasSentRegionStart = false;

  private volatile boolean isReleased;

  CelebornTierProducerAgent(
      CelebornConf conf,
      TieredStoragePartitionId partitionId,
      int numPartitions,
      int numSubPartitions,
      int numBytesPerSegment,
      int bufferSizeBytes,
      TieredStorageMemoryManager memoryManager,
      TieredStorageResourceRegistry resourceRegistry,
      List<TierShuffleDescriptor> shuffleDescriptors) {
    checkArgument(
        numBytesPerSegment >= bufferSizeBytes, "One segment should contain at least one buffer.");
    checkArgument(shuffleDescriptors.size() == 1, "There should be only one shuffle descriptor.");
    TierShuffleDescriptor descriptor = shuffleDescriptors.get(0);
    checkArgument(
        descriptor instanceof TierShuffleDescriptorImpl,
        "Wrong shuffle descriptor type " + descriptor.getClass());
    TierShuffleDescriptorImpl shuffleDesc = (TierShuffleDescriptorImpl) descriptor;

    this.numBuffersPerSegment = numBytesPerSegment / bufferSizeBytes;
    this.bufferSizeBytes = bufferSizeBytes;
    this.memoryManager = memoryManager;
    this.numPartitions = numPartitions;
    this.numSubPartitions = numSubPartitions;
    this.celebornConf = conf;
    this.subPartitionSegmentIds = new int[numSubPartitions];
    this.subPartitionSegmentBuffers = new int[numSubPartitions];
    this.maxReviveTimes = conf.clientPushMaxReviveTimes();

    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 = getShuffleClient();

    Arrays.fill(subPartitionSegmentIds, -1);
    Arrays.fill(subPartitionSegmentBuffers, 0);

    this.bufferPacker = new ReceivedNoHeaderBufferPacker(this::write);
    resourceRegistry.registerResource(partitionId, this::releaseResources);
    registerShuffle();
    try {
      handshake();
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
  }

  @Override
  public boolean tryStartNewSegment(
      TieredStorageSubpartitionId tieredStorageSubpartitionId, int segmentId, int minNumBuffers) {
    int subPartitionId = tieredStorageSubpartitionId.getSubpartitionId();
    checkState(
        segmentId >= subPartitionSegmentIds[subPartitionId], "Wrong segment id " + segmentId);
    subPartitionSegmentIds[subPartitionId] = segmentId;
    // If the start segment rpc is sent, the worker side will expect that
    // there must be at least one buffer will be written in the next moment.
    try {
      flinkShuffleClient.segmentStart(
          shuffleId, mapId, attemptId, subPartitionId, segmentId, partitionLocation);
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
    return true;
  }

  @Override
  public boolean tryWrite(
      TieredStorageSubpartitionId tieredStorageSubpartitionId,
      Buffer buffer,
      Object bufferOwner,
      int numRemainingConsecutiveBuffers) {
    // It should be noted that, unlike RemoteShuffleOutputGate#write, the received buffer contains
    // only
    // and does not have any remaining space for writing the celeborn header.

    int subPartitionId = tieredStorageSubpartitionId.getSubpartitionId();

    if (subPartitionSegmentBuffers[subPartitionId] + 1 + numRemainingConsecutiveBuffers
        >= numBuffersPerSegment) {
      // End the current segment if the segment buffer count reaches the threshold
      subPartitionSegmentBuffers[subPartitionId] = 0;
      try {
        bufferPacker.drain();
      } catch (InterruptedException e) {
        buffer.recycleBuffer();
        ExceptionUtils.rethrow(e, "Failed to process buffer.");
      }
      appendEndOfSegmentBuffer(subPartitionId);
      return false;
    }

    if (buffer.isBuffer()) {
      memoryManager.transferBufferOwnership(
          bufferOwner, CelebornTierFactory.getCelebornTierName(), buffer);
    }

    // write buffer to BufferPacker and record buffer count per subPartition per segment
    processBuffer(buffer, subPartitionId);
    subPartitionSegmentBuffers[subPartitionId]++;
    return true;
  }

  @Override
  public void close() {
    if (hasSentRegionStart) {
      regionFinish();
    }
    try {
      if (hasRegisteredShuffle && partitionLocation != null) {
        flinkShuffleClient.mapPartitionMapperEnd(
            shuffleId,
            mapId,
            attemptId,
            numPartitions,
            numSubPartitions,
            partitionLocation.getId());
      }
    } catch (Exception e) {
      Utils.rethrowAsRuntimeException(e);
    }
    bufferPacker.close();
    bufferPacker = null;
    flinkShuffleClient.cleanup(shuffleId, mapId, attemptId);
    flinkShuffleClient = null;
  }

  private void regionStartOrFinish(int subPartitionId) {
    // check whether the region should be started or finished
    regionStart();
    if (subPartitionId < currentSubpartition) {
      // if the consumed subPartitionId is out of order, it means that should the previous region
      // should be finished, and starting a new region.
      regionFinish();
      LOG.debug(
          "Check region finish sub partition id {} and start next region {}",
          subPartitionId,
          currentRegionIndex);
      regionStart();
    }
  }

  private void regionStart() {
    if (hasSentRegionStart) {
      return;
    }
    regionStartWithRevive();
  }

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

  void regionFinish() {
    try {
      bufferPacker.drain();
      flinkShuffleClient.regionFinish(shuffleId, mapId, attemptId, partitionLocation);
      hasSentRegionStart = false;
      currentRegionIndex++;
    } catch (Exception e) {
      Utils.rethrowAsRuntimeException(e);
    }
  }

  private void handshake() throws IOException {
    try {
      int remainingReviveTimes = maxReviveTimes;
      while (remainingReviveTimes-- > 0 && !hasSentHandshake) {
        // In the Flink hybrid shuffle integration strategy, the data buffer sent to the Celeborn
        // workers consists of two components: the Celeborn header and the data buffers.
        // In this scenario, the maximum byte size of the buffer received by the Celeborn worker is
        // equal to the sum of the Flink buffer size and the Celeborn header size.
        Optional<PartitionLocation> revivePartition =
            flinkShuffleClient.pushDataHandShake(
                shuffleId,
                mapId,
                attemptId,
                numSubPartitions,
                bufferSizeBytes + BufferUtils.HEADER_LENGTH,
                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);
    }
  }

  private void releaseResources() {
    if (!isReleased) {
      isReleased = true;
    }
  }

  private void registerShuffle() {
    try {
      if (!hasRegisteredShuffle) {
        partitionLocation =
            flinkShuffleClient.registerMapPartitionTask(
                shuffleId, numPartitions, mapId, attemptId, partitionId, true);
        Utils.checkNotNull(partitionLocation);
        hasRegisteredShuffle = true;
      }
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
  }

  private void write(ByteBuf byteBuf, BufferHeader bufferHeader) {
    try {
      // create a composite buffer and write a header into it. This composite buffer will serve as
      // the result packed buffer.
      CompositeByteBuf compositeByteBuf = Unpooled.compositeBuffer();
      ByteBuf headerBuf = Unpooled.buffer(BufferUtils.HEADER_LENGTH);

      // write celeborn buffer header (subpartitionid(4) + attemptId(4) + nextBatchId(4) +
      // compressedsize)
      headerBuf.writeInt(bufferHeader.getSubPartitionId());
      headerBuf.writeInt(attemptId);
      headerBuf.writeInt(0);
      headerBuf.writeInt(
          byteBuf.readableBytes() + (BufferUtils.HEADER_LENGTH - BufferUtils.HEADER_LENGTH_PREFIX));

      // write flink buffer header (dataType(1) + isCompress(1) + size(4))
      headerBuf.writeByte(bufferHeader.getDataType().ordinal());
      headerBuf.writeBoolean(bufferHeader.isCompressed());
      headerBuf.writeInt(bufferHeader.getSize());

      // composite the headerBuf and data buffer together
      compositeByteBuf.addComponents(true, headerBuf, byteBuf);
      io.netty.buffer.ByteBuf wrappedBuffer =
          io.netty.buffer.Unpooled.wrappedBuffer(compositeByteBuf.nioBuffer());

      int numWritten =
          flinkShuffleClient.pushDataToLocation(
              shuffleId,
              mapId,
              attemptId,
              bufferHeader.getSubPartitionId(),
              wrappedBuffer,
              partitionLocation,
              compositeByteBuf::release);
      checkState(
          numWritten == byteBuf.readableBytes() + BufferUtils.HEADER_LENGTH, "Wrong written size.");
    } catch (IOException e) {
      Utils.rethrowAsRuntimeException(e);
    }
  }

  private void appendEndOfSegmentBuffer(int subPartitionId) {
    try {
      checkState(bufferPacker.isEmpty(), "BufferPacker is not empty");
      MemorySegment endSegmentMemorySegment =
          MemorySegmentFactory.wrap(
              EventSerializer.toSerializedEvent(EndOfSegmentEvent.INSTANCE).array());
      Buffer endOfSegmentBuffer =
          new NetworkBuffer(
              endSegmentMemorySegment,
              FreeingBufferRecycler.INSTANCE,
              END_OF_SEGMENT,
              endSegmentMemorySegment.size());
      processBuffer(endOfSegmentBuffer, subPartitionId);
      // drain the bufferPacker to ensure that the EndOfSegment event is not bundled with the data
      // buffer
      bufferPacker.drain();
    } catch (Exception e) {
      ExceptionUtils.rethrow(e, "Failed to append end of segment event.");
    }
  }

  private void processBuffer(Buffer originBuffer, int subPartitionId) {
    try {
      regionStartOrFinish(subPartitionId);
      currentSubpartition = subPartitionId;

      Buffer buffer = originBuffer;
      if (originBuffer.isCompressed()) {
        // Flink hybrid shuffle will send a compressed buffer to tier. However, since we need to
        // write data to this buffer and the compressed buffer is read-only, we must create a
        // new Buffer object to the wrap origin buffer.
        NetworkBuffer networkBuffer =
            new NetworkBuffer(
                originBuffer.getMemorySegment(),
                originBuffer.getRecycler(),
                originBuffer.getDataType(),
                originBuffer.getSize());
        networkBuffer.writerIndex(originBuffer.asByteBuf().writerIndex());
        buffer = networkBuffer;
      }

      // set the buffer meta
      BufferUtils.setCompressedDataWithoutHeader(buffer, originBuffer);

      bufferPacker.process(buffer, subPartitionId);
    } catch (InterruptedException e) {
      originBuffer.recycleBuffer();
      ExceptionUtils.rethrow(e, "Failed to process buffer.");
    }
  }

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