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

import java.io.IOException;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import io.netty.channel.Channel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.network.protocol.*;
import org.apache.celeborn.common.network.server.MessageHandler;
import org.apache.celeborn.common.network.util.NettyUtils;
import org.apache.celeborn.common.network.util.TransportConf;
import org.apache.celeborn.common.protocol.TransportModuleConstants;
import org.apache.celeborn.common.protocol.message.StatusCode;
import org.apache.celeborn.common.read.FetchRequestInfo;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.ThreadUtils;
import org.apache.celeborn.common.write.PushRequestInfo;

/**
 * Handler that processes server responses, in response to requests issued from a
 * [[TransportClient]]. It works by tracking the list of outstanding requests (and their callbacks).
 *
 * <p>Concurrency: thread safe and can be called from multiple threads.
 */
public class TransportResponseHandler extends MessageHandler<ResponseMessage> {
  private static final Logger logger = LoggerFactory.getLogger(TransportResponseHandler.class);

  private final TransportConf conf;
  private final Channel channel;

  private final ConcurrentHashMap<StreamChunkSlice, FetchRequestInfo> outstandingFetches;

  private final ConcurrentHashMap<Long, RpcResponseCallback> outstandingRpcs;
  private final ConcurrentHashMap<Long, PushRequestInfo> outstandingPushes;

  /** Records the time (in system nanoseconds) that the last fetch or RPC request was sent. */
  private final AtomicLong timeOfLastRequestNs;

  private final long pushTimeoutCheckerInterval;
  private static ScheduledExecutorService pushTimeoutChecker = null;
  private ScheduledFuture pushCheckerScheduleFuture;

  private final long fetchTimeoutCheckerInterval;
  private static ScheduledExecutorService fetchTimeoutChecker = null;
  private ScheduledFuture fetchCheckerScheduleFuture;

  @SuppressWarnings("StaticAssignmentInConstructor")
  public TransportResponseHandler(TransportConf conf, Channel channel) {
    this.conf = conf;
    this.channel = channel;
    this.outstandingFetches = JavaUtils.newConcurrentHashMap();
    this.outstandingRpcs = JavaUtils.newConcurrentHashMap();
    this.outstandingPushes = JavaUtils.newConcurrentHashMap();
    this.timeOfLastRequestNs = new AtomicLong(0);
    this.pushTimeoutCheckerInterval = conf.pushDataTimeoutCheckIntervalMs();
    this.fetchTimeoutCheckerInterval = conf.fetchDataTimeoutCheckIntervalMs();

    String module = conf.getModuleName();
    boolean checkPushTimeout = false;
    boolean checkFetchTimeout = false;
    if (TransportModuleConstants.DATA_MODULE.equals(module)) {
      checkPushTimeout = true;
      checkFetchTimeout = true;
    } else if (TransportModuleConstants.PUSH_MODULE.equals(module)
        || TransportModuleConstants.REPLICATE_MODULE.equals(module)) {
      checkPushTimeout = true;
    }
    synchronized (TransportResponseHandler.class) {
      if (checkPushTimeout) {
        if (pushTimeoutChecker == null) {
          pushTimeoutChecker =
              ThreadUtils.newDaemonThreadPoolScheduledExecutor(
                  "celeborn-push-timeout-checker", conf.pushDataTimeoutCheckerThreads());
        }
      }

      if (checkFetchTimeout) {
        if (fetchTimeoutChecker == null) {
          fetchTimeoutChecker =
              ThreadUtils.newDaemonThreadPoolScheduledExecutor(
                  "celeborn-fetch-timeout-checker", conf.fetchDataTimeoutCheckerThreads());
        }
      }
    }

    if (checkPushTimeout) {
      pushCheckerScheduleFuture =
          pushTimeoutChecker.scheduleWithFixedDelay(
              () -> failExpiredPushRequest(),
              pushTimeoutCheckerInterval,
              pushTimeoutCheckerInterval,
              TimeUnit.MILLISECONDS);
    }

    if (checkFetchTimeout) {
      fetchCheckerScheduleFuture =
          fetchTimeoutChecker.scheduleWithFixedDelay(
              () -> failExpiredFetchRequest(),
              fetchTimeoutCheckerInterval,
              fetchTimeoutCheckerInterval,
              TimeUnit.MILLISECONDS);
    }
  }

