/*
 * 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.service.deploy.worker.storage;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.FileChannel;
import java.util.ArrayDeque;
import java.util.Queue;
import java.util.concurrent.atomic.AtomicInteger;

import javax.annotation.concurrent.GuardedBy;

import com.google.common.annotations.VisibleForTesting;
import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFutureListener;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.exception.FileCorruptedException;
import org.apache.celeborn.common.meta.DiskFileInfo;
import org.apache.celeborn.common.meta.FileInfo;
import org.apache.celeborn.common.meta.MapFileMeta;
import org.apache.celeborn.common.network.buffer.NioManagedBuffer;
import org.apache.celeborn.common.network.client.TransportClient;
import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
import org.apache.celeborn.common.network.protocol.ReadData;
import org.apache.celeborn.common.network.protocol.RequestMessage;
import org.apache.celeborn.common.network.protocol.RpcRequest;
import org.apache.celeborn.common.network.protocol.TransportMessage;
import org.apache.celeborn.common.network.protocol.TransportableError;
import org.apache.celeborn.common.network.util.NettyUtils;
import org.apache.celeborn.common.protocol.MessageType;
import org.apache.celeborn.common.protocol.PbBufferStreamEnd;
import org.apache.celeborn.common.protocol.StreamType;
import org.apache.celeborn.common.util.ExceptionUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.service.deploy.worker.memory.BufferQueue;
import org.apache.celeborn.service.deploy.worker.memory.BufferRecycler;
import org.apache.celeborn.service.deploy.worker.memory.RecyclableBuffer;

public class MapPartitionDataReader implements Comparable<MapPartitionDataReader> {
  private static final Logger logger = LoggerFactory.getLogger(MapPartitionDataReader.class);

  private final ByteBuffer indexBuffer;
  private final ByteBuffer headerBuffer;
  private final int startPartitionIndex;
  private final int endPartitionIndex;
  private int numRegions;
  private int numRemainingPartitions;
  private int currentDataRegion = -1;
  private long dataConsumingOffset;
  private volatile long currentPartitionRemainingBytes;
  private DiskFileInfo fileInfo;
  protected MapFileMeta mapFileMeta;
  private int INDEX_ENTRY_SIZE = 16;
  protected long streamId;
  protected final Object lock = new Object();

  protected final AtomicInteger credits = new AtomicInteger();

  @GuardedBy("lock")
  protected final Queue<RecyclableBuffer> buffersToSend = new ArrayDeque<>();

  /** Whether all the data has been successfully read or not. */
  @GuardedBy("lock")
  private volatile boolean readFinished;

  /** Whether this partition reader has been released or not. */
  @GuardedBy("lock")
  protected boolean isReleased;

  /** Exception causing the release of this partition reader. */
  @GuardedBy("lock")
  protected Throwable errorCause;

  /** Whether there is any error at the consumer side or not. */
  @GuardedBy("lock")
  protected boolean errorNotified;

  private FileChannel dataFileChannel;

  // The size of the data file, it is initialized in the open method and remains unchanged
  // afterward.
  private long dataFileChannelSize;
  private FileChannel indexFileChannel;

  private Channel associatedChannel;

  private Runnable recycleStream;

  protected AtomicInteger numInUseBuffers = new AtomicInteger(0);
  private boolean isOpen = false;

  public MapPartitionDataReader(
      int startPartitionIndex,
      int endPartitionIndex,
      DiskFileInfo fileInfo,
      long streamId,
      Channel associatedChannel,
      Runnable recycleStream) {
    this.startPartitionIndex = startPartitionIndex;
    this.endPartitionIndex = endPartitionIndex;

    int indexBufferSize = 16 * (endPartitionIndex - startPartitionIndex + 1);
    this.indexBuffer = ByteBuffer.allocateDirect(indexBufferSize);

    this.headerBuffer = ByteBuffer.allocateDirect(16);
    this.streamId = streamId;
    this.associatedChannel = associatedChannel;
    this.recycleStream = recycleStream;

    this.fileInfo = fileInfo;
    this.mapFileMeta = ((MapFileMeta) fileInfo.getFileMeta());
    this.readFinished = false;
  }

  public void open(FileChannel dataFileChannel, FileChannel indexFileChannel, long indexSize)
      throws IOException {
    if (!isOpen) {
      this.dataFileChannel = dataFileChannel;
      this.dataFileChannelSize = dataFileChannel.size();
      this.indexFileChannel = indexFileChannel;

      int numSubpartitions = mapFileMeta.getNumSubpartitions();
      // If numSubpartitions is 0, it means pushDataHandShake was never successfully called.
      // This can happen when the first handshake failed before any data was written.
      // In this case, check if data file is empty, and if so, treat this as an empty partition.
      if (numSubpartitions == 0) {
        if (dataFileChannelSize == 0) {
          logger.warn(
              "Partition {} has numSubpartitions=0 and empty data file, this indicates a failed "
                  + "handshake before any data was written. Treating as empty partition.",
              fileInfo.getFilePath());
          isOpen = true;
          // Mark as finished and notify consumer with backlog=1 so they can complete normally
          closeReader();
          notifyBacklog(1);
          return;
        } else {
          // Data file has content but numSubpartitions is 0, this is a corrupted state
          throw new FileCorruptedException(
              "Partition "
                  + fileInfo.getFilePath()
                  + " has numSubpartitions=0 but data file size is "
                  + dataFileChannelSize);
        }
      }

      // index is (offset,length)
      long indexRegionSize = numSubpartitions * (long) INDEX_ENTRY_SIZE;
      this.numRegions = Utils.checkedDownCast(indexSize / indexRegionSize);

      updateConsumingOffset();
      isOpen = true;
    }
  }

  public void addCredit(int credit) {
    credits.getAndAdd(credit);
  }

  public void readData(BufferQueue bufferQueue, BufferRecycler bufferRecycler) throws IOException {
    boolean hasRemaining = hasRemaining();
    boolean continueReading = hasRemaining;
    int numDataBuffers = 0;
    while (continueReading) {

      ByteBuf buffer = bufferQueue.poll();
      if (buffer == null) {
        // if there are no buffers available, halt current read and waiting for next triggered read
        break;
      } else {
        buffer.retain();
        numInUseBuffers.incrementAndGet();
      }

      try {
        continueReading = readBuffer(buffer);
      } catch (Throwable throwable) {
        bufferRecycler.recycle(buffer);
        numInUseBuffers.decrementAndGet();
        throw throwable;
      }

      hasRemaining = hasRemaining();
      addBuffer(buffer, bufferRecycler);
      ++numDataBuffers;
    }

    if (!hasRemaining) {
      closeReader();
    }

    tryNotifyBacklog(numDataBuffers);
  }

  protected void tryNotifyBacklog(int numDataBuffers) {
    if (numDataBuffers > 0) {
      notifyBacklog(numDataBuffers);
    }
  }

  protected void addBuffer(ByteBuf buffer, BufferRecycler bufferRecycler) {
    if (buffer == null) {
      return;
    }
    synchronized (lock) {
      if (!isReleased) {
        buffersToSend.add(new RecyclableBuffer(buffer, -1, bufferRecycler));
      } else {
        bufferRecycler.recycle(buffer);
        numInUseBuffers.decrementAndGet();
        throw new RuntimeException("Partition reader has been failed or finished.", errorCause);
      }
    }
  }

  protected RecyclableBuffer fetchBufferToSend() {
    synchronized (lock) {
      if (!buffersToSend.isEmpty() && credits.get() > 0 && !isReleased) {
        return buffersToSend.poll();
      } else {
        return null;
      }
    }
  }

  protected int getNumBuffersToSend() {
    synchronized (lock) {
      return buffersToSend.size();
    }
  }

  public synchronized void sendData() {
    RecyclableBuffer buffer;
    while (null != (buffer = fetchBufferToSend())) {
      sendDataInternal(buffer);
    }

    tryRecycleReader();
  }

  public void tryRecycleReader() {
    boolean shouldRecycle = false;
    synchronized (lock) {
      if (isReleased) return;
      if (readFinished && buffersToSend.isEmpty()) {
        shouldRecycle = true;
      }
    }

    if (shouldRecycle) {
      recycle();
    }
  }

  public RequestMessage generateReadDataMessage(
      long streamId, int subPartitionId, ByteBuf byteBuf) {
    return new ReadData(streamId, byteBuf);
  }

  protected void sendDataInternal(RecyclableBuffer buffer) {
    final RecyclableBuffer wrappedBuffer = buffer;
    int readableBytes = wrappedBuffer.byteBuf.readableBytes();
    if (logger.isDebugEnabled()) {
      logger.debug("send data start: {}, {}", streamId, readableBytes);
    }

    RequestMessage readData =
        generateReadDataMessage(streamId, wrappedBuffer.subPartitionId, wrappedBuffer.byteBuf);
    associatedChannel
        .writeAndFlush(readData)
        .addListener(
            (ChannelFutureListener)
                future -> {
                  try {
                    if (!future.isSuccess()) {
                      recycleOnError(future.cause());
                    }
                  } finally {
                    if (logger.isDebugEnabled()) {
                      logger.debug("send data end: {}, {}", streamId, readableBytes);
                    }
                    wrappedBuffer.recycle();
                    numInUseBuffers.decrementAndGet();
                  }
                });
    int currentCredit = credits.decrementAndGet();
    logger.debug("Current credit is {} after stream {}", currentCredit, streamId);
  }

  protected long getIndexRegionSize() {
    return mapFileMeta.getNumSubpartitions() * (long) INDEX_ENTRY_SIZE;
  }

  protected void readHeaderOrIndexBuffer(FileChannel channel, ByteBuffer buffer, int length)
      throws IOException {
    Utils.checkFileIntegrity(channel, length);
    buffer.clear();
    buffer.limit(length);
    while (buffer.hasRemaining()) {
      channel.read(buffer);
    }
    buffer.flip();
  }

  protected void readBufferIntoReadBuffer(FileChannel channel, ByteBuf buf, int length)
      throws IOException {
    Utils.checkFileIntegrity(channel, length);
    ByteBuffer tmpBuffer = ByteBuffer.allocate(length);
    while (tmpBuffer.hasRemaining()) {
      channel.read(tmpBuffer);
    }
    tmpBuffer.flip();
    buf.writeBytes(tmpBuffer);
  }

  protected int readBuffer(
      String filename, FileChannel channel, ByteBuffer header, ByteBuf buffer, int headerSize)
      throws IOException {
    readHeaderOrIndexBuffer(channel, header, headerSize);
    // header is combined of mapId(4),attemptId(4),nextBatchId(4) and total Compressed Length(4)
    // we need size here,so we read length directly
    int bufferLength = header.getInt(12);
    if (bufferLength <= 0 || bufferLength > buffer.capacity()) {
      logger.error("Incorrect buffer header, buffer length: {}.", bufferLength);
      throw new FileCorruptedException("File " + filename + " is corrupted");
    }
    buffer.writeBytes(header);
    readBufferIntoReadBuffer(channel, buffer, bufferLength);
    return bufferLength + headerSize;
  }

  protected void updateConsumingOffset() throws IOException {
    while (currentPartitionRemainingBytes == 0
        && (currentDataRegion < numRegions - 1 || numRemainingPartitions > 0)) {
      if (numRemainingPartitions <= 0) {
        ++currentDataRegion;
        numRemainingPartitions = endPartitionIndex - startPartitionIndex + 1;

        // read the target index entry to the target index buffer
        indexFileChannel.position(
            currentDataRegion * getIndexRegionSize()
                + (long) startPartitionIndex * INDEX_ENTRY_SIZE);
        readHeaderOrIndexBuffer(indexFileChannel, indexBuffer, indexBuffer.capacity());
      }

      // get the data file offset and the data size
      dataConsumingOffset = indexBuffer.getLong();
      currentPartitionRemainingBytes = indexBuffer.getLong();
      --numRemainingPartitions;

      logger.debug(
          "readBuffer updateConsumingOffset, {},  {}, {}, {}",
          streamId,
          dataFileChannelSize,
          dataConsumingOffset,
          currentPartitionRemainingBytes);

      // if these checks fail, the partition file must be corrupted
      if (dataConsumingOffset < 0
          || dataConsumingOffset + currentPartitionRemainingBytes > dataFileChannelSize
          || currentPartitionRemainingBytes < 0) {
        throw new FileCorruptedException("File " + fileInfo.getFilePath() + " is corrupted");
      }
    }
  }

  private synchronized boolean readBuffer(ByteBuf buffer) throws IOException {
    try {
      dataFileChannel.position(dataConsumingOffset);

      int readSize =
          readBuffer(
              fileInfo.getFilePath(),
              dataFileChannel,
              headerBuffer,
              buffer,
              headerBuffer.capacity());
      currentPartitionRemainingBytes -= readSize;
      dataConsumingOffset = dataFileChannel.position();

      logger.debug(
          "readBuffer data: {}, {}, {}, {}, {}, {}",
          streamId,
          currentPartitionRemainingBytes,
          readSize,
          dataConsumingOffset,
          fileInfo.getFilePath(),
          System.identityHashCode(buffer));

      // if this check fails, the partition file must be corrupted
      if (currentPartitionRemainingBytes < 0) {
        throw new FileCorruptedException("File is corrupted");
      } else if (currentPartitionRemainingBytes == 0) {
        logger.debug(
            "readBuffer end, {},  {}, {}, {}",
            streamId,
            dataFileChannelSize,
            dataConsumingOffset,
            currentPartitionRemainingBytes);
        int prevDataRegion = currentDataRegion;
        updateConsumingOffset();
        return prevDataRegion == currentDataRegion && currentPartitionRemainingBytes > 0;
      }

      logger.debug(
          "readBuffer run: {}, {}, {}, {}",
          streamId,
          dataFileChannelSize,
          dataConsumingOffset,
          currentPartitionRemainingBytes);
      return true;
    } catch (Throwable throwable) {
      logger.error("Failed to read partition file.", throwable);
      synchronized (lock) {
        isReleased = true;
      }
      throw throwable;
    }
  }

  public boolean hasRemaining() {
    return currentPartitionRemainingBytes > 0;
  }

  protected void notifyBacklog(int backlog) {
    logger.debug("stream manager stream id {} backlog:{}", streamId, backlog);
    associatedChannel
        .writeAndFlush(new BacklogAnnouncement(streamId, backlog))
        .addListener(
            future -> {
              if (!future.isSuccess()) {
                logger.error("send backlog {} to stream {} failed", backlog, streamId);
              }
            });
  }

  private void notifyError(Throwable throwable) {
    logger.error(
        "Read file: {} error from {}, stream id {}",
        fileInfo.getFilePath(),
        NettyUtils.getRemoteAddress(this.associatedChannel),
        streamId,
        throwable);
    if (throwable instanceof ClosedChannelException) {
      return;
    }

    if (this.associatedChannel.isActive()) {
      // If a stream is failed, send exceptions with the best effort, do not expect response.
      // And do not close channel because multiple streams are using the very same channel.
      // wrapIOException to PartitionUnRetryAbleException, client may choose regenerate the data.
      this.associatedChannel.writeAndFlush(
          new TransportableError(streamId, ExceptionUtils.wrapIOExceptionToUnRetryable(throwable)));
    }
  }

  public long getPriority() {
    return dataConsumingOffset;
  }

  @Override
  public int compareTo(MapPartitionDataReader that) {
    return Long.compare(getPriority(), that.getPriority());
  }

  public FileInfo getFileInfo() {
    return fileInfo;
  }

  public void closeReader() {
    synchronized (lock) {
      readFinished = true;
    }

    logger.debug("Closed read for stream {}", this.streamId);
  }

  public void recycle() {
    synchronized (lock) {
      if (!isReleased) {
        release();
        recycleStream.run();
      }
    }
  }

  public void recycleOnError(Throwable throwable) {
    synchronized (lock) {
      if (!errorNotified) {
        errorNotified = true;
        errorCause = throwable;
        notifyError(throwable);
        recycle();
      }
    }
  }

  public void release() {
    // we can safely release if reader reaches error or (read/send finished)
    synchronized (lock) {
      if (!isReleased) {
        logger.debug("release reader for stream {}", streamId);
        // old client can't support BufferStreamEnd, so for new client it tells client that this
        // stream is finished.
        if (fileInfo.isPartitionSplitEnabled() && !errorNotified) {
          associatedChannel.writeAndFlush(
              new RpcRequest(
                  TransportClient.requestId(),
                  new NioManagedBuffer(
                      new TransportMessage(
                              MessageType.BUFFER_STREAM_END,
                              PbBufferStreamEnd.newBuilder()
                                  .setStreamType(StreamType.CreditStream)
                                  .setStreamId(streamId)
                                  .build()
                                  .toByteArray())
                          .toByteBuffer())));
        }
        if (!buffersToSend.isEmpty()) {
          int dataBufferInUse = 0;
          RecyclableBuffer buffer;
          while ((buffer = buffersToSend.poll()) != null) {
            if (buffer.isDataBuffer()) {
              dataBufferInUse++;
            }
            buffer.recycle();
          }
          numInUseBuffers.addAndGet(-1 * dataBufferInUse);
        }
        isReleased = true;
      }
    }
  }

  public boolean isFinished() {
    synchronized (lock) {
      // ensure every buffer are return to bufferQueue or release in buffersRead
      return numInUseBuffers.get() == 0 && isReleased;
    }
  }

  @Override
  public String toString() {
    final StringBuilder sb = new StringBuilder("DataPartitionReader{");
    sb.append("startPartitionIndex=").append(startPartitionIndex);
    sb.append(", endPartitionIndex=").append(endPartitionIndex);
    sb.append(", streamId=").append(streamId);
    sb.append('}');
    return sb.toString();
  }

  public long getStreamId() {
    return streamId;
  }

  @VisibleForTesting
  public AtomicInteger getNumInUseBuffers() {
    return numInUseBuffers;
  }

  public boolean shouldReadData() {
    synchronized (lock) {
      return !isReleased && !readFinished;
    }
  }
}
