/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.

Licensed 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 "tensorflow/tools/proto_splitter/merge.h"

#include <algorithm>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "riegeli/base/object.h"  // from @riegeli
#include "riegeli/bytes/fd_reader.h"  // from @riegeli
#include "riegeli/records/record_reader.h"  // from @riegeli
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/file_system_helper.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/protobuf.h"
#include "tsl/platform/statusor.h"

namespace tensorflow::tools::proto_splitter {

using ::tensorflow::proto_splitter::ChunkedField;
using ::tensorflow::proto_splitter::ChunkedMessage;
using ::tensorflow::proto_splitter::ChunkInfo;
using ::tensorflow::proto_splitter::ChunkMetadata;
using ::tensorflow::proto_splitter::FieldIndex;
using tools::proto_splitter::GetChunkMetadata;
using tools::proto_splitter::GetRiegeliReader;
using tools::proto_splitter::OnlyContainsPb;
using tsl::protobuf::FieldDescriptor;
using tsl::protobuf::Message;
using tsl::protobuf::Reflection;

absl::Status Merger::Merge(const std::vector<std::unique_ptr<Message>>& chunks,
                           const ChunkedMessage& chunked_message,
                           Message* merged_message) {
  riegeli::RecordReader<riegeli::FdReader<>> null_reader{riegeli::kClosed};

  if (chunked_message.has_chunk_index()) {
    // Chunks referenced by fields should be merged into the parent chunk.
    merged_message->MergeFrom(*chunks[chunked_message.chunk_index()].get());
  }

  // Use each chunked_field within the chunked_message to merge its
  // corresponding chunk into merged_message.
  for (const auto& chunked_field : chunked_message.chunked_fields()) {
    absl::Status s = ProcessField(chunked_field, merged_message, {}, chunks,
                                  null_reader, MergerOp::MERGE);
    if (!s.ok()) return s;
  }

  return absl::OkStatus();
}

absl::Status Merger::Read(std::string prefix, Message* merged_message) {
  uint64_t start_time = Env::Default()->NowMicros();

  TF_ASSIGN_OR_RETURN(bool only_contains_pb, OnlyContainsPb(prefix));
  if (only_contains_pb) {
    return ReadPb(absl::StrCat(prefix, ".pb"), merged_message);
  }

  // Create riegeli reader for file.cpb
  TF_ASSIGN_OR_RETURN(auto reader,
                      GetRiegeliReader(absl::StrCat(prefix, ".cpb")));

  auto read_metadata = GetChunkMetadata(reader);
  if (!read_metadata.ok()) {
    reader.Close();
    return absl::FailedPreconditionError(
        absl::StrCat("Couldn't read ChunkMetadata from chunked proto.\n",
                     read_metadata.status().ToString()));
  }
  ChunkMetadata chunk_metadata = read_metadata.value();

  std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
      chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());

  // Read the remaining chunks.
  absl::Status s =
      ReadFields(chunk_metadata.message(), reader, chunks_info, merged_message);

  reader.Close();

  uint64_t end_time = Env::Default()->NowMicros();
  LOG(INFO) << "Finished reading and merging chunked proto, took "
            << HumanReadableDuration(end_time - start_time) << ".";
  return s;
}

absl::Status Merger::ReadPartial(absl::string_view prefix,
                                 const ChunkMetadata& chunk_metadata,
                                 Message* merged_message) {
  uint64_t start_time = Env::Default()->NowMicros();

  TF_ASSIGN_OR_RETURN(bool only_contains_pb, OnlyContainsPb(prefix));
  if (only_contains_pb) {
    return absl::FailedPreconditionError(
        absl::StrCat("Attempting to read part of a chunked proto .cpb file, "
                     "but only found a regular proto: ",
                     prefix, ".pb"));
  }

  // Create riegeli reader for file.cpb
  TF_ASSIGN_OR_RETURN(auto reader,
                      GetRiegeliReader(absl::StrCat(prefix, ".cpb")));

  std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
      chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());

  // Read the remaining chunks.
  absl::Status s =
      ReadFields(chunk_metadata.message(), reader, chunks_info, merged_message);

  reader.Close();

  uint64_t end_time = Env::Default()->NowMicros();
  LOG(INFO) << "Finished reading and merging chunked proto, took "
            << HumanReadableDuration(end_time - start_time) << ".";
  return s;
}

