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

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageInputChannelId;
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.shaded.netty4.io.netty.buffer.ByteBuf;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
import org.apache.celeborn.common.network.protocol.BufferStreamEnd;
import org.apache.celeborn.common.network.protocol.RequestMessage;
import org.apache.celeborn.common.network.protocol.TransportableError;
import org.apache.celeborn.common.network.util.NettyUtils;
import org.apache.celeborn.common.protocol.PbNotifyRequiredSegment;
import org.apache.celeborn.common.protocol.PbReadAddCredit;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.plugin.flink.ShuffleResourceDescriptor;
import org.apache.celeborn.plugin.flink.client.CelebornBufferStream;
import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
import org.apache.celeborn.plugin.flink.protocol.SubPartitionReadData;

/** Wrap the {@link CelebornBufferStream}, used in flink hybrid shuffle integration strategy now. */
public class CelebornChannelBufferReader {
  private static final Logger LOG = LoggerFactory.getLogger(CelebornChannelBufferReader.class);

  private CelebornChannelBufferManager bufferManager;

  private final FlinkShuffleClientImpl client;

  private final int shuffleId;

  private final int partitionId;

  private final TieredStorageInputChannelId inputChannelId;

  private final int subPartitionIndexStart;

  private final int subPartitionIndexEnd;

  private final BiConsumer<ByteBuf, TieredStorageSubpartitionId> dataListener;

  private final BiConsumer<Throwable, TieredStorageSubpartitionId> failureListener;

  private final Consumer<RequestMessage> messageConsumer;

  private CelebornBufferStream bufferStream;

  private boolean isOpened;

  private volatile boolean closed = false;

  private volatile ConcurrentHashMap<Integer, Integer> subPartitionRequiredSegmentIds;

  /** Note this field is to record the number of backlog before the read is set up. */
  private int numBackLog = 0;

  public CelebornChannelBufferReader(
      FlinkShuffleClientImpl client,
      ShuffleResourceDescriptor shuffleDescriptor,
      TieredStorageInputChannelId inputChannelId,
      int startSubIdx,
      int endSubIdx,
      BiConsumer<ByteBuf, TieredStorageSubpartitionId> dataListener,
      BiConsumer<Throwable, TieredStorageSubpartitionId> failureListener) {
    this.client = client;
    this.shuffleId = shuffleDescriptor.getShuffleId();
    this.partitionId = shuffleDescriptor.getPartitionId();
    this.inputChannelId = inputChannelId;
    this.subPartitionIndexStart = startSubIdx;
    this.subPartitionIndexEnd = endSubIdx;
    this.dataListener = dataListener;
    this.failureListener = failureListener;
    this.subPartitionRequiredSegmentIds = JavaUtils.newConcurrentHashMap();
    for (int subPartitionId = subPartitionIndexStart;
        subPartitionId <= subPartitionIndexEnd;
        subPartitionId++) {
      subPartitionRequiredSegmentIds.put(subPartitionId, -1);
    }
    this.messageConsumer =
        requestMessage -> {
          // Note that we need to use SubPartitionReadData because the isSegmentGranularityVisible
          // is set as true when opening stream
          if (requestMessage instanceof SubPartitionReadData) {
            dataReceived((SubPartitionReadData) requestMessage);
          } else if (requestMessage instanceof BacklogAnnouncement) {
            backlogReceived(((BacklogAnnouncement) requestMessage).getBacklog());
          } else if (requestMessage instanceof TransportableError) {
            errorReceived(((TransportableError) requestMessage).getErrorMessage());
          } else if (requestMessage instanceof BufferStreamEnd) {
            onStreamEnd((BufferStreamEnd) requestMessage);
          }
        };
  }

  public void setup(TieredStorageMemoryManager memoryManager) {
    this.bufferManager = new CelebornChannelBufferManager(memoryManager, this);
    if (numBackLog > 0) {
      int numRequestedBuffers = bufferManager.requestBuffers(numBackLog);
      if (numRequestedBuffers > 0) {
        bufferManager.decreaseRequiredCredits(numRequestedBuffers);
        notifyAvailableCredits(numRequestedBuffers);
      }
      numBackLog = 0;
    }
  }

  public void open(int initialCredit, boolean sync) {
    try {
      bufferStream =
          client.readBufferedPartition(
              shuffleId, partitionId, subPartitionIndexStart, subPartitionIndexEnd, true);
      bufferStream.open(this::requestBuffer, initialCredit, messageConsumer, sync);
    } catch (Exception e) {
      messageConsumer.accept(new TransportableError(0L, e));
      LOG.error("Failed to open reader", e);
    }
  }

  public void close() {
    // It may be call multiple times because subPartitions can share the same reader, as a single
    // reader can consume multiple subPartitions
    if (closed) {
      return;
    }

    // need set closed first before remove Handler
    closed = true;
    if (!CelebornBufferStream.isEmptyStream(bufferStream)) {
      bufferStream.close();
      bufferStream = null;
    } else {
      LOG.warn(
          "bufferStream is null when closed, shuffleId: {}, partitionId: {}",
          shuffleId,
          partitionId);
    }

    try {
      if (bufferManager != null) {
        bufferManager.close();
        bufferManager = null;
      }
    } catch (Throwable throwable) {
      LOG.warn("Failed to close buffer manager.", throwable);
    }

    subPartitionRequiredSegmentIds.clear();
    subPartitionRequiredSegmentIds = null;
  }

  public boolean isOpened() {
    return isOpened;
  }

  public void setOpened(boolean opened) {
    isOpened = opened;
  }

  boolean isClosed() {
    return closed;
  }

