// Copyright 2020 Google Inc. 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 <stdexcept>

#include "absl/strings/match.h"
#include "absl/strings/str_split.h"
#include "absl/strings/substitute.h"

#include "google/protobuf/arena.h"

#include "bloaty.h"
#include "report.pb.h"

namespace bloaty {

void RollupOutput::PrintToProtobuf(std::ostream* out) const {
  using namespace ::bloaty_report;

  google::protobuf::Arena arena;
  Report* report = google::protobuf::Arena::CreateMessage<Report>(&arena);
  report->set_file_total(toplevel_row_.filesize);
  report->set_vm_total(toplevel_row_.vmsize);

  ([&] {
    // If 2 levels, assume compileunits,symbols.
    // If 3 levels, assume accesspattern,compileunits,symbols; and pick the cold accesses.
    auto* compile_unit_row = &toplevel_row_;
    if (toplevel_row_.sorted_children.front().sorted_children.front().sorted_children.size() > 0) {
      bool found_hot = false;
      for (const auto& access : toplevel_row_.sorted_children) {
        if (access.name == "Hot") {
          found_hot = true;
        }
        if (access.name == "Cold") {
          compile_unit_row = &access;
          break;
        }
      }
      if (compile_unit_row == &toplevel_row_) {
        if (!found_hot) {
          THROWF("Could not find hot/cold regions in the binary, top-level row has $0 children "
                 "([$1, ...]), "
                 "next level row has $2 children ([$3, ...])",
                 toplevel_row_.sorted_children.size(),
                 toplevel_row_.sorted_children.front().name,
                 toplevel_row_.sorted_children.front().sorted_children.size(),
                 toplevel_row_.sorted_children.front().sorted_children.front().name);
        } else {
          // No cold regions, hence nothing to output.
          return;
        }
      }
    }

    for (const auto& child_row : compile_unit_row->sorted_children) {
      CompileUnit* compile_unit = report->add_compile_units();
      for (const auto& symbol_row : child_row.sorted_children) {
        SizeInfo* info = google::protobuf::Arena::CreateMessage<SizeInfo>(&arena);
        info->set_file_actual(symbol_row.filesize);
        info->set_vm_actual(symbol_row.vmsize);
        Symbol* symbol = compile_unit->add_symbols();
        symbol->set_allocated_sizes(info);
        auto decoded_symbol = DecodeSymbolWithCrateId(symbol_row.name);
        symbol->set_name(decoded_symbol.symbol);
        symbol->set_maybe_rust_crate(decoded_symbol.crate);
      }
      if (child_row.sorted_children.empty()) {
        // Add a "fake symbol" that is the same as the compile unit,
        // to make sizes add up.
        SizeInfo* info = google::protobuf::Arena::CreateMessage<SizeInfo>(&arena);
        info->set_file_actual(child_row.filesize);
        info->set_vm_actual(child_row.vmsize);
        Symbol* symbol = compile_unit->add_symbols();
        symbol->set_allocated_sizes(info);
        symbol->set_name(child_row.name);
      }

      SizeInfo* info = google::protobuf::Arena::CreateMessage<SizeInfo>(&arena);
      info->set_file_actual(child_row.filesize);
      info->set_vm_actual(child_row.vmsize);
      compile_unit->set_allocated_sizes(info);
      compile_unit->set_name(child_row.name);
    }
  })();

  bool good = report->SerializeToOstream(out);
  if (!good) {
    throw std::runtime_error("Failed to serialize report");
  }
  out->flush();
}

namespace {

const std::string kSeparator = ", in crate ";

}  // namespace

std::string EncodeSymbolWithCrateId(absl::string_view symbol, absl::string_view crate) {
  // The encoding format is "$symbol, in crate $crate". Since regular C++ and
  // Rust code is extremely unlikely to produce the word "in crate", This format
  // is unlikely to cause ambiguities, while staying readable if bloaty is
  // invoked from the command line.
  return std::string(symbol) + kSeparator + std::string(crate);
}

DecodeCrateIdResult DecodeSymbolWithCrateId(absl::string_view symbol) {
  if (absl::StrContains(symbol, kSeparator)) {
    std::vector<std::string> parts = absl::StrSplit(symbol, kSeparator);
    if (parts.size() != 2) {
      throw std::runtime_error("Unexpected symbol when decoding: " + std::string(symbol));
    }
    return DecodeCrateIdResult{.symbol = parts[0], .crate = parts[1]};
  } else {
    return DecodeCrateIdResult{.symbol = std::string(symbol)};
  }
}

}  // namespace bloaty