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

import static org.apache.celeborn.common.protocol.PartitionLocation.Mode.PRIMARY;

import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

import scala.Tuple2;

import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.client.read.CelebornInputStream;
import org.apache.celeborn.client.read.MetricsCallback;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.network.client.TransportClientFactory;
import org.apache.celeborn.common.protocol.PartitionLocation;
import org.apache.celeborn.common.protocol.PbStreamHandler;
import org.apache.celeborn.common.rpc.RpcEndpointRef;
import org.apache.celeborn.common.util.ExceptionMaker;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.write.LocationPushFailedBatches;
import org.apache.celeborn.common.write.PushState;

public class DummyShuffleClient extends ShuffleClient {

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

  private final OutputStream os;
  private final CelebornConf conf;

  private final Map<Integer, ConcurrentHashMap<Integer, PartitionLocation>> reducePartitionMap =
      new HashMap<>();

  public AtomicInteger fetchFailureCount = new AtomicInteger();

  public DummyShuffleClient(CelebornConf conf, File file) throws Exception {
    this.os = new BufferedOutputStream(new FileOutputStream(file));
    this.conf = conf;
  }

  @Override
  public void setupLifecycleManagerRef(String host, int port) {}

  @Override
  public void setupLifecycleManagerRef(RpcEndpointRef endpointRef) {}

  @Override
  public void setExtension(byte[] extension) {}

  @Override
  public int pushData(
      int shuffleId,
      int mapId,
      int attemptId,
      int partitionId,
      byte[] data,
      int offset,
      int length,
      int numMappers,
      int numPartitions)
      throws IOException {
    os.write(data, offset, length);
    return length;
  }

  @Override
  public void prepareForMergeData(int shuffleId, int mapId, int attemptId) throws IOException {}

  @Override
  public int mergeData(
      int shuffleId,
      int mapId,
      int attemptId,
      int partitionId,
      byte[] data,
      int offset,
      int length,
      int numMappers,
      int numPartitions)
      throws IOException {
    os.write(data, offset, length);
    return length;
  }

  @Override
  public void pushMergedData(int shuffleId, int mapId, int attemptId) {}

  @Override
  public void mapperEnd(
      int shuffleId, int mapId, int attemptId, int numMappers, int numPartitions) {}

  @Override
  public void readReducerPartitionEnd(
      int shuffleId, int partitionId, int startMapIndex, int endMapIndex, int crc32, long bytes)
      throws IOException {}

  @Override
  public void mapPartitionMapperEnd(
      int shuffleId, int mapId, int attemptId, int numMappers, int numPartitions, int partitionId)
      throws IOException {}

  @Override
  public void cleanup(int shuffleId, int mapId, int attemptId) {}

  @Override
  public ShuffleClientImpl.ReduceFileGroups updateFileGroup(int shuffleId, int partitionId)
      throws CelebornIOException {
    return null;
  }

  @Override
  public boolean isShuffleStageEnd(int shuffleId) throws Exception {
    return true;
  }

  @Override
  public CelebornInputStream readPartition(
      int shuffleId,
      int appShuffleId,
      int partitionId,
      int attemptNumber,
      long taskId,
      int startMapIndex,
      int endMapIndex,
      ExceptionMaker exceptionMaker,
      ArrayList<PartitionLocation> locations,
      ArrayList<PbStreamHandler> streamHandlers,
      Map<String, LocationPushFailedBatches> failedBatchSetMap,
      Map<String, Pair<Integer, Integer>> chunksRange,
      int[] mapAttempts,
      MetricsCallback metricsCallback,
      boolean needDecompress)
      throws IOException {
    return null;
  }

  @Override
  public boolean cleanupShuffle(int shuffleId) {
    return false;
  }

  @Override
  public void shutdown() {
    try {
      os.close();
    } catch (IOException e) {
      LOG.error("Closing file failed.", e);
    }
  }

  @Override
  public PartitionLocation registerMapPartitionTask(
      int shuffleId, int numMappers, int mapId, int attemptId, int partitionId) {
    return null;
  }

  @Override
  public ConcurrentHashMap<Integer, PartitionLocation> getPartitionLocation(
      int shuffleId, int numMappers, int numPartitions) {
    return reducePartitionMap.get(shuffleId);
  }

  @Override
  public PushState getPushState(String mapKey) {
    return new PushState(conf);
  }

  @Override
  public Tuple2<Integer, Boolean> getShuffleId(
      int appShuffleId, String appShuffleIdentifier, boolean isWriter, boolean isBarrierStage) {
    return Tuple2.apply(appShuffleId, true);
  }

  @Override
  public boolean reportShuffleFetchFailure(int appShuffleId, int shuffleId, long taskId) {
    fetchFailureCount.incrementAndGet();
    return true;
  }

  @Override
  public boolean reportBarrierTaskFailure(int appShuffleId, String appShuffleIdentifier) {
    return true;
  }

  @Override
  public TransportClientFactory getDataClientFactory() {
    return null;
  }

  @Override
  public void excludeFailedFetchLocation(String hostAndFetchPort, Exception e) {}

  public void initReducePartitionMap(int shuffleId, int numPartitions, int workerNum) {
    ConcurrentHashMap<Integer, PartitionLocation> map = JavaUtils.newConcurrentHashMap();
    String host = "host";
    List<PartitionLocation> partitionLocationList = new ArrayList<>();
    for (int i = 0; i < workerNum; i++) {
      partitionLocationList.add(
          new PartitionLocation(0, 0, host, 1000 + i, 2000 + i, 3000 + i, 4000 + i, PRIMARY));
    }
    for (int i = 0; i < numPartitions; i++) {
      map.put(i, partitionLocationList.get(i % workerNum));
    }
    reducePartitionMap.put(shuffleId, map);
  }

  public Map<Integer, ConcurrentHashMap<Integer, PartitionLocation>> getReducePartitionMap() {
    return reducePartitionMap;
  }
}