  public void failExpiredPushRequest() {
    long currentTime = System.currentTimeMillis();
    Iterator<Map.Entry<Long, PushRequestInfo>> iter = outstandingPushes.entrySet().iterator();
    while (iter.hasNext()) {
      Map.Entry<Long, PushRequestInfo> entry = iter.next();
      if (entry.getValue().dueTime <= currentTime) {
        PushRequestInfo info = removePushRequest(entry.getKey());
        if (info != null) {
          if (info.channelFuture != null) {
            info.channelFuture.cancel(true);
          }
          String module = conf.getModuleName();
          // When module name equals to DATA_MODULE, mean shuffle client push data, else means
          // do data replication.
          if (TransportModuleConstants.DATA_MODULE.equals(module)) {
            info.callback.onFailure(new CelebornIOException(StatusCode.PUSH_DATA_TIMEOUT_PRIMARY));
          } else if (TransportModuleConstants.PUSH_MODULE.equals(module)
              || TransportModuleConstants.REPLICATE_MODULE.equals(module)) {
            info.callback.onFailure(new CelebornIOException(StatusCode.PUSH_DATA_TIMEOUT_REPLICA));
          }
          info.channelFuture = null;
          info.callback = null;
        }
      }
    }
  }

  public void failExpiredFetchRequest() {
    long currentTime = System.currentTimeMillis();
    Iterator<Map.Entry<StreamChunkSlice, FetchRequestInfo>> iter =
        outstandingFetches.entrySet().iterator();
    while (iter.hasNext()) {
      Map.Entry<StreamChunkSlice, FetchRequestInfo> entry = iter.next();
      if (entry.getValue().dueTime <= currentTime) {
        FetchRequestInfo info = removeFetchRequest(entry.getKey());
        if (info != null) {
          if (info.channelFuture != null) {
            info.channelFuture.cancel(true);
          }
          logger.info(
              "Fail expire fetch request {},{},{},{}",
              entry.getKey().streamId,
              entry.getKey().chunkIndex,
              entry.getKey().offset,
              entry.getKey().len);
          info.callback.onFailure(
              entry.getKey().chunkIndex, new CelebornIOException(StatusCode.FETCH_DATA_TIMEOUT));
          info.channelFuture = null;
          info.callback = null;
        }
      }
    }
  }

  public void addFetchRequest(StreamChunkSlice streamChunkSlice, FetchRequestInfo info) {
    updateTimeOfLastRequest();
    if (outstandingFetches.containsKey(streamChunkSlice)) {
      logger.warn("[addFetchRequest] streamChunkSlice {} already exists!", streamChunkSlice);
    }
    outstandingFetches.put(streamChunkSlice, info);
  }

  public FetchRequestInfo removeFetchRequest(StreamChunkSlice streamChunkSlice) {
    return outstandingFetches.remove(streamChunkSlice);
  }

  public void addRpcRequest(long requestId, RpcResponseCallback callback) {
    updateTimeOfLastRequest();
    if (outstandingRpcs.containsKey(requestId)) {
      logger.warn("[addRpcRequest] requestId {} already exists!", requestId);
    }
    outstandingRpcs.put(requestId, callback);
  }

  public RpcResponseCallback removeRpcRequest(long requestId) {
    return outstandingRpcs.remove(requestId);
  }

  public void addPushRequest(long requestId, PushRequestInfo info) {
    updateTimeOfLastRequest();
    if (outstandingPushes.containsKey(requestId)) {
      logger.warn("[addPushRequest] requestId {} already exists!", requestId);
    }
    outstandingPushes.put(requestId, info);
  }

  public PushRequestInfo removePushRequest(long requestId) {
    return outstandingPushes.remove(requestId);
  }

