/*
 * 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.master.clustermeta.ha;

import java.io.File;
import java.io.IOException;
import java.util.Optional;

import com.google.protobuf.InvalidProtocolBufferException;
import org.apache.ratis.protocol.Message;
import org.apache.ratis.statemachine.impl.SimpleStateMachineStorage;
import org.apache.ratis.thirdparty.com.google.protobuf.ByteString;

import org.apache.celeborn.common.client.MasterNotLeaderException;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.rpc.RpcCallContext;
import org.apache.celeborn.service.deploy.master.clustermeta.AbstractMetaManager;
import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos;

public class HAHelper {

  public static boolean checkShouldProcess(
      RpcCallContext context, AbstractMetaManager masterStatusSystem, boolean bindPreferIp) {
    HARaftServer ratisServer = getRatisServer(masterStatusSystem);
    if (ratisServer != null) {
      if (ratisServer.isLeader()) {
        return true;
      }
      sendFailure(context, ratisServer, null, bindPreferIp);
      return false;
    }
    return true;
  }

  public static void sendFailure(
      RpcCallContext context, HARaftServer ratisServer, Throwable cause, boolean bindPreferIp) {
    if (context != null) {
      if (ratisServer != null) {
        Optional<HARaftServer.LeaderPeerEndpoints> leaderPeer =
            ratisServer.getCachedLeaderPeerRpcEndpoint();
        if (leaderPeer.isPresent()) {
          context.sendFailure(
              new MasterNotLeaderException(
                  ratisServer.getRpcEndpoint(),
                  leaderPeer.get().rpcEndpoints,
                  leaderPeer.get().rpcInternalEndpoints,
                  bindPreferIp,
                  cause));
        } else {
          context.sendFailure(
              new MasterNotLeaderException(
                  ratisServer.getRpcEndpoint(),
                  MasterNotLeaderException.LEADER_NOT_PRESENTED,
                  cause));
        }
      } else {
        context.sendFailure(new CelebornIOException(cause.getMessage(), cause));
      }
    }
  }

  public static long getWorkerTimeoutDeadline(AbstractMetaManager masterStatusSystem) {
    HARaftServer ratisServer = getRatisServer(masterStatusSystem);
    if (ratisServer != null) {
      return ratisServer.getWorkerTimeoutDeadline();
    } else {
      return -1;
    }
  }

  public static long getAppTimeoutDeadline(AbstractMetaManager masterStatusSystem) {
    HARaftServer ratisServer = getRatisServer(masterStatusSystem);
    if (ratisServer != null) {
      return ratisServer.getAppTimeoutDeadline();
    } else {
      return -1;
    }
  }

  public static HARaftServer getRatisServer(AbstractMetaManager masterStatusSystem) {
    if ((masterStatusSystem instanceof HAMasterMetaManager)) {
      HARaftServer ratisServer = ((HAMasterMetaManager) masterStatusSystem).getRatisServer();
      return ratisServer;
    }

    return null;
  }

  public static ByteString convertRequestToByteString(ResourceProtos.ResourceRequest request) {
    byte[] requestBytes = request.toByteArray();
    return ByteString.copyFrom(requestBytes);
  }

  public static ResourceProtos.ResourceRequest convertByteStringToRequest(ByteString byteString)
      throws InvalidProtocolBufferException {
    byte[] bytes = byteString.toByteArray();
    return ResourceProtos.ResourceRequest.parseFrom(bytes);
  }

  public static Message convertResponseToMessage(ResourceProtos.ResourceResponse response) {
    byte[] requestBytes = response.toByteArray();
    return Message.valueOf(ByteString.copyFrom(requestBytes));
  }

  /**
   * Creates a temporary snapshot file.
   *
   * @param storage the snapshot storage
   * @return the temporary snapshot file
   * @throws IOException if error occurred while creating the snapshot file
   */
  public static File createTempSnapshotFile(SimpleStateMachineStorage storage) throws IOException {
    File tempDir = storage.getTmpDir();
    if (!tempDir.isDirectory() && !tempDir.mkdir()) {
      throw new IOException(
          "Cannot create temporary snapshot directory at " + tempDir.getAbsolutePath());
    }
    return File.createTempFile(
        "raft_snapshot_" + System.currentTimeMillis() + "_", ".dat", tempDir);
  }
}
