/*
 * 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.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;

import javax.annotation.concurrent.GuardedBy;

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

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.util.JavaUtils;
import org.apache.celeborn.common.util.ThreadUtils;
import org.apache.celeborn.service.deploy.worker.memory.MemoryManager;
import org.apache.celeborn.service.deploy.worker.storage.segment.SegmentMapPartitionData;

public class CreditStreamManager {
  private static final Logger logger = LoggerFactory.getLogger(CreditStreamManager.class);
  private final AtomicLong nextStreamId;
  private final ConcurrentHashMap<Long, StreamState> streams;
  private final ConcurrentHashMap<FileInfo, MapPartitionData> activeMapPartitions;
  private final ConcurrentHashMap<String, ExecutorService> storageFetcherPool =
      JavaUtils.newConcurrentHashMap();
  private int minReadBuffers;
  private int maxReadBuffers;
  private int threadsPerMountPoint;
  private int minBuffersToTriggerRead;
  private final BlockingQueue<DelayedStreamId> recycleStreamIds = new DelayQueue<>();

  @GuardedBy("lock")
  private volatile ExecutorService recycleThread;

  private final Object lock = new Object();

  public CreditStreamManager(
      int minReadBuffers,
      int maxReadBuffers,
      int threadsPerMountPoint,
      int minBuffersToTriggerRead) {
    nextStreamId = new AtomicLong((long) new Random().nextInt(Integer.MAX_VALUE) * 1000);
    streams = JavaUtils.newConcurrentHashMap();
    activeMapPartitions = JavaUtils.newConcurrentHashMap();
    this.minReadBuffers = minReadBuffers;
    this.maxReadBuffers = maxReadBuffers;
    this.threadsPerMountPoint = threadsPerMountPoint;
    this.minBuffersToTriggerRead = minBuffersToTriggerRead;
    MemoryManager.instance().setCreditStreamManager(this);
    logger.debug(
        "Initialize buffer stream manager with {} {} {}",
        this.minReadBuffers,
        this.maxReadBuffers,
        threadsPerMountPoint);
  }

  public long registerStream(
      Consumer<Long> notifyStreamHandlerCallback,
      Channel channel,
      String shuffleKey,
      int initialCredit,
      int startSubIndex,
      int endSubIndex,
      DiskFileInfo fileInfo)
      throws IOException {
    long streamId = nextStreamId.getAndIncrement();
    logger.debug(
        "Register stream start from {}, streamId: {}, fileInfo: {}",
        channel.remoteAddress(),
        streamId,
        fileInfo);

    AtomicReference<IOException> exception = new AtomicReference<>();
    MapPartitionData mapPartitionData =
        activeMapPartitions.compute(
            fileInfo,
            (k, v) -> {
              if (v == null) {
                try {
                  MapFileMeta fileMeta = (MapFileMeta) fileInfo.getFileMeta();
                  v =
                      fileMeta.isSegmentGranularityVisible()
                          ? new SegmentMapPartitionData(
                              minReadBuffers,
                              maxReadBuffers,
                              storageFetcherPool,
                              threadsPerMountPoint,
                              fileInfo,
                              id -> recycleStream(id),
                              minBuffersToTriggerRead)
                          : new MapPartitionData(
                              minReadBuffers,
                              maxReadBuffers,
                              storageFetcherPool,
                              threadsPerMountPoint,
                              fileInfo,
                              id -> recycleStream(id),
                              minBuffersToTriggerRead);
                } catch (IOException e) {
                  exception.set(e);
                  return null;
                }
              }
              initializeStreamStateAndPartitionReader(
                  channel, shuffleKey, startSubIndex, endSubIndex, fileInfo, streamId, v);
              return v;
            });
    if (exception.get() != null) {
      throw exception.get();
    }
    mapPartitionData.tryRequestBufferOrRead();

    notifyStreamHandlerCallback.accept(streamId);
    addCredit(initialCredit, streamId);

    logger.debug("Register stream streamId: {}, fileInfo: {}", streamId, fileInfo);

    return streamId;
  }

  private void initializeStreamStateAndPartitionReader(
      Channel channel,
      String shuffleKey,
      int startSubIndex,
      int endSubIndex,
      FileInfo fileInfo,
      long streamId,
      MapPartitionData mapPartitionData) {
    StreamState streamState =
        new StreamState(
            channel,
            shuffleKey,
            ((MapFileMeta) fileInfo.getFileMeta()).getBufferSize(),
            mapPartitionData);
    streams.put(streamId, streamState);
    mapPartitionData.setupDataPartitionReader(startSubIndex, endSubIndex, streamId, channel);
  }

  private void addCredit(MapPartitionData mapPartitionData, int numCredit, long streamId) {
    logger.debug("streamId: {}, add credit: {}", streamId, numCredit);
    try {
      if (mapPartitionData != null && numCredit > 0) {
        mapPartitionData.addReaderCredit(numCredit, streamId);
      }
    } catch (Throwable e) {
      logger.error("streamId: {}, failed to add credit: {}", streamId, numCredit, e);
    }
  }

  private void notifyRequiredSegment(
      MapPartitionData mapPartitionData, int requiredSegmentId, long streamId, int subPartitionId) {
    logger.debug(
        "Receive RequiredSegment from client, streamId: {}, requiredSegmentId: {}, subPartitionId: {}",
        streamId,
        requiredSegmentId,
        subPartitionId);
    try {
      if (mapPartitionData instanceof SegmentMapPartitionData) {
        ((SegmentMapPartitionData) mapPartitionData)
            .notifyRequiredSegmentId(requiredSegmentId, streamId, subPartitionId);
      } else {
        logger.warn("Only non-null SegmentMapPartitionData is expected for notifyRequiredSegment.");
      }
    } catch (Throwable e) {
      logger.error("Failed to notify segmentId {} for stream {}.", requiredSegmentId, streamId, e);
      throw e;
    }
  }

  public void addCredit(int numCredit, long streamId) {
    StreamState streamState = streams.get(streamId);
    if (streamState == null) {
      // In flink hybrid shuffle integration strategy, the stream may release in worker before
      // client receive bufferStreamEnd,
      // and the client may send request with old streamId, so ignore non-exist streams.
      logger.warn("Ignore AddCredit from stream {}, numCredit {}.", streamId, numCredit);
      return;
    }
    addCredit(streamState.getMapPartitionData(), numCredit, streamId);
  }

  public void notifyRequiredSegment(int requiredSegmentId, long streamId, int subPartitionId) {
    StreamState streamState = streams.get(streamId);
    if (streamState != null) {
      notifyRequiredSegment(
          streamState.getMapPartitionData(), requiredSegmentId, streamId, subPartitionId);
    } else {
      // In flink hybrid shuffle integration strategy, the stream may release in worker before
      // client receive bufferStreamEnd,
      // and the client may send request with old streamId, so ignore non-exist streams.
      logger.warn(
          "Ignore RequiredSegment from stream {}, subPartition {}, segmentId {}.",
          streamId,
          subPartitionId,
          requiredSegmentId);
    }
  }

  public void connectionTerminated(Channel channel) {
    for (Map.Entry<Long, StreamState> entry : streams.entrySet()) {
      if (entry.getValue().getAssociatedChannel() == channel) {
        logger.info("connection closed, clean streamId: {}", entry.getKey());
        recycleStream(entry.getKey());
      }
    }
  }

  public void notifyStreamEndByClient(long streamId) {
    recycleStream(streamId);
  }

  public void recycleStream(long streamId) {
    recycleStreamIds.add(new DelayedStreamId(streamId));
    startRecycleThread(); // lazy start thread
  }

  @VisibleForTesting
  public int numRecycleStreams() {
    return recycleStreamIds.size();
  }

  @VisibleForTesting
  public ConcurrentHashMap<Long, StreamState> getStreams() {
    return streams;
  }

  public String getStreamShuffleKey(Long streamId) {
    StreamState streamState = streams.get(streamId);
    return streamState == null ? null : streamState.getShuffleKey();
  }

  private void startRecycleThread() {
    synchronized (lock) {
      if (recycleThread == null) {
        recycleThread =
            ThreadUtils.newDaemonSingleThreadExecutor("worker-credit-stream-manager-recycler");
        recycleThread.submit(
            () -> {
              while (true) {
                try {
                  DelayedStreamId delayedStreamId = recycleStreamIds.take();
                  cleanResource(delayedStreamId.streamId);
                } catch (Throwable e) {
                  logger.warn(e.getMessage(), e);
                }
              }
            });

        logger.info("start stream recycle thread");
      }
    }
  }

  public void cleanResource(Long streamId) {
    logger.debug("received clean stream: {}", streamId);
    StreamState streamState = streams.get(streamId);
    if (streamState != null) {
      MapPartitionData mapPartitionData = streamState.getMapPartitionData();
      if (mapPartitionData != null) {
        if (mapPartitionData.releaseReader(streamId)) {
          streams.remove(streamId);
          if (mapPartitionData.getReaders().isEmpty()) {
            FileInfo fileInfo = mapPartitionData.getDiskFileInfo();
            activeMapPartitions.computeIfPresent(
                fileInfo,
                (k, v) -> {
                  if (v.getReaders().isEmpty()) {
                    v.close();
                    return null;
                  }
                  return v;
                });
          }
        } else {
          logger.debug("retry clean stream: {}", streamId);
          recycleStreamIds.add(new DelayedStreamId(streamId));
        }
      }
    }
  }

  public long getStreamsCount() {
    return streams.size();
  }

  public int getActiveMapPartitionCount() {
    return activeMapPartitions.size();
  }

  protected static class StreamState {
    private Channel associatedChannel;
    private String shuffleKey;
    private int bufferSize;
    private MapPartitionData mapPartitionData;

    public StreamState(
        Channel associatedChannel,
        String shuffleKey,
        int bufferSize,
        MapPartitionData mapPartitionData) {
      this.associatedChannel = associatedChannel;
      this.shuffleKey = shuffleKey;
      this.bufferSize = bufferSize;
      this.mapPartitionData = mapPartitionData;
    }

    public Channel getAssociatedChannel() {
      return associatedChannel;
    }

    public String getShuffleKey() {
      return shuffleKey;
    }

    public int getBufferSize() {
      return bufferSize;
    }

    public MapPartitionData getMapPartitionData() {
      return mapPartitionData;
    }
  }

  public static class DelayedStreamId implements Delayed {
    private static final long delayTime = 100; // 100ms
    private long createMillis = System.currentTimeMillis();

    private long streamId;

    public DelayedStreamId(long streamId) {
      this.createMillis = createMillis + delayTime;
      this.streamId = streamId;
    }

    @Override
    public long getDelay(TimeUnit unit) {
      long diff = createMillis - System.currentTimeMillis();
      return unit.convert(diff, TimeUnit.MILLISECONDS);
    }

    public long getCreateMillis() {
      return createMillis;
    }

    @Override
    public int compareTo(Delayed o) {
      long otherCreateMillis = ((DelayedStreamId) o).getCreateMillis();
      if (this.createMillis < otherCreateMillis) {
        return -1;
      } else if (this.createMillis > otherCreateMillis) {
        return 1;
      }

      return 0;
    }

    @Override
    public String toString() {
      return "DelayedStreamId{" + "createMillis=" + createMillis + ", streamId=" + streamId + '}';
    }
  }
}
