/*
 * 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.util.Random;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

import com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.meta.TimeWindow;
import org.apache.celeborn.common.network.buffer.ChunkBuffers;
import org.apache.celeborn.common.network.buffer.ManagedBuffer;
import org.apache.celeborn.common.util.JavaUtils;

/**
 * StreamManager which allows registration of an Iterator&lt;ManagedBuffer&gt;, which are
 * individually fetched as chunks by the client. Each registered buffer is one chunk.
 */
public class ChunkStreamManager {
  private static final Logger logger = LoggerFactory.getLogger(ChunkStreamManager.class);

  private final AtomicLong nextStreamId;
  // StreamId -> StreamState
  protected final ConcurrentHashMap<Long, StreamState> streams;
  // ShuffleKey -> StreamId
  protected final ConcurrentHashMap<String, Set<Long>> shuffleStreamIds;

  /** State of a single stream. */
  public static class StreamState {
    public final ChunkBuffers buffers;
    public final String shuffleKey;
    public final String fileName;
    public final TimeWindow fetchTimeMetric;
    // Used to keep track of the number of chunks being transferred and not finished yet.
    volatile long chunksBeingTransferred = 0L;

    StreamState(
        String shuffleKey, ChunkBuffers buffers, String fileName, TimeWindow fetchTimeMetric) {
      this.buffers = buffers;
      this.shuffleKey = shuffleKey;
      this.fileName = fileName;
      this.fetchTimeMetric = fetchTimeMetric;
    }
  }

  public ChunkStreamManager() {
    // For debugging purposes, start with a random stream id to help identifying different streams.
    // This does not need to be globally unique, only unique to this class.
    nextStreamId = new AtomicLong((long) new Random().nextInt(Integer.MAX_VALUE) * 1000);
    streams = JavaUtils.newConcurrentHashMap();
    shuffleStreamIds = JavaUtils.newConcurrentHashMap();
  }

  public ManagedBuffer getChunk(long streamId, int chunkIndex, int offset, int len) {
    StreamState state = streams.get(streamId);
    if (state == null) {
      throw new IllegalStateException(
          String.format(
              "Stream %s for chunk %s is not registered(Maybe removed).", streamId, chunkIndex));
    } else if (chunkIndex >= state.buffers.numChunks()) {
      throw new IllegalStateException(
          String.format("Requested chunk index beyond end %s", chunkIndex));
    }

    ChunkBuffers buffers = state.buffers;
    return buffers.chunk(chunkIndex, offset, len);
  }

  public TimeWindow getFetchTimeMetric(long streamId) {
    StreamState state = streams.get(streamId);
    if (state != null) {
      return state.fetchTimeMetric;
    } else {
      return null;
    }
  }

  public void chunkBeingSent(long streamId) {
    StreamState streamState = streams.get(streamId);
    if (streamState != null) {
      synchronized (streamState) {
        streamState.chunksBeingTransferred++;
      }
    }
  }

  public void chunkSent(long streamId) {
    StreamState streamState = streams.get(streamId);
    if (streamState != null) {
      synchronized (streamState) {
        streamState.chunksBeingTransferred--;
      }
    }
  }

  public long chunksBeingTransferred() {
    long sum = 0L;
    for (StreamState streamState : streams.values()) {
      sum += streamState.chunksBeingTransferred;
    }
    return sum;
  }

  /**
   * Registers a stream with shuffle key and disk file when reading local or dfs shuffle, which is
   * served to obtain disk file via registered stream id to close stream.
   *
   * <p>This stream could be reused again when other channel of the client is reconnected. If a
   * stream is not properly closed, it will eventually be cleaned up by `cleanupExpiredShuffleKey`.
   */
  public long registerStream(long streamId, String shuffleKey, String fileName) {
    return registerStream(streamId, shuffleKey, null, fileName, null);
  }

  /**
   * Registers a stream of ManagedBuffers which are served as individual chunks one at a time to
   * callers. Each ManagedBuffer will be release()'d after it is transferred on the wire. If a
   * client connection is closed before the iterator is fully drained, then the remaining buffers
   * will all be release()'d.
   *
   * <p>If an app ID is provided, only callers who've authenticated with the given app ID will be
   * allowed to fetch from this stream.
   *
   * <p>This stream could be reused again when other channel of the client is reconnected. If a
   * stream is not properly closed, it will eventually be cleaned up by `cleanupExpiredShuffleKey`.
   */
  public long registerStream(
      String shuffleKey, ChunkBuffers buffers, String fileName, TimeWindow fetchTimeMetric) {
    long myStreamId = nextStreamId.getAndIncrement();
    return registerStream(myStreamId, shuffleKey, buffers, fileName, fetchTimeMetric);
  }

  public long registerStream(
      long streamId,
      String shuffleKey,
      ChunkBuffers buffers,
      String fileName,
      TimeWindow fetchTimeMetric) {
    StreamState streamState = new StreamState(shuffleKey, buffers, fileName, fetchTimeMetric);
    streams.put(streamId, streamState);
    shuffleStreamIds.compute(
        shuffleKey,
        (key, value) -> {
          if (value == null) {
            value = ConcurrentHashMap.newKeySet();
          }
          value.add(streamId);
          return value;
        });

    return streamId;
  }

  public long nextStreamId() {
    return nextStreamId.getAndIncrement();
  }

  public void cleanupExpiredShuffleKey(Set<String> expiredShuffleKeys) {
    logger.info(
        "Clean up expired shuffle keys {}",
        String.join(",", expiredShuffleKeys.toArray(new String[0])));
    for (String expiredShuffleKey : expiredShuffleKeys) {
      Set<Long> expiredStreamIds = shuffleStreamIds.remove(expiredShuffleKey);

      // normally expiredStreamIds set will be empty as streamId will be removed when be fully read
      if (expiredStreamIds != null && !expiredStreamIds.isEmpty()) {
        expiredStreamIds.forEach(streams::remove);
      }
    }
    logger.info(
        "Cleaned up expired shuffle keys. The count of shuffle keys and streams: {}, {}",
        shuffleStreamIds.size(),
        streams.size());
  }

  public StreamState getStreamState(long streamId) {
    return streams.get(streamId);
  }

  public StreamState removeStreamState(long streamId) {
    return streams.remove(streamId);
  }

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

  @VisibleForTesting
  public long numShuffleSteams() {
    return shuffleStreamIds.values().stream().mapToLong(Set::size).sum();
  }
}