  /**
   * Fire the failure callback for all outstanding requests. This is called when we have an uncaught
   * exception or pre-mature connection termination.
   */
  private void failOutstandingRequests(Throwable cause) {
    Iterator<StreamChunkSlice> fetchRequestIter = outstandingFetches.keySet().iterator();
    while (fetchRequestIter.hasNext()) {
      try {
        StreamChunkSlice slice = fetchRequestIter.next();
        FetchRequestInfo info = removeFetchRequest(slice);
        if (info != null) {
          info.callback.onFailure(slice.chunkIndex, cause);
        }
      } catch (Exception e) {
        logger.warn("ChunkReceivedCallback.onFailure throws exception", e);
      }
    }

    Iterator<Long> rpcRequestIter = outstandingRpcs.keySet().iterator();
    while (rpcRequestIter.hasNext()) {
      try {
        long requestId = rpcRequestIter.next();
        RpcResponseCallback callback = removeRpcRequest(requestId);
        if (callback != null) {
          callback.onFailure(cause);
        }
      } catch (Exception e) {
        logger.warn("RpcResponseCallback.onFailure throws exception", e);
      }
    }

    Iterator<Long> pushRequestIter = outstandingPushes.keySet().iterator();
    while (pushRequestIter.hasNext()) {
      try {
        long requestId = pushRequestIter.next();
        PushRequestInfo info = removePushRequest(requestId);
        if (info != null) {
          info.callback.onFailure(cause);
        }
      } catch (Exception e) {
        logger.warn("RpcResponseCallback.onFailure throws exception", e);
      }
    }
  }

  @Override
  public void channelActive() {}

  @Override
  public void channelInactive() {
    if (hasOutstandingRequests()) {
      // show the details of outstanding Fetches
      if (logger.isDebugEnabled()) {
        if (outstandingFetches.size() > 0) {
          for (Map.Entry<StreamChunkSlice, FetchRequestInfo> e : outstandingFetches.entrySet()) {
            StreamChunkSlice key = e.getKey();
            logger.debug("The channel is closed, but there is still outstanding Fetch {}", key);
          }
        } else {
          logger.debug("The channel is closed, the outstanding Fetches are empty");
        }
      }
      String remoteAddress = NettyUtils.getRemoteAddress(channel);
      logger.error(
          "Still have {} requests outstanding when connection from {} is closed",
          numOutstandingRequests(),
          remoteAddress);
      failOutstandingRequests(new IOException("Connection from " + remoteAddress + " closed"));
    }
    if (pushCheckerScheduleFuture != null) {
      pushCheckerScheduleFuture.cancel(false);
    }
    if (fetchCheckerScheduleFuture != null) {
      fetchCheckerScheduleFuture.cancel(false);
    }
  }

  @Override
  public void exceptionCaught(Throwable cause) {
    if (hasOutstandingRequests()) {
      String remoteAddress = NettyUtils.getRemoteAddress(channel);
      logger.error(
          "Still have {} requests outstanding when connection from {} is closed",
          numOutstandingRequests(),
          remoteAddress);
      failOutstandingRequests(cause);
    }
    if (pushCheckerScheduleFuture != null) {
      pushCheckerScheduleFuture.cancel(false);
    }
    if (fetchCheckerScheduleFuture != null) {
      fetchCheckerScheduleFuture.cancel(false);
    }
  }

