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

import java.io.IOException;
import java.util.ArrayList;
import java.util.Objects;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.atomic.LongAdder;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Consumer;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.util.ThreadUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.common.write.PushState;

public class DataPusher {
  private static final Logger logger = LoggerFactory.getLogger(DataPusher.class);

  private final long WAIT_TIME_NANOS = TimeUnit.MILLISECONDS.toNanos(500);

  private LinkedBlockingQueue<PushTask> idleQueue;
  // partition -> PushTask Queue
  private final PushState pushState;
  private final DataPushQueue dataPushQueue;
  private final ReentrantLock idleLock = new ReentrantLock();
  private final Condition idleFull = idleLock.newCondition();

  private final AtomicReference<IOException> exceptionRef = new AtomicReference<>();

  private final int shuffleId;
  private final int mapId;
  private final int attemptId;
  private final int numMappers;
  private final int numPartitions;
  private final ShuffleClient client;
  private final Consumer<Integer> afterPush;

  private volatile boolean terminated;
  private final LongAdder[] mapStatusLengths;
  private Thread pushThread;

  public DataPusher(
      int shuffleId,
      int mapId,
      int attemptId,
      long taskId,
      int numMappers,
      int numPartitions,
      CelebornConf conf,
      ShuffleClient client,
      LinkedBlockingQueue<PushTask> pushTasks,
      Consumer<Integer> afterPush,
      LongAdder[] mapStatusLengths)
      throws InterruptedException {
    final int pushQueueCapacity = conf.clientPushQueueCapacity();
    final int pushBufferMaxSize = conf.clientPushBufferMaxSize();

    if (pushTasks == null) {
      idleQueue = new LinkedBlockingQueue<>(pushQueueCapacity);
    } else {
      idleQueue = pushTasks;
    }
    dataPushQueue =
        new DataPushQueue(
            conf, this, client, shuffleId, mapId, attemptId, numMappers, numPartitions);

    for (int i = idleQueue.size(); i < pushQueueCapacity; i++) {
      idleQueue.put(new PushTask(pushBufferMaxSize));
    }

    this.shuffleId = shuffleId;
    this.mapId = mapId;
    this.attemptId = attemptId;
    this.numMappers = numMappers;
    this.numPartitions = numPartitions;
    this.client = client;
    this.afterPush = afterPush;
    this.mapStatusLengths = mapStatusLengths;
    final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
    this.pushState = client.getPushState(mapKey);

    pushThread =
        ThreadUtils.newDaemonThread(
            new Runnable() {
              private void reclaimTask(PushTask task) throws InterruptedException {
                idleLock.lockInterruptibly();
                try {
                  idleQueue.put(task);
                  if (idleQueue.remainingCapacity() == 0) {
                    idleFull.signal();
                  }
                } catch (InterruptedException e) {
                  logger.error("DataPusher thread interrupted while reclaiming data.");
                  throw e;
                } finally {
                  idleLock.unlock();
                }
              }

              @Override
              public void run() {
                while (stillRunning()) {
                  try {
                    ArrayList<PushTask> tasks = dataPushQueue.takePushTasks();
                    for (int i = 0; i < tasks.size(); i++) {
                      PushTask task = tasks.get(i);
                      pushData(task);
                      reclaimTask(task);
                    }
                  } catch (CelebornIOException e) {
                    exceptionRef.set(e);
                  } catch (IOException e) {
                    exceptionRef.set(new CelebornIOException(e));
                  } catch (InterruptedException e) {
                    logger.error("DataPusher push thread interrupted while pushing data.");
                    break;
                  }
                }
              }
            },
            "celeborn-client-data-pusher-" + taskId);

    pushThread.start();
  }

  public void addTask(int partitionId, byte[] buffer, int size)
      throws IOException, InterruptedException {
    try {
      PushTask task = null;
      while (task == null) {
        checkException();
        task = idleQueue.poll(WAIT_TIME_NANOS, TimeUnit.NANOSECONDS);
      }
      task.setSize(size);
      task.setPartitionId(partitionId);
      System.arraycopy(buffer, 0, task.getBuffer(), 0, size);
      while (!dataPushQueue.addPushTask(task)) {
        checkException();
      }
    } catch (InterruptedException e) {
      logger.error("DataPusher thread interrupted while adding push task.");
      pushThread.interrupt();
      throw e;
    }
  }

  public void waitOnTermination() throws IOException, InterruptedException {
    try {
      idleLock.lockInterruptibly();
      waitIdleQueueFullWithLock();
    } catch (InterruptedException e) {
      logger.error("DataPusher thread interrupted while waitOnTermination.");
      pushThread.interrupt();
      throw e;
    }

    terminated = true;
    try {
      pushThread.join();
    } catch (InterruptedException e) {
      logger.info("Thread interrupted while joining pushThread");
      Thread.currentThread().interrupt();
      throw e;
    }
    dataPushQueue.clear();
    checkException();
  }

  public void setException(IOException ie) {
    exceptionRef.compareAndSet(null, ie);
  }

  public void checkException() throws IOException {
    if (exceptionRef.get() != null) {
      throw exceptionRef.get();
    }
    if (pushState.exception.get() != null) {
      throw pushState.exception.get();
    }
  }

  protected void pushData(PushTask task) throws IOException {
    int bytesWritten =
        client.pushData(
            shuffleId,
            mapId,
            attemptId,
            task.getPartitionId(),
            task.getBuffer(),
            0,
            task.getSize(),
            numMappers,
            numPartitions);
    afterPush.accept(bytesWritten);
    mapStatusLengths[task.getPartitionId()].add(bytesWritten);
  }

  private void waitIdleQueueFullWithLock() throws InterruptedException {
    try {
      while (idleQueue != null
          && idleQueue.remainingCapacity() > 0
          && exceptionRef.get() == null
          && pushState.exception.get() == null
          && (pushThread != null && pushThread.isAlive())) {
        idleFull.await(WAIT_TIME_NANOS, TimeUnit.NANOSECONDS);
      }
    } catch (InterruptedException e) {
      logger.error("Thread interrupted while waitIdleQueueFullWithLock.", e);
      throw e;
    } finally {
      idleLock.unlock();
    }
  }

  protected boolean stillRunning() {
    return !terminated
        && !Objects.nonNull(exceptionRef.get())
        && !Objects.nonNull(pushState.exception.get());
  }

  public DataPushQueue getDataPushQueue() {
    return dataPushQueue;
  }

  public LinkedBlockingQueue<PushTask> getAndResetIdleQueue() {
    LinkedBlockingQueue<PushTask> queue = idleQueue;
    idleQueue = null;
    return queue;
  }
}
