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

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicInteger;

import javax.annotation.Nullable;

import io.netty.buffer.ByteBuf;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

// Assume that max-managed memory for a MapPartitionData is (2^31 * buffersize)
public class BufferQueue {
  public static final Logger logger = LoggerFactory.getLogger(BufferQueue.class);

  private final Queue<ByteBuf> buffers = new ConcurrentLinkedQueue<>();

  private final MemoryManager memoryManager = MemoryManager.instance();

  /** Number of buffers occupied by this buffer queue (added but still not recycled). */
  private final AtomicInteger numBuffersOccupied = new AtomicInteger();

  private final AtomicInteger pendingRequestBuffers = new AtomicInteger();

  /** Whether this buffer queue is released or not. */
  private volatile boolean isReleased = false;

  private volatile int localBuffersTarget = 0;

  /** Returns the number of available buffers in this buffer queue. */
  public int size() {
    return buffers.size();
  }

  /**
   * Returns an available buffer from this buffer queue or returns null if no buffer is available
   * currently.
   */
  @Nullable
  public ByteBuf poll() {
    synchronized (buffers) {
      return buffers.poll();
    }
  }

  /**
   * Add buffers and increment numBufferOccupied. Free all buffers to global memory pool if this
   * buffer queue is released.
   */
  public void add(Collection<ByteBuf> availableBuffers) {
    synchronized (buffers) {
      if (!isReleased) {
        buffers.addAll(availableBuffers);
        numBuffersOccupied.addAndGet(availableBuffers.size());
        pendingRequestBuffers.addAndGet(-1 * availableBuffers.size());
      } else {
        for (ByteBuf availableBuffer : availableBuffers) {
          memoryManager.recycleReadBuffer(availableBuffer);
        }
      }
    }
  }

  public void recycle(ByteBuf buffer) {
    if (numBuffersOccupied.get() > localBuffersTarget) {
      recycleToGlobalPool(buffer);
    } else {
      recycleToLocalPool(buffer);
    }
  }

  public void recycleToGlobalPool(ByteBuf buffer) {
    numBuffersOccupied.decrementAndGet();
    memoryManager.recycleReadBuffer(buffer);
  }

  public void recycleToLocalPool(ByteBuf buffer) {
    buffer.clear();
    buffers.add(buffer);
  }

  // free unused buffer to the main pool if possible
  public void trim() {
    List<ByteBuf> buffersToFree = new ArrayList<>();
    synchronized (this) {
      while (numBuffersOccupied.get() > localBuffersTarget) {
        ByteBuf buffer = poll();
        if (buffer != null) {
          buffersToFree.add(buffer);
          numBuffersOccupied.decrementAndGet();
        } else {
          // there are no unused buffers here
          break;
        }
      }
    }

    if (!buffersToFree.isEmpty()) {
      buffersToFree.forEach(memoryManager::recycleReadBuffer);
    }
  }

  /**
   * Releases this buffer queue and recycles all available buffers. After released, no buffer can be
   * added to or polled from this buffer queue.
   */
  public void release() {
    synchronized (buffers) {
      isReleased = true;
      buffers.forEach(this::recycleToGlobalPool);
      buffers.clear();
    }
    pendingRequestBuffers.set(0);
    numBuffersOccupied.set(0);
  }

  /** Returns true is this buffer queue has been released. */
  public boolean isReleased() {
    return isReleased;
  }

  public int getLocalBuffersTarget() {
    return localBuffersTarget;
  }

  public void setLocalBuffersTarget(int localBuffersTarget) {
    this.localBuffersTarget = localBuffersTarget;
  }

  public void tryApplyNewBuffers(
      int readerSize, int bufferSize, ReadBufferListener readBufferListener) {
    if (readerSize != 0) {
      synchronized (this) {
        int occupiedSnapshot = numBuffersOccupied.get();
        int pendingSnapShot = pendingRequestBuffers.get();
        if (occupiedSnapshot + pendingSnapShot < localBuffersTarget) {
          int newBuffersCount = (localBuffersTarget - occupiedSnapshot - pendingSnapShot);
          logger.debug(
              "apply new buffers {} while current buffer queue size {} with read count {}",
              newBuffersCount,
              numBuffersOccupied.get(),
              readerSize);
          pendingRequestBuffers.addAndGet(newBuffersCount);
          memoryManager.requestReadBuffers(
              new ReadBufferRequest(newBuffersCount, bufferSize, readBufferListener));
        }
      }
    }
  }

  public boolean bufferAvailable() {
    return !buffers.isEmpty();
  }
}
