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

#include "celeborn/protocol/ControlMessages.h"
#include "celeborn/utils/CelebornUtils.h"

namespace celeborn {
namespace protocol {

namespace {
std::vector<std::unique_ptr<const PartitionLocation>>
fromPbPackedPartitionLocationsPair(
    const PbPackedPartitionLocationsPair& pbPackedPartitionLocationsPair) {
  std::vector<std::unique_ptr<const PartitionLocation>> finalLocations;
  std::vector<std::unique_ptr<PartitionLocation>> partialLocations;
  int inputLocationSize = pbPackedPartitionLocationsPair.inputlocationsize();
  auto& pbPackedPartitionLocations = pbPackedPartitionLocationsPair.locations();
  auto& pbIds = pbPackedPartitionLocations.ids();
  for (int idx = 0; idx < pbIds.size(); idx++) {
    partialLocations.push_back(
        PartitionLocation::fromPackedPb(pbPackedPartitionLocations, idx));
  }
  for (int idx = 0; idx < inputLocationSize; idx++) {
    auto replicaIdx = pbPackedPartitionLocationsPair.peerindexes(idx);
    // Has peer.
    if (replicaIdx != INT_MAX) {
      CELEBORN_CHECK_GE(replicaIdx, inputLocationSize);
      CELEBORN_CHECK_LT(replicaIdx, partialLocations.size());
      auto location = std::move(partialLocations[idx]);
      auto peerLocation = std::move(partialLocations[replicaIdx]);
      // Make sure the location is primary and peer location is replica.
      if (location->mode == PartitionLocation::Mode::REPLICA) {
        std::swap(location, peerLocation);
      }
      CELEBORN_CHECK(location->mode == PartitionLocation::Mode::PRIMARY);
      CELEBORN_CHECK(peerLocation->mode == PartitionLocation::Mode::REPLICA);
      location->replicaPeer = std::move(peerLocation);
      finalLocations.push_back(std::move(location));
    }
    // Has no peer.
    else {
      finalLocations.push_back(std::move(partialLocations[idx]));
    }
  }
  return finalLocations;
}

} // namespace

TransportMessage RegisterShuffle::toTransportMessage() const {
  MessageType type = REGISTER_SHUFFLE;
  PbRegisterShuffle pb;
  pb.set_shuffleid(shuffleId);
  pb.set_nummappers(numMappers);
  pb.set_numpartitions(numPartitions);
  std::string payload = pb.SerializeAsString();
  return TransportMessage(type, std::move(payload));
}

std::unique_ptr<RegisterShuffleResponse>
RegisterShuffleResponse::fromTransportMessage(
    const TransportMessage& transportMessage) {
  CELEBORN_CHECK(
      transportMessage.type() == REGISTER_SHUFFLE_RESPONSE,
      "transportMessageType mismatch");
  auto payload = transportMessage.payload();
  auto pbRegisterShuffleResponse = utils::parseProto<PbRegisterShuffleResponse>(
      reinterpret_cast<const uint8_t*>(payload.c_str()), payload.size());
  auto response = std::make_unique<RegisterShuffleResponse>();
  response->status = toStatusCode(pbRegisterShuffleResponse->status());

  // Legacy mode is deprecated.
  auto& pbPartitionLocations = pbRegisterShuffleResponse->partitionlocations();
  CELEBORN_CHECK_EQ(
      pbPartitionLocations.size(),
      0,
      "legacy PartitionLocation pb is deprecated");

  // Packed mode: must use packedPartitionLocations.
  const auto& pbPackedPartitionLocationsPair =
      pbRegisterShuffleResponse->packedpartitionlocationspair();
  response->partitionLocations = std::move(
      fromPbPackedPartitionLocationsPair(pbPackedPartitionLocationsPair));
  return std::move(response);
}

TransportMessage MapperEnd::toTransportMessage() const {
  MessageType type = MAPPER_END;
  PbMapperEnd pb;
  pb.set_shuffleid(shuffleId);
  pb.set_mapid(mapId);
  pb.set_attemptid(attemptId);
  pb.set_nummappers(numMappers);
  pb.set_partitionid(partitionId);
  auto* pbPushFailureBatches = pb.mutable_pushfailurebatches();
  for (const auto& [partitionUniqueId, attemptBatches] : pushFailedBatches) {
    PbLocationPushFailedBatches pbLocationBatches;
    auto* pbFailedBatchesMap = pbLocationBatches.mutable_failedbatches();
    for (const auto& [attemptKey, batchIds] : attemptBatches) {
      PbFailedBatches pbFailedBatches;
      for (int batchId : batchIds) {
        pbFailedBatches.add_failedbatches(batchId);
      }
      (*pbFailedBatchesMap)[attemptKey] = std::move(pbFailedBatches);
    }
    (*pbPushFailureBatches)[partitionUniqueId] = std::move(pbLocationBatches);
  }
  std::string payload = pb.SerializeAsString();
  return TransportMessage(type, std::move(payload));
}

std::unique_ptr<MapperEndResponse> MapperEndResponse::fromTransportMessage(
    const TransportMessage& transportMessage) {
  CELEBORN_CHECK(
      transportMessage.type() == MAPPER_END_RESPONSE,
      "transportMessageType mismatch");
  auto payload = transportMessage.payload();
  auto pbMapperEndResponse = utils::parseProto<PbMapperEndResponse>(
      reinterpret_cast<const uint8_t*>(payload.c_str()), payload.size());
  auto response = std::make_unique<MapperEndResponse>();
  response->status = toStatusCode(pbMapperEndResponse->status());
  return std::move(response);
}

ReviveRequest::ReviveRequest(
    int _shuffleId,
    int _mapId,
    int _attemptId,
    int _partitionId,
    int _epoch,
    std::shared_ptr<const PartitionLocation> _loc,
    StatusCode _cause)
    : shuffleId(_shuffleId),
      mapId(_mapId),
      attemptId(_attemptId),
      partitionId(_partitionId),
      epoch(_epoch),
      loc(std::move(_loc)),
      cause(_cause) {}

TransportMessage Revive::toTransportMessage() const {
  MessageType type = CHANGE_LOCATION;
  PbRevive pb;
  pb.set_shuffleid(shuffleId);
  for (auto mapId : mapIds) {
    pb.add_mapid(mapId);
  }
  for (auto& reviveRequest : reviveRequests) {
    auto pbRevivePartitionInfo = pb.add_partitioninfo();
    pbRevivePartitionInfo->set_partitionid(reviveRequest->partitionId);
    pbRevivePartitionInfo->set_epoch(reviveRequest->epoch);
    pbRevivePartitionInfo->set_status(reviveRequest->cause);
    if (reviveRequest->loc) {
      pbRevivePartitionInfo->set_allocated_partition(
          reviveRequest->loc->toPb().release());
    }
  }
  std::string payload = pb.SerializeAsString();
  return TransportMessage(type, std::move(payload));
}

std::unique_ptr<ChangeLocationResponse>
ChangeLocationResponse::fromTransportMessage(
    const TransportMessage& transportMessage) {
  CELEBORN_CHECK(
      transportMessage.type() == CHANGE_LOCATION_RESPONSE,
      "transportMessageType mismatch");
  auto payload = transportMessage.payload();
  auto pbChangeLocationResponse = utils::parseProto<PbChangeLocationResponse>(
      reinterpret_cast<const uint8_t*>(payload.c_str()), payload.size());
  auto response = std::make_unique<ChangeLocationResponse>();
  response->endedMapIds.reserve(pbChangeLocationResponse->endedmapid_size());
  for (auto endedMapId : pbChangeLocationResponse->endedmapid()) {
    response->endedMapIds.push_back(endedMapId);
  }
  int numPartitionInfo = pbChangeLocationResponse->partitioninfo_size();
  response->partitionInfos.resize(numPartitionInfo);
  for (int i = 0; i < numPartitionInfo; i++) {
    auto& partitionInfo = response->partitionInfos[i];
    auto& pbPartitionInfo = pbChangeLocationResponse->partitioninfo(i);
    partitionInfo.partitionId = pbPartitionInfo.partitionid();
    partitionInfo.status = toStatusCode(pbPartitionInfo.status());
    partitionInfo.partition =
        PartitionLocation::fromPb(pbPartitionInfo.partition());
    partitionInfo.oldAvailable = pbPartitionInfo.oldavailable();
  }
  return std::move(response);
}

TransportMessage GetReducerFileGroup::toTransportMessage() const {
  MessageType type = GET_REDUCER_FILE_GROUP;
  PbGetReducerFileGroup pb;
  pb.set_shuffleid(shuffleId);
  std::string payload = pb.SerializeAsString();
  return TransportMessage(type, std::move(payload));
}

std::unique_ptr<GetReducerFileGroupResponse>
GetReducerFileGroupResponse::fromTransportMessage(
    const TransportMessage& transportMessage) {
  CELEBORN_CHECK(
      transportMessage.type() == GET_REDUCER_FILE_GROUP_RESPONSE,
      "transportMessageType mismatch");
  auto payload = transportMessage.payload();
  auto pbGetReducerFileGroupResponse =
      utils::parseProto<PbGetReducerFileGroupResponse>(
          reinterpret_cast<const uint8_t*>(payload.c_str()), payload.size());
  auto response = std::make_unique<GetReducerFileGroupResponse>();
  response->status = toStatusCode(pbGetReducerFileGroupResponse->status());
  auto fileGroups = pbGetReducerFileGroupResponse->filegroups();
  for (auto& kv : fileGroups) {
    auto& fileGroup = response->fileGroups[kv.first];
    // Legacy mode is deprecated.
    CELEBORN_CHECK_EQ(
        kv.second.locations().size(),
        0,
        "legecy PartitionLocation pb is deprecated");
    // Packed mode: must use packedPartitionLocations.
    const auto& pbPackedPartitionLocationsPair =
        kv.second.partitionlocationspair();
    auto locations =
        fromPbPackedPartitionLocationsPair(pbPackedPartitionLocationsPair);
    for (auto& location : locations) {
      fileGroup.insert(std::move(location));
    }
  }
  auto attempts = pbGetReducerFileGroupResponse->attempts();
  response->attempts.reserve(attempts.size());
  for (auto attempt : attempts) {
    response->attempts.push_back(attempt);
  }
  auto partitionIds = pbGetReducerFileGroupResponse->partitionids();
  for (auto partitionId : partitionIds) {
    response->partitionIds.insert(partitionId);
  }
  const auto& pushFailedBatches =
      pbGetReducerFileGroupResponse->pushfailedbatches();
  for (const auto& locationEntry : pushFailedBatches) {
    auto& attemptBatches = response->pushFailedBatches[locationEntry.first];
    for (const auto& attemptEntry : locationEntry.second.failedbatches()) {
      auto& batchIds = attemptBatches[attemptEntry.first];
      for (int batchId : attemptEntry.second.failedbatches()) {
        batchIds.insert(batchId);
      }
    }
  }
  return std::move(response);
}

OpenStream::OpenStream(
    const std::string& shuffleKey,
    const std::string& filename,
    int32_t startMapIndex,
    int32_t endMapIndex)
    : shuffleKey(shuffleKey),
      filename(filename),
      startMapIndex(startMapIndex),
      endMapIndex(endMapIndex) {}

TransportMessage OpenStream::toTransportMessage() const {
  MessageType type = OPEN_STREAM;
  PbOpenStream pb;
  pb.set_shufflekey(shuffleKey);
  pb.set_filename(filename);
  pb.set_startindex(startMapIndex);
  pb.set_endindex(endMapIndex);
  std::string payload = pb.SerializeAsString();
  return TransportMessage(type, std::move(payload));
}

std::unique_ptr<StreamHandler> StreamHandler::fromTransportMessage(
    const TransportMessage& transportMessage) {
  CELEBORN_CHECK(
      transportMessage.type() == STREAM_HANDLER,
      "transportMessageType should be STREAM_HANDLER");
  auto payload = transportMessage.payload();
  auto pbStreamHandler = utils::parseProto<PbStreamHandler>(
      reinterpret_cast<const uint8_t*>(payload.c_str()), payload.size());
  auto streamHandler = std::make_unique<StreamHandler>();
  streamHandler->streamId = pbStreamHandler->streamid();
  streamHandler->numChunks = pbStreamHandler->numchunks();
  for (auto chunkOffset : pbStreamHandler->chunkoffsets()) {
    streamHandler->chunkOffsets.push_back(chunkOffset);
  }
  streamHandler->fullPath = pbStreamHandler->fullpath();
  return std::move(streamHandler);
}

std::unique_ptr<PbStreamChunkSlice> StreamChunkSlice::toProto() const {
  auto pb = std::make_unique<PbStreamChunkSlice>();
  pb->set_streamid(streamId);
  pb->set_chunkindex(chunkIndex);
  pb->set_offset(offset);
  pb->set_len(len);
  return std::move(pb);
}

StreamChunkSlice StreamChunkSlice::decodeFrom(
    memory::ReadOnlyByteBuffer& data) {
  CELEBORN_CHECK_GE(data.remainingSize(), 20);
  StreamChunkSlice slice;
  slice.streamId = data.read<long>();
  slice.chunkIndex = data.read<int>();
  slice.offset = data.read<int>();
  slice.len = data.read<int>();
  return slice;
}

size_t StreamChunkSlice::Hasher::operator()(const StreamChunkSlice& lhs) const {
  const auto hashStreamId = std::hash<long>()(lhs.streamId);
  const auto hashChunkIndex = std::hash<int>()(lhs.chunkIndex) << 1;
  const auto hashOffset = std::hash<int>()(lhs.offset) << 2;
  const auto hashLen = std::hash<int>()(lhs.len) << 3;
  return hashStreamId ^ hashChunkIndex ^ hashOffset ^ hashLen;
}

TransportMessage ChunkFetchRequest::toTransportMessage() const {
  MessageType type = CHUNK_FETCH_REQUEST;
  PbChunkFetchRequest pb;
  pb.unsafe_arena_set_allocated_streamchunkslice(
      streamChunkSlice.toProto().release());
  std::string payload = pb.SerializeAsString();
  return TransportMessage(type, std::move(payload));
}

TransportMessage BufferStreamEnd::toTransportMessage() const {
  MessageType type = BUFFER_STREAM_END;
  PbBufferStreamEnd pb;
  pb.set_streamtype(ChunkStream);
  pb.set_streamid(streamId);
  std::string payload = pb.SerializeAsString();
  return TransportMessage(type, std::move(payload));
}
} // namespace protocol
} // namespace celeborn
