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

import static org.apache.celeborn.plugin.flink.metric.RemoteShuffleMetricFactory.createShuffleIOOwnerMetricGroup;
import static org.apache.celeborn.plugin.flink.utils.Utils.checkNotNull;
import static org.apache.celeborn.plugin.flink.utils.Utils.checkState;
import static org.apache.flink.runtime.io.network.metrics.NettyShuffleMetricFactory.METRIC_GROUP_INPUT;
import static org.apache.flink.runtime.io.network.metrics.NettyShuffleMetricFactory.METRIC_GROUP_OUTPUT;
import static org.apache.flink.runtime.metrics.groups.ShuffleIOMetricGroup.createShuffleIOMetricGroup;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
import org.apache.flink.runtime.deployment.ResultPartitionDeploymentDescriptor;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
import org.apache.flink.runtime.executiongraph.PartitionInfo;
import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
import org.apache.flink.runtime.io.network.metrics.InputChannelMetrics;
import org.apache.flink.runtime.io.network.partition.PartitionProducerStateProvider;
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.io.network.partition.ResultPartitionManager;
import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate;
import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
import org.apache.flink.runtime.jobgraph.IntermediateResultPartitionID;
import org.apache.flink.runtime.metrics.groups.ShuffleIOMetricGroup;
import org.apache.flink.runtime.shuffle.ShuffleEnvironment;
import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
import org.apache.flink.util.FlinkRuntimeException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.plugin.flink.netty.NettyShuffleEnvironmentWrapper;

/**
 * The implementation of {@link ShuffleEnvironment} based on the remote shuffle service, providing
 * shuffle environment on flink TM side.
 */
