/*
 * 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.util.List;
import java.util.Map;
import java.util.stream.Collectors;

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

import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.client.MasterClient;
import org.apache.celeborn.common.exception.CelebornRuntimeException;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.meta.ApplicationMeta;
import org.apache.celeborn.common.meta.DiskInfo;
import org.apache.celeborn.common.meta.WorkerInfo;
import org.apache.celeborn.common.meta.WorkerStatus;
import org.apache.celeborn.common.network.CelebornRackResolver;
import org.apache.celeborn.common.quota.ResourceConsumption;
import org.apache.celeborn.common.rpc.RpcEnv;
import org.apache.celeborn.service.deploy.master.clustermeta.AbstractMetaManager;
import org.apache.celeborn.service.deploy.master.clustermeta.MetaUtil;
import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos;
import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos.ResourceRequest;
import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos.Type;

public class HAMasterMetaManager extends AbstractMetaManager {
  private static final Logger LOG = LoggerFactory.getLogger(HAMasterMetaManager.class);

  protected HARaftServer ratisServer;

  public HAMasterMetaManager(RpcEnv rpcEnv, CelebornConf conf) {
    this(rpcEnv, conf, new CelebornRackResolver(conf));
  }

  public HAMasterMetaManager(RpcEnv rpcEnv, CelebornConf conf, CelebornRackResolver rackResolver) {
    this.rpcEnv = rpcEnv;
    this.conf = conf;
    this.initialEstimatedPartitionSize = conf.initialEstimatedPartitionSize();
    this.estimatedPartitionSize = initialEstimatedPartitionSize;
    this.unhealthyDiskRatioThreshold = conf.masterExcludeWorkerUnhealthyDiskRatioThreshold();
    this.rackResolver = rackResolver;
  }

  public HARaftServer getRatisServer() {
    return ratisServer;
  }

  public void setRatisServer(HARaftServer ratisServer) {
    this.ratisServer = ratisServer;
  }

  @Override
  public void handleRegisterApplicationInfo(
      String appId,
      UserIdentifier userIdentifier,
      Map<String, String> extraInfo,
      String requestId) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.RegisterApplicationInfo)
              .setRequestId(requestId)
              .setRegisterApplicationInfoRequest(
                  ResourceProtos.RegisterApplicationInfoRequest.newBuilder()
                      .setAppId(appId)
                      .setUserIdentifier(
                          ResourceProtos.UserIdentifier.newBuilder()
                              .setTenantId(userIdentifier.tenantId())
                              .setName(userIdentifier.name())
                              .build())
                      .putAllExtraInfo(extraInfo)
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle app lost for {} failed!", appId, e);
      throw e;
    }
  }

  @Override
  public void handleRequestSlots(
      String shuffleKey,
      String hostName,
      Map<String, Map<String, Integer>> workerToAllocatedSlots,
      String requestId) {
    try {
      ResourceProtos.RequestSlotsRequest.Builder builder =
          ResourceProtos.RequestSlotsRequest.newBuilder()
              .setShuffleKey(shuffleKey)
              .setHostName(hostName);
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.RequestSlots)
              .setRequestId(requestId)
              .setRequestSlotsRequest(builder.build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle request slots for {} failed!", shuffleKey, e);
      throw e;
    }
  }

  @Override
  public void handleUnRegisterShuffle(String shuffleKey, String requestId) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.UnRegisterShuffle)
              .setRequestId(requestId)
              .setUnregisterShuffleRequest(
                  ResourceProtos.UnregisterShuffleRequest.newBuilder()
                      .setShuffleKey(shuffleKey)
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle unregister shuffle for {} failed!", shuffleKey, e);
      throw e;
    }
  }

  @Override
  public void handleBatchUnRegisterShuffles(List<String> shuffleKeys, String requestId) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.BatchUnRegisterShuffle)
              .setRequestId(requestId)
              .setBatchUnregisterShuffleRequest(
                  ResourceProtos.BatchUnregisterShuffleRequest.newBuilder()
                      .addAllShuffleKeys(shuffleKeys)
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Batch handle unregister shuffle for {} failed!", shuffleKeys, e);
      throw e;
    }
  }

  @Override
  public void handleAppHeartbeat(
      String appId,
      long totalWritten,
      long fileCount,
      long shuffleCount,
      long applicationCount,
      Map<String, Long> shuffleFallbackCounts,
      Map<String, Long> applicationFallbackCounts,
      long time,
      String requestId) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.AppHeartbeat)
              .setRequestId(requestId)
              .setAppHeartbeatRequest(
                  ResourceProtos.AppHeartbeatRequest.newBuilder()
                      .setAppId(appId)
                      .setTime(time)
                      .setTotalWritten(totalWritten)
                      .setFileCount(fileCount)
                      .setShuffleCount(shuffleCount)
                      .setApplicationCount(applicationCount)
                      .putAllShuffleFallbackCounts(shuffleFallbackCounts)
                      .putAllApplicationFallbackCounts(applicationFallbackCounts)
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle heartbeat for {} failed!", appId, e);
      throw e;
    }
  }

  @Override
  public void handleAppLost(String appId, String requestId) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.AppLost)
              .setRequestId(requestId)
              .setAppLostRequest(ResourceProtos.AppLostRequest.newBuilder().setAppId(appId).build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle app lost for {} failed!", appId, e);
      throw e;
    }
  }

  @Override
  public void handleWorkerExclude(
      List<WorkerInfo> workersToAdd, List<WorkerInfo> workersToRemove, String requestId) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.WorkerExclude)
              .setRequestId(requestId)
              .setWorkerExcludeRequest(
                  ResourceProtos.WorkerExcludeRequest.newBuilder()
                      .addAllWorkersToAdd(
                          workersToAdd.stream()
                              .map(MetaUtil::infoToAddr)
                              .collect(Collectors.toList()))
                      .addAllWorkersToRemove(
                          workersToRemove.stream()
                              .map(MetaUtil::infoToAddr)
                              .collect(Collectors.toList()))
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error(
          "Handle worker exclude for workersToAdd {} and workersToRemove {} failed!",
          workersToAdd.stream().map(WorkerInfo::toString).collect(Collectors.joining(",")),
          workersToRemove.stream().map(WorkerInfo::toString).collect(Collectors.joining(",")),
          e);
      throw e;
    }
  }

  @Override
  public void handleReviseLostShuffles(String appId, List<Integer> shuffles, String requestId) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.ReviseLostShuffles)
              .setRequestId(requestId)
              .setReviseLostShufflesRequest(
                  ResourceProtos.ReviseLostShufflesRequest.newBuilder()
                      .setAppId(appId)
                      .addAllLostShuffles(shuffles)
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle revise lost shuffle failed!", e);
      throw e;
    }
  }

  @Override
  public void handleWorkerLost(
      String host, int rpcPort, int pushPort, int fetchPort, int replicatePort, String requestId) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.WorkerLost)
              .setRequestId(requestId)
              .setWorkerLostRequest(
                  ResourceProtos.WorkerLostRequest.newBuilder()
                      .setHost(host)
                      .setRpcPort(rpcPort)
                      .setPushPort(pushPort)
                      .setFetchPort(fetchPort)
                      .setReplicatePort(replicatePort)
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle worker lost for {} failed!", host, e);
      throw e;
    }
  }

  @Override
  public void handleRemoveWorkersUnavailableInfo(
      List<WorkerInfo> unavailableWorkers, String requestId) {
    try {
      List<ResourceProtos.WorkerAddress> addrs =
          unavailableWorkers.stream().map(MetaUtil::infoToAddr).collect(Collectors.toList());
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.RemoveWorkersUnavailableInfo)
              .setRequestId(requestId)
              .setRemoveWorkersUnavailableInfoRequest(
                  ResourceProtos.RemoveWorkersUnavailableInfoRequest.newBuilder()
                      .addAllUnavailable(addrs)
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle remove workers unavailable info failed!", e);
      throw e;
    }
  }

  @Override
  public void handleWorkerHeartbeat(
      String host,
      int rpcPort,
      int pushPort,
      int fetchPort,
      int replicatePort,
      Map<String, DiskInfo> disks,
      Map<UserIdentifier, ResourceConsumption> userResourceConsumption,
      long time,
      boolean highWorkload,
      WorkerStatus workerStatus,
      String requestId) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.WorkerHeartbeat)
              .setRequestId(requestId)
              .setWorkerHeartbeatRequest(
                  ResourceProtos.WorkerHeartbeatRequest.newBuilder()
                      .setHost(host)
                      .setRpcPort(rpcPort)
                      .setPushPort(pushPort)
                      .setFetchPort(fetchPort)
                      .setReplicatePort(replicatePort)
                      .putAllDisks(MetaUtil.toPbDiskInfos(disks))
                      .setWorkerStatus(MetaUtil.toPbWorkerStatus(workerStatus))
                      .setTime(time)
                      .setHighWorkload(highWorkload)
                      .build())
              .build());
      updateWorkerResourceConsumptions(
          host, rpcPort, pushPort, fetchPort, replicatePort, userResourceConsumption);
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle worker heartbeat for {} failed!", host, e);
      throw e;
    }
  }

  @Override
  public void handleRegisterWorker(
      String host,
      int rpcPort,
      int pushPort,
      int fetchPort,
      int replicatePort,
      int internalPort,
      String networkLocation,
      Map<String, DiskInfo> disks,
      Map<UserIdentifier, ResourceConsumption> userResourceConsumption,
      String requestId) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.RegisterWorker)
              .setRequestId(requestId)
              .setRegisterWorkerRequest(
                  ResourceProtos.RegisterWorkerRequest.newBuilder()
                      .setHost(host)
                      .setRpcPort(rpcPort)
                      .setPushPort(pushPort)
                      .setFetchPort(fetchPort)
                      .setReplicatePort(replicatePort)
                      .setInternalPort(internalPort)
                      .setNetworkLocation(networkLocation)
                      .putAllDisks(MetaUtil.toPbDiskInfos(disks))
                      .build())
              .build());
      updateWorkerResourceConsumptions(
          host, rpcPort, pushPort, fetchPort, replicatePort, userResourceConsumption);
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle worker register for {} failed!", host, e);
      throw e;
    }
  }

  @Override
  public void handleReportWorkerUnavailable(List<WorkerInfo> failedNodes, String requestId) {
    try {
      List<ResourceProtos.WorkerAddress> addrs =
          failedNodes.stream().map(MetaUtil::infoToAddr).collect(Collectors.toList());
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.ReportWorkerUnavailable)
              .setRequestId(requestId)
              .setReportWorkerUnavailableRequest(
                  ResourceProtos.ReportWorkerUnavailableRequest.newBuilder()
                      .addAllUnavailable(addrs)
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle report node failure for {} failed!", failedNodes, e);
      throw e;
    }
  }

  @Override
  public void handleWorkerEvent(
      int workerEventTypeValue, List<WorkerInfo> workerInfoList, String requestId) {
    try {
      List<ResourceProtos.WorkerAddress> addrs =
          workerInfoList.stream().map(MetaUtil::infoToAddr).collect(Collectors.toList());
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.WorkerEvent)
              .setRequestId(requestId)
              .setWorkerEventRequest(
                  ResourceProtos.WorkerEventRequest.newBuilder()
                      .setWorkerEventType(
                          ResourceProtos.WorkerEventType.forNumber(workerEventTypeValue))
                      .addAllWorkerAddress(addrs)
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error(
          "Handle worker event {} failure for {} failed!", workerEventTypeValue, workerInfoList, e);
      throw e;
    }
  }

  @Override
  public void handleReportWorkerDecommission(List<WorkerInfo> workers, String requestId) {
    try {
      List<ResourceProtos.WorkerAddress> addrs =
          workers.stream().map(MetaUtil::infoToAddr).collect(Collectors.toList());
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.ReportWorkerDecommission)
              .setRequestId(MasterClient.genRequestId())
              .setReportWorkerDecommissionRequest(
                  ResourceProtos.ReportWorkerDecommissionRequest.newBuilder()
                      .addAllWorkers(addrs)
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle worker decommission failed!", e);
      throw e;
    }
  }

  @Override
  public void handleUpdatePartitionSize() {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.UpdatePartitionSize)
              .setRequestId(MasterClient.genRequestId())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle update partition size failed!", e);
      throw e;
    }
  }

  @Override
  public void handleApplicationMeta(ApplicationMeta applicationMeta) {
    try {
      ratisServer.submitRequest(
          ResourceRequest.newBuilder()
              .setCmdType(Type.ApplicationMeta)
              .setRequestId(MasterClient.genRequestId())
              .setApplicationMetaRequest(
                  ResourceProtos.ApplicationMetaRequest.newBuilder()
                      .setAppId(applicationMeta.appId())
                      .setSecret(applicationMeta.secret())
                      .build())
              .build());
    } catch (CelebornRuntimeException e) {
      LOG.error("Handle app meta for {} failed!", applicationMeta.appId(), e);
      throw e;
    }
  }
}