absl::Status Merger::ReadChunkedFromString(
    absl::string_view data, tsl::protobuf::Message* merged_message) {
  uint64_t start_time = Env::Default()->NowMicros();

  // Create riegeli string reader.
  TF_ASSIGN_OR_RETURN(auto reader, GetRiegeliStringReader(data));

  auto read_metadata = GetChunkMetadata(reader);
  if (!read_metadata.ok()) {
    reader.Close();
    return absl::FailedPreconditionError(
        absl::StrCat("Couldn't read ChunkMetadata from chunked proto.\n",
                     read_metadata.status().ToString()));
  }
  ChunkMetadata chunk_metadata = read_metadata.value();

  std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
      chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());

  // Read the remaining chunks.
  absl::Status s =
      ReadFields(chunk_metadata.message(), reader, chunks_info, merged_message);

  reader.Close();

  uint64_t end_time = Env::Default()->NowMicros();
  LOG(INFO) << "Finished reading and merging chunked proto, took "
            << HumanReadableDuration(end_time - start_time) << ".";
  return s;
}

absl::Status Merger::ReadPb(const std::string& pb_file,
                            Message* merged_message) {
  uint64_t start_time = Env::Default()->NowMicros();
  TF_ASSIGN_OR_RETURN(bool file_exists,
                      internal::FileExists(Env::Default(), pb_file));
  if (!file_exists)
    return absl::NotFoundError(absl::StrCat("File not found: ", pb_file));
  LOG(INFO) << "Reading binary proto from " << pb_file;
  auto ret = ReadBinaryProto(Env::Default(), pb_file, merged_message);
  uint64_t end_time = Env::Default()->NowMicros();
  LOG(INFO) << "Finished reading binary proto, took "
            << HumanReadableDuration(end_time - start_time) << ".";
  return ret;
}

namespace {

bool CompareMapKey(const FieldIndex::MapKey& key1,
                   const FieldIndex::MapKey& key2) {
  if (key1.type_case() != key2.type_case()) {
    return key1.type_case() < key2.type_case();
  }
  switch (key1.type_case()) {
    case FieldIndex::MapKey::TypeCase::kS:
      return key1.s() < key2.s();
    case FieldIndex::MapKey::TypeCase::kBoolean:
      return key1.boolean() < key2.boolean();
    case FieldIndex::MapKey::TypeCase::kUi32:
      return key1.ui32() < key2.ui32();
    case FieldIndex::MapKey::TypeCase::kUi64:
      return key1.ui64() < key2.ui64();
    case FieldIndex::MapKey::TypeCase::kI32:
      return key1.i32() < key2.i32();
    case FieldIndex::MapKey::TypeCase::kI64:
      return key1.i64() < key2.i64();
    default:
      return false;
  }
}

bool CompareFieldIndex(const FieldIndex& tag1, const FieldIndex& tag2) {
  if (tag1.kind_case() != tag2.kind_case()) {
    return tag1.kind_case() < tag2.kind_case();
  }
  switch (tag1.kind_case()) {
    case FieldIndex::KindCase::kField:
      return tag1.field() < tag2.field();
    case FieldIndex::KindCase::kIndex:
      return tag1.index() < tag2.index();
    case FieldIndex::KindCase::kMapKey:
      return CompareMapKey(tag1.map_key(), tag2.map_key());
    default:
      return false;
  }
}

}  // namespace

