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

import static org.apache.celeborn.plugin.flink.utils.Utils.checkArgument;
import static org.apache.celeborn.plugin.flink.utils.Utils.checkState;

import java.nio.ByteBuffer;
import java.util.*;

import javax.annotation.concurrent.GuardedBy;

import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** A buffer pool which will dispatch buffers to all {@link CreditListener}s. */
public class TransferBufferPool implements BufferRecycler {
  private static Logger logger = LoggerFactory.getLogger(TransferBufferPool.class);

  private static final int MIN_CREDITS_TO_NOTIFY = 2;

  private final Object lock = new Object();

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

  @GuardedBy("lock")
  private final Queue<CreditListener> listeners = new ArrayDeque<>();

  @GuardedBy("lock")
  private int numAvailableBuffers;

  @GuardedBy("lock")
  private boolean isDestroyed;

  public TransferBufferPool(Collection<ByteBuf> initialBuffers) {
    synchronized (lock) {
      buffers.addAll(initialBuffers);
      numAvailableBuffers += initialBuffers.size();
    }
  }

  /** Requests a data transmitting unit. */
  public ByteBuf requestBuffer() {
    synchronized (lock) {
      checkState(!isDestroyed, "Buffer pool has been destroyed.");

      return buffers.poll();
    }
  }

  /** Adds an available buffer to this buffer pool. */
  public void addBuffers(List<? extends ByteBuf> byteBufs) {
    List<CreditAssignment> creditAssignments;
    synchronized (lock) {
      if (isDestroyed) {
        byteBufs.forEach(ByteBuf::release);
        return;
      }

      buffers.addAll(byteBufs);
      numAvailableBuffers += byteBufs.size();
      creditAssignments = dispatchReservedCredits();
    }
    for (CreditAssignment creditAssignment : creditAssignments) {
      creditAssignment.getCreditListener().notifyAvailableCredits(creditAssignment.getNumCredits());
    }
  }

  /** Tries to reserve buffers for the target {@link CreditListener}. */
  public void reserveBuffers(CreditListener creditListener, int numRequiredBuffers) {
    int numCredits;
    CreditListener listener = null;
    synchronized (lock) {
      if (isDestroyed) {
        throw new IllegalStateException("Buffer pool has been destroyed.");
      }

      if (numRequiredBuffers > numAvailableBuffers) {
        creditListener.increaseNumCreditsNeeded(numRequiredBuffers - numAvailableBuffers);
      }

      if (!creditListener.isRegistered() && creditListener.getNumCreditsNeeded() > 0) {
        listeners.add(creditListener);
        creditListener.setRegistered(true);
      }

      numCredits = Math.min(numAvailableBuffers, numRequiredBuffers);
      if (numCredits > 0) {
        numAvailableBuffers -= numCredits;
        listener = creditListener;
      }

      logger.debug(
          "reserveBuffers: numCredits: {}, requiredBuffers: {}", numCredits, numRequiredBuffers);
    }
    if (listener != null) {
      listener.notifyAvailableCredits(numCredits);
    }
  }

  /** Returns the number of available buffers. */
  public int numBuffers() {
    synchronized (lock) {
      return buffers.size();
    }
  }

  /** Destroys buffer pool. */
  public void destroy() {
    synchronized (lock) {
      isDestroyed = true;
      listeners.clear();
      buffers.forEach(ByteBuf::release);
      buffers.clear();
    }
  }

  /** Returns true if this buffer pool has been destroyed. */
  public boolean isDestroyed() {
    synchronized (lock) {
      return isDestroyed;
    }
  }

  @Override
  public void recycle(ByteBuffer buffer) {
    List<CreditAssignment> creditAssignments;
    synchronized (lock) {
      // unmanaged memory no need to recycle, currently it is used only by tests
      if (isDestroyed) {
        return;
      }

      buffers.add(new Buffer(buffer, this, 0));
      ++numAvailableBuffers;
      creditAssignments = dispatchReservedCredits();
    }
    for (CreditAssignment creditAssignment : creditAssignments) {
      creditAssignment.getCreditListener().notifyAvailableCredits(creditAssignment.getNumCredits());
    }
  }

  @GuardedBy("lock")
  private int assignCredits(CreditListener creditListener) {
    assert Thread.holdsLock(lock);

    if (creditListener == null) {
      return 0;
    }

    int numCredits = Math.min(creditListener.getNumCreditsNeeded(), numAvailableBuffers);
    if (numCredits > 0) {
      creditListener.decreaseNumCreditsNeeded(numCredits);
      numAvailableBuffers -= numCredits;
    }

    if (creditListener.getNumCreditsNeeded() > 0) {
      listeners.add(creditListener);
    } else {
      creditListener.setRegistered(false);
    }
    return numCredits;
  }

  @GuardedBy("lock")
  private List<CreditAssignment> dispatchReservedCredits() {
    assert Thread.holdsLock(lock);

    List<CreditAssignment> creditAssignments = new ArrayList<>();
    if (numAvailableBuffers >= MIN_CREDITS_TO_NOTIFY && !listeners.isEmpty()) {
      while (numAvailableBuffers > 0 && !listeners.isEmpty()) {
        CreditListener creditListener = listeners.poll();
        int numCredits = assignCredits(creditListener);
        if (numCredits > 0) {
          creditAssignments.add(new CreditAssignment(numCredits, creditListener));
        }
      }
    }
    return creditAssignments;
  }

  private static class CreditAssignment {

    private final int numCredits;
    private final CreditListener creditListener;

    CreditAssignment(int numCredits, CreditListener creditListener) {
      checkArgument(numCredits > 0, "Must be positive.");
      checkArgument(creditListener != null, "Must be not null.");

      this.numCredits = numCredits;
      this.creditListener = creditListener;
    }

    public int getNumCredits() {
      return numCredits;
    }

    public CreditListener getCreditListener() {
      return creditListener;
    }
  }
}