  public void notifyAvailableCredits(int numCredits) {
    if (numCredits <= 0) {
      return;
    }
    if (!closed && !CelebornBufferStream.isEmptyStream(bufferStream)) {
      bufferStream.addCredit(
          PbReadAddCredit.newBuilder()
              .setStreamId(bufferStream.getStreamId())
              .setCredit(numCredits)
              .build());
      return;
    }
    LOG.warn(
        "The buffer stream is null or closed, ignore the credits for shuffleId: {}, partitionId: {}",
        shuffleId,
        partitionId);
  }

  public void notifyRequiredSegmentIfNeeded(int requiredSegmentId, int subPartitionId) {
    Integer lastRequiredSegmentId =
        subPartitionRequiredSegmentIds.computeIfAbsent(subPartitionId, id -> -1);
    if (requiredSegmentId >= 0 && requiredSegmentId != lastRequiredSegmentId) {
      LOG.debug(
          "Notify required segment id {} for {} {}, the last segment id is {}",
          requiredSegmentId,
          partitionId,
          subPartitionId,
          lastRequiredSegmentId);
      subPartitionRequiredSegmentIds.put(subPartitionId, requiredSegmentId);
      if (!this.notifyRequiredSegment(requiredSegmentId, subPartitionId)) {
        // if fail to notify reader segment, restore the last required segment id
        subPartitionRequiredSegmentIds.putIfAbsent(subPartitionId, lastRequiredSegmentId);
      }
    }
  }

  public boolean notifyRequiredSegment(int requiredSegmentId, int subPartitionId) {
    this.subPartitionRequiredSegmentIds.put(subPartitionId, requiredSegmentId);
    if (!closed && !CelebornBufferStream.isEmptyStream(bufferStream)) {
      LOG.debug(
          "Notify required segmentId {} for {} {} {}",
          requiredSegmentId,
          partitionId,
          subPartitionId,
          shuffleId);
      PbNotifyRequiredSegment notifyRequiredSegment =
          PbNotifyRequiredSegment.newBuilder()
              .setStreamId(bufferStream.getStreamId())
              .setRequiredSegmentId(requiredSegmentId)
              .setSubPartitionId(subPartitionId)
              .build();
      bufferStream.notifyRequiredSegment(notifyRequiredSegment);
      return true;
    }
    return false;
  }

  public ByteBuf requestBuffer() {
    Buffer buffer = bufferManager.requestBuffer();
    return buffer == null ? null : buffer.asByteBuf();
  }

  public void backlogReceived(int backlog) {
    if (!closed) {
      if (bufferManager == null) {
        numBackLog += backlog;
        return;
      }
      int numRequestedBuffers = bufferManager.requestBuffers(backlog);
      if (numRequestedBuffers > 0) {
        bufferManager.decreaseRequiredCredits(numRequestedBuffers);
        notifyAvailableCredits(numRequestedBuffers);
      }
      numBackLog = 0;
      return;
    }
    LOG.warn(
        "The buffer stream {} is null or closed, ignore the backlog for shuffleId: {}, partitionId: {}",
        bufferStream.getStreamId(),
        shuffleId,
        partitionId);
  }

  public void errorReceived(String errorMsg) {
    if (!closed) {
      closed = true;
      LOG.debug("Error received, " + errorMsg);
      if (!CelebornBufferStream.isEmptyStream(bufferStream) && bufferStream.getClient() != null) {
        LOG.error(
            "Received error from {} message {}",
            NettyUtils.getRemoteAddress(bufferStream.getClient().getChannel()),
            errorMsg);
      }
      for (int subPartitionId = subPartitionIndexStart;
          subPartitionId <= subPartitionIndexEnd;
          subPartitionId++) {
        failureListener.accept(
            new IOException(errorMsg), new TieredStorageSubpartitionId(subPartitionId));
      }
    }
  }

  public void dataReceived(SubPartitionReadData readData) {
    LOG.debug(
        "Remote buffer stream reader get stream id {} subPartitionId {} received readable bytes {}.",
        readData.getStreamId(),
        readData.getSubPartitionId(),
        readData.getFlinkBuffer().readableBytes());
    checkState(
        readData.getSubPartitionId() >= subPartitionIndexStart
            && readData.getSubPartitionId() <= subPartitionIndexEnd,
        "Wrong sub partition id: " + readData.getSubPartitionId());
    dataListener.accept(
        readData.getFlinkBuffer(), new TieredStorageSubpartitionId(readData.getSubPartitionId()));
    int numRequested = bufferManager.tryRequestBuffersIfNeeded();
    if (numRequested > 0) {
      bufferManager.decreaseRequiredCredits(numRequested);
      notifyAvailableCredits(numRequested);
    }
  }

  public void onStreamEnd(BufferStreamEnd streamEnd) {
    long streamId = streamEnd.getStreamId();
    LOG.debug("Buffer stream reader get stream end for {}", streamId);
    if (!closed && !CelebornBufferStream.isEmptyStream(bufferStream)) {
      // TOOD: Update the partition locations here if support reading and writing shuffle data
      // simultaneously
      bufferStream.moveToNextPartitionIfPossible(streamId, this::sendRequireSegmentId, true);
    }
  }

  public TieredStorageInputChannelId getInputChannelId() {
    return inputChannelId;
  }

  private void sendRequireSegmentId(long streamId, int subPartitionId) {
    if (subPartitionRequiredSegmentIds.containsKey(subPartitionId)) {
      int currentSegmentId = subPartitionRequiredSegmentIds.get(subPartitionId);
      if (bufferStream.isOpened() && currentSegmentId >= 0) {
        LOG.debug(
            "Buffer stream {} is opened, notify required segment id {} ",
            streamId,
            currentSegmentId);
        notifyRequiredSegment(currentSegmentId, subPartitionId);
      }
    }
  }
}