template <typename RecordReader>
absl::Status Merger::ReadFields(const ChunkedMessage& chunked_message,
                                RecordReader& reader,
                                const std::vector<ChunkInfo>& chunks_info,
                                tsl::protobuf::Message* merged_message) {
  if (chunked_message.has_chunk_index()) {
    // Chunks referenced by fields should be merged into the parent chunk.
    TF_ASSIGN_OR_RETURN(
        std::string chunk,
        ReadChunk(reader, chunks_info[chunked_message.chunk_index()]));
    if (!merged_message->MergeFromString(chunk)) {
      return absl::FailedPreconditionError(
          "Couldn't merge chunk into message.");
    }
  }

  // Sort the chunked_fields by depth and index.
  // For example, this ensures that GraphDef.library is merged before
  //   GraphDef.library.function[0], which will be merged before
  //   GraphDef.library.function[1].
  std::vector<ChunkedField> chunked_fields(
      chunked_message.chunked_fields().begin(),
      chunked_message.chunked_fields().end());
  std::sort(
      chunked_fields.begin(), chunked_fields.end(),
      // Ensure proper strict weak ordering.
      [](const ChunkedField& cf1, const ChunkedField& cf2) {
        int tag_depth =
            std::min(cf1.field_tag().size(), cf2.field_tag().size());
        for (int depth = 0; depth < tag_depth; ++depth) {
          const FieldIndex& tag1 = cf1.field_tag()[depth];
          const FieldIndex& tag2 = cf2.field_tag()[depth];
          bool is_equiv = false;
          if (tag1.kind_case() == tag2.kind_case()) {
            if (tag1.kind_case() == FieldIndex::KindCase::kField) {
              is_equiv = (tag1.field() == tag2.field());
            } else if (tag1.kind_case() == FieldIndex::KindCase::kIndex) {
              is_equiv = (tag1.index() == tag2.index());
            } else if (tag1.kind_case() == FieldIndex::KindCase::kMapKey) {
              is_equiv = !CompareMapKey(tag1.map_key(), tag2.map_key()) &&
                         !CompareMapKey(tag2.map_key(), tag1.map_key());
            } else {
              is_equiv = true;  // KIND_NOT_SET
            }
          }
          if (!is_equiv) {
            return CompareFieldIndex(tag1, tag2);
          }
        }
        if (cf1.field_tag().size() == cf2.field_tag().size()) {
          // If the fields are the same, merge the earlier chunks first.
          return cf1.message().chunk_index() < cf2.message().chunk_index();
        }
        return cf1.field_tag().size() < cf2.field_tag().size();
      });

  // Use each chunked_field within the chunked_message to merge its
  // corresponding chunk into merged_message.
  for (const auto& chunked_field : chunked_fields) {
    absl::Status s = ProcessField(chunked_field, merged_message, chunks_info,
                                  {}, reader, MergerOp::READ);
    if (!s.ok()) return s;
  }
  return absl::OkStatus();
}