  @Override
  public void handle(ResponseMessage message) throws Exception {
    if (message instanceof ChunkFetchSuccess) {
      ChunkFetchSuccess resp = (ChunkFetchSuccess) message;
      logger.debug("Chunk {} fetch succeeded", resp.streamChunkSlice);
      FetchRequestInfo info = removeFetchRequest(resp.streamChunkSlice);
      if (info == null) {
        logger.warn(
            "Ignoring response for block {} from {} since it is not outstanding",
            resp.streamChunkSlice,
            NettyUtils.getRemoteAddress(channel));
        resp.body().release();
      } else {
        try {
          info.callback.onSuccess(resp.streamChunkSlice.chunkIndex, resp.body());
        } finally {
          resp.body().release();
        }
      }
    } else if (message instanceof ChunkFetchFailure) {
      ChunkFetchFailure resp = (ChunkFetchFailure) message;
      logger.error(
          "chunk {} fetch failed, errorMessage {}", resp.streamChunkSlice, resp.errorString);
      FetchRequestInfo info = removeFetchRequest(resp.streamChunkSlice);
      if (info == null) {
        logger.warn(
            "Ignoring response for block {} from {} ({}) since it is not outstanding",
            resp.streamChunkSlice,
            NettyUtils.getRemoteAddress(channel),
            resp.errorString);
      } else {
        logger.warn("Receive ChunkFetchFailure, errorMsg {}", resp.errorString);
        info.callback.onFailure(
            resp.streamChunkSlice.chunkIndex,
            new ChunkFetchFailureException(
                "Failure while fetching " + resp.streamChunkSlice + ": " + resp.errorString));
      }
    } else if (message instanceof RpcResponse) {
      RpcResponse resp = (RpcResponse) message;
      PushRequestInfo info = removePushRequest(resp.requestId);
      if (info == null) {
        RpcResponseCallback listener = removeRpcRequest(resp.requestId);
        if (listener == null) {
          logger.warn(
              "Ignoring response for RPC {} from {} ({} bytes) since it is not outstanding",
              resp.requestId,
              NettyUtils.getRemoteAddress(channel),
              resp.body().size());
          resp.body().release();
        } else {
          try {
            listener.onSuccess(resp.body().nioByteBuffer());
          } finally {
            resp.body().release();
          }
        }
      } else {
        try {
          info.callback.onSuccess(resp.body().nioByteBuffer());
        } finally {
          resp.body().release();
        }
      }
    } else if (message instanceof RpcFailure) {
      RpcFailure resp = (RpcFailure) message;
      PushRequestInfo info = removePushRequest(resp.requestId);
      if (info == null) {
        RpcResponseCallback listener = removeRpcRequest(resp.requestId);
        if (listener == null) {
          logger.warn(
              "Ignoring response for RPC {} from {} ({}) since it is not outstanding",
              resp.requestId,
              NettyUtils.getRemoteAddress(channel),
              resp.errorString);
        } else {
          listener.onFailure(new IOException(resp.errorString));
        }
      } else {
        info.callback.onFailure(new CelebornIOException(resp.errorString));
      }
    } else {
      throw new IllegalStateException("Unknown response type: " + message.type());
    }
  }

  /** Returns total number of outstanding requests (fetch requests + rpcs) */
  public int numOutstandingRequests() {
    return outstandingFetches.size() + outstandingRpcs.size() + outstandingPushes.size();
  }

  /** Check if there are any outstanding requests (fetch requests + rpcs + push requests) */
  public Boolean hasOutstandingRequests() {
    return !outstandingFetches.isEmpty()
        || !outstandingRpcs.isEmpty()
        || !outstandingPushes.isEmpty();
  }

  /** Returns the time in nanoseconds of when the last request was sent out. */
  public long getTimeOfLastRequestNs() {
    return timeOfLastRequestNs.get();
  }

  /** Updates the time of the last request to the current system time. */
  public void updateTimeOfLastRequest() {
    timeOfLastRequestNs.set(System.nanoTime());
  }

  public void handleRpcFailure(long rpcRequestId, String errorMsg, Throwable cause) {
    RpcResponseCallback callback = removeRpcRequest(rpcRequestId);
    if (callback != null) {
      callback.onFailure(new CelebornIOException(errorMsg, cause));
    } else {
      logger.warn(
          "RpcResponseCallback {} not found/already addressed when listener handles rpc request failure",
          rpcRequestId);
    }
  }

  public void handlePushFailure(long pushRequestId, String errorMsg, Throwable cause) {
    PushRequestInfo info = removePushRequest(pushRequestId);
    if (info != null) {
      RpcResponseCallback callback = info.callback;
      if (callback != null) {
        callback.onFailure(new CelebornIOException(errorMsg, cause));
      } else {
        logger.warn(
            "PushRequestInfo {} callback is null when handle push request failure", pushRequestId);
      }
    } else {
      logger.warn(
          "PushRequestInfo {} not found/already addressed when listener handles push request failure",
          pushRequestId);
    }
  }

  public void handleFetchFailure(
      StreamChunkSlice streamChunkSlice, String errorMsg, Throwable cause) {
    FetchRequestInfo info = removeFetchRequest(streamChunkSlice);
    if (info != null) {
      ChunkReceivedCallback callback = info.callback;
      if (callback != null) {
        callback.onFailure(streamChunkSlice.chunkIndex, new IOException(errorMsg, cause));
      } else {
        logger.warn(
            "FetchRequestInfo ({}) callback is null when listener handles fetch request failure",
            streamChunkSlice);
      }
    } else {
      logger.warn(
          "FetchRequestInfo ({}) not found/already addressed when listener handles fetch request failure",
          streamChunkSlice);
    }
  }
}