public class RemoteShuffleEnvironment
    implements ShuffleEnvironment<ResultPartitionWriter, IndexedInputGate> {

  private static final Logger LOG = LoggerFactory.getLogger(RemoteShuffleEnvironment.class);

  /** Network buffer pool for shuffle read and shuffle write. */
  private final NetworkBufferPool networkBufferPool;

  /** A trivial {@link ResultPartitionManager}. */
  private final ResultPartitionManager resultPartitionManager;

  /** Factory class to create {@link AbstractRemoteShuffleResultPartition}. */
  private final AbstractRemoteShuffleResultPartitionFactory resultPartitionFactory;

  private final AbstractRemoteShuffleInputGateFactory inputGateFactory;

  private final CelebornConf conf;

  private final NettyShuffleEnvironmentWrapper shuffleEnvironmentWrapper;

  /** Whether the shuffle environment is closed. */
  private boolean isClosed;

  private final Object lock = new Object();

  private final ConcurrentHashMap.KeySetView<IntermediateDataSetID, Boolean> nettyResultIds =
      ConcurrentHashMap.newKeySet();

  private final ConcurrentHashMap.KeySetView<IntermediateResultPartitionID, Boolean>
      nettyResultPartitionIds = ConcurrentHashMap.newKeySet();

  private final Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups =
      JavaUtils.newConcurrentHashMap();

  /**
   * @param networkBufferPool Network buffer pool for shuffle read and shuffle write.
   * @param resultPartitionManager A trivial {@link ResultPartitionManager}.
   */
  public RemoteShuffleEnvironment(
      NetworkBufferPool networkBufferPool,
      ResultPartitionManager resultPartitionManager,
      AbstractRemoteShuffleResultPartitionFactory resultPartitionFactory,
      AbstractRemoteShuffleInputGateFactory inputGateFactory,
      CelebornConf conf,
      NettyShuffleEnvironmentWrapper shuffleEnvironmentWrapper) {
    this.networkBufferPool = networkBufferPool;
    this.resultPartitionManager = resultPartitionManager;
    this.resultPartitionFactory = resultPartitionFactory;
    this.inputGateFactory = inputGateFactory;
    this.conf = conf;
    this.shuffleEnvironmentWrapper = shuffleEnvironmentWrapper;
    this.isClosed = false;
  }

  @Override
  public void close() {
    LOG.info("Close RemoteShuffleEnvironment.");
    synchronized (lock) {
      nettyResultIds.clear();
      nettyResultPartitionIds.clear();
      shuffleIOMetricGroups.clear();
      try {
        networkBufferPool.destroyAllBufferPools();
      } catch (Throwable t) {
        LOG.error("Close RemoteShuffleEnvironment failure.", t);
      }
      try {
        resultPartitionManager.shutdown();
      } catch (Throwable t) {
        LOG.error("Close RemoteShuffleEnvironment failure.", t);
      }
      try {
        networkBufferPool.destroy();
      } catch (Throwable t) {
        LOG.error("Close RemoteShuffleEnvironment failure.", t);
      }
      isClosed = true;
    }
  }

  @Override
  public int start() throws IOException {
    synchronized (lock) {
      checkState(!isClosed, "The RemoteShuffleEnvironment has already been shut down.");
      LOG.info("Starting the network environment and its components.");
      // trivial value.
      return 1;
    }
  }

  @Override
  public boolean updatePartitionInfo(ExecutionAttemptID consumerID, PartitionInfo partitionInfo) {
    throw new FlinkRuntimeException("Not implemented yet.");
  }

  @Override
  public ShuffleIOOwnerContext createShuffleIOOwnerContext(
      String ownerName, ExecutionAttemptID executionAttemptID, MetricGroup parentGroup) {
    MetricGroup remoteGroup = createShuffleIOOwnerMetricGroup(checkNotNull(parentGroup));
    return new ShuffleIOOwnerContext(
        checkNotNull(ownerName),
        checkNotNull(executionAttemptID),
        parentGroup,
        remoteGroup.addGroup(METRIC_GROUP_OUTPUT),
        remoteGroup.addGroup(METRIC_GROUP_INPUT));
  }

  @Override
  public Collection<ResultPartitionID> getPartitionsOccupyingLocalResources() {
    return resultPartitionManager.getUnreleasedPartitions();
  }

  @Override
  public List<ResultPartitionWriter> createResultPartitionWriters(
      ShuffleIOOwnerContext ownerContext,
      List<ResultPartitionDeploymentDescriptor> resultPartitionDeploymentDescriptors) {

    synchronized (lock) {
      checkState(!isClosed, "The RemoteShuffleEnvironment has already been shut down.");

      ResultPartitionWriter[] resultPartitions =
          new ResultPartitionWriter[resultPartitionDeploymentDescriptors.size()];
      for (int index = 0; index < resultPartitions.length; index++) {
        resultPartitions[index] =
            createResultPartitionWriterInternal(
                ownerContext, index, resultPartitionDeploymentDescriptors.get(index), conf);
      }
      return Arrays.asList(resultPartitions);
    }
  }

  private ResultPartitionWriter createResultPartitionWriterInternal(
      ShuffleIOOwnerContext ownerContext,
      int index,
      ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor,
      CelebornConf conf) {
    if (resultPartitionDeploymentDescriptor.getShuffleDescriptor()
        instanceof RemoteShuffleDescriptor) {
      int shuffleId =
          ((RemoteShuffleDescriptor) resultPartitionDeploymentDescriptor.getShuffleDescriptor())
              .getShuffleResource()
              .getMapPartitionShuffleDescriptor()
              .getShuffleId();
      ;
      return resultPartitionFactory.create(
          ownerContext.getOwnerName(),
          index,
          resultPartitionDeploymentDescriptor,
          conf,
          shuffleIOMetricGroups.computeIfAbsent(
              shuffleId, k -> createShuffleIOMetricGroup(ownerContext, shuffleId, conf)));
    } else {
      nettyResultIds.add(resultPartitionDeploymentDescriptor.getResultId());
      nettyResultPartitionIds.add(resultPartitionDeploymentDescriptor.getPartitionId());
      return shuffleEnvironmentWrapper
          .nettyResultPartitionFactory()
          .create(ownerContext.getOwnerName(), index, resultPartitionDeploymentDescriptor);
    }
  }

  @Override
  public void releasePartitionsLocally(Collection<ResultPartitionID> partitionIds) {
    List<ResultPartitionID> resultPartitionIds =
        partitionIds.stream()
            .filter(partitionId -> nettyResultPartitionIds.contains(partitionId.getPartitionId()))
            .collect(Collectors.toList());
    if (!resultPartitionIds.isEmpty()) {
      shuffleEnvironmentWrapper
          .nettyShuffleEnvironment()
          .releasePartitionsLocally(resultPartitionIds);
    }
  }

  @Override
  public List<IndexedInputGate> createInputGates(
      ShuffleIOOwnerContext ownerContext,
      PartitionProducerStateProvider producerStateProvider,
      List<InputGateDeploymentDescriptor> inputGateDescriptors) {
    synchronized (lock) {
      checkState(!isClosed, "The RemoteShuffleEnvironment has already been shut down.");

      InputChannelMetrics inputChannelMetrics =
          new InputChannelMetrics(ownerContext.getInputGroup(), ownerContext.getParentGroup());
      IndexedInputGate[] inputGates = new IndexedInputGate[inputGateDescriptors.size()];
      for (int gateIndex = 0; gateIndex < inputGates.length; gateIndex++) {
        InputGateDeploymentDescriptor igdd = inputGateDescriptors.get(gateIndex);
        IndexedInputGate inputGate =
            createInputGateInternal(
                ownerContext,
                producerStateProvider,
                gateIndex,
                igdd,
                inputChannelMetrics,
                shuffleIOMetricGroups);
        inputGates[gateIndex] = inputGate;
      }
      return Arrays.asList(inputGates);
    }
  }

  private IndexedInputGate createInputGateInternal(
      ShuffleIOOwnerContext ownerContext,
      PartitionProducerStateProvider producerStateProvider,
      int gateIndex,
      InputGateDeploymentDescriptor igdd,
      InputChannelMetrics inputChannelMetrics,
      Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
    return nettyResultIds.contains(igdd.getConsumedResultId())
        ? shuffleEnvironmentWrapper
            .nettyInputGateFactory()
            .create(ownerContext, gateIndex, igdd, producerStateProvider, inputChannelMetrics)
        : inputGateFactory.create(ownerContext, gateIndex, igdd, shuffleIOMetricGroups);
  }

  @VisibleForTesting
  AbstractRemoteShuffleResultPartitionFactory getResultPartitionFactory() {
    return resultPartitionFactory;
  }
}