template <typename RecordReader>
absl::Status Merger::ProcessField(
    const ChunkedField& chunked_field, Message* merged_message,
    const std::vector<ChunkInfo>& chunks_info,
    const std::vector<std::unique_ptr<Message>>& chunks, RecordReader& reader,
    MergerOp op) {
  std::string chunk;
  switch (op) {
    case MergerOp::READ: {
      TF_ASSIGN_OR_RETURN(
          chunk, ReadChunk(reader,
                           chunks_info[chunked_field.message().chunk_index()]));
      break;
    }
    case MergerOp::MERGE: {
      chunk =
          chunks[chunked_field.message().chunk_index()]->SerializeAsString();
      break;
    }
  }

  if (chunked_field.field_tag().empty()) {
    // Chunk is not a field within the parent, but instead a portion of the
    // parent itself. Needs to be concatenated.
    merged_message->MergeFromString(chunk);
    return absl::OkStatus();
  }

  uint64_t field_index;
  Message* curr_message = merged_message;
  // Use field tags to navigate the merged_message, constructing necessary
  // fields along the way.
  TF_ASSIGN_OR_RETURN(const std::vector<Field> fields,
                      GetFieldTypes(chunked_field.field_tag()));
  const FieldDescriptor* field_desc = nullptr;
  for (const auto& field : fields) {
    merged_message = curr_message;
    field_desc = merged_message->GetDescriptor()->FindFieldByNumber(
        std::get<int>(field.first));

    if (field_desc == nullptr) {
      return absl::FailedPreconditionError(
          absl::StrCat("Field with number ", std::get<int>(field.first),
                       " not found in message descriptor.",
                       merged_message->GetDescriptor()->full_name(), ".\n"));
    }

    auto res = GetMutableField(merged_message, field);
    if (!res.ok()) {
      if (!absl::IsNotFound(res.status())) return res.status();
      // Add missing field.
      if (field_desc->is_map()) {
        TF_RETURN_IF_ERROR(
            AddMapEntry(curr_message, field_desc, field.second.value()));
        res = GetMutableField(curr_message, field);
      } else {
        curr_message->GetReflection()->AddMessage(curr_message, field_desc);
        res = GetMutableField(curr_message, field);
      }
    }

    auto [parent, mutable_field, mutable_field_index] = res.value();
    if (mutable_field->is_repeated() && mutable_field_index != -1) {
      field_index = mutable_field_index;
      // Update merged_message to repeated element.
      curr_message = parent->GetReflection()->MutableRepeatedMessage(
          parent, mutable_field, std::max(0, mutable_field_index));
      if (mutable_field->is_map()) {
        // messages of map type have the value at field #2
        field_desc = mutable_field->message_type()->FindFieldByNumber(2);
        merged_message = curr_message;
        curr_message = curr_message->GetReflection()->MutableMessage(
            curr_message, field_desc);
      }
    } else if (mutable_field->type() == FieldDescriptor::Type::TYPE_MESSAGE) {
      // Update merged_message.
      curr_message =
          parent->GetReflection()->MutableMessage(parent, mutable_field);
    }
  }
  // merged_message now points to the Message whose field (described by
  // field_desc) will be added/set.

  const Reflection* reflection = merged_message->GetReflection();
  if (field_desc->is_repeated()) {
    // field may contain multiple elements
    auto message_callback = [&reflection, &merged_message, &field_index, &op,
                             &chunks, &chunked_field, &reader, &chunks_info,
                             &field_desc]() -> absl::Status {
      for (int _ = reflection->FieldSize(*merged_message, field_desc);
           _ <= field_index; _++) {
        reflection->AddMessage(merged_message, field_desc);
      }
      switch (op) {
        case MergerOp::MERGE:
          TF_RETURN_IF_ERROR(
              Merge(chunks, chunked_field.message(),
                    reflection->MutableRepeatedMessage(
                        merged_message, field_desc, field_index)));
          break;
        case MergerOp::READ:
          TF_RETURN_IF_ERROR(
              ReadFields(chunked_field.message(), reader, chunks_info,
                         reflection->MutableRepeatedMessage(
                             merged_message, field_desc, field_index)));
          break;
        default:
          return absl::InternalError("Encountered unknown MergerOp.");
      }
      return absl::OkStatus();
    };
    TF_RETURN_IF_ERROR(SetRepeatedFieldElement(
        merged_message, field_desc, field_index, chunk, message_callback));
  } else {
    // regular field
    auto message_callback = [&reflection, &merged_message, &op, &chunks,
                             &chunked_field, &reader, &chunks_info,
                             &field_desc]() -> absl::Status {
      switch (op) {
        case MergerOp::MERGE:
          TF_RETURN_IF_ERROR(
              Merge(chunks, chunked_field.message(),
                    reflection->MutableMessage(merged_message, field_desc)));
          break;
        case MergerOp::READ:
          TF_RETURN_IF_ERROR(ReadFields(
              chunked_field.message(), reader, chunks_info,
              reflection->MutableMessage(merged_message, field_desc)));
          break;
        default:
          return absl::InternalError("Encountered unknown MergerOp.");
      }
      return absl::OkStatus();
    };
    TF_RETURN_IF_ERROR(
        SetFieldElement(merged_message, field_desc, chunk, message_callback));
  }

  return absl::OkStatus();
}

}  // namespace tensorflow::tools::proto_splitter