/*
 * Copyright (C) 2019 The Android Open Source Project
 *
 * 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 "src/trace_processor/importers/proto/args_table_utils.h"

#include "perfetto/ext/base/string_view.h"
#include "perfetto/protozero/scattered_heap_buffer.h"
#include "protos/perfetto/common/descriptor.pbzero.h"
#include "protos/perfetto/trace/interned_data/interned_data.pbzero.h"
#include "protos/perfetto/trace/track_event/source_location.pbzero.h"
#include "src/protozero/test/example_proto/test_messages.pbzero.h"
#include "src/trace_processor/importers/common/args_tracker.h"
#include "src/trace_processor/storage/trace_storage.h"
#include "src/trace_processor/test_messages.descriptor.h"
#include "test/gtest_and_gmock.h"

namespace perfetto {
namespace trace_processor {
namespace {

constexpr size_t kChunkSize = 42;

using ::testing::_;
using ::testing::Eq;
using ::testing::Invoke;
using ::testing::NiceMock;

class ArgsTableUtilsTest : public ::testing::Test {
 protected:
  ArgsTableUtilsTest() {
    context_.storage.reset(new TraceStorage);
    storage_ = context_.storage.get();
    context_.global_args_tracker.reset(new GlobalArgsTracker(&context_));
    context_.args_tracker.reset(new ArgsTracker(&context_));
    sequence_state_.reset(new PacketSequenceState(&context_));
  }

  /**
   * Check whether the argument set contains the value with given flat_key and
   * key and is equal to the given value.
   */
  bool HasArg(ArgSetId set_id,
              const base::StringView& flat_key,
              const base::StringView& key,
              Variadic value) {
    const auto& args = storage_->arg_table();
    auto key_id = storage_->string_pool().GetId(key);
    EXPECT_TRUE(key_id);
    auto flat_key_id = storage_->string_pool().GetId(flat_key);
    EXPECT_TRUE(flat_key_id);

    RowMap rm = args.FilterToRowMap({args.arg_set_id().eq(set_id)});
    bool found = false;
    for (auto it = rm.IterateRows(); it; it.Next()) {
      if (args.key()[it.row()] == key_id) {
        EXPECT_EQ(args.flat_key()[it.row()], flat_key_id);
        if (storage_->GetArgValue(it.row()) == value) {
          found = true;
          break;
        }
      }
    }
    return found;
  }

  /**
   * Implementation of HasArg for a simple case when flat_key is equals to key,
   * so that two won't have to be repeated for each assertion.
   */
  bool HasArg(ArgSetId set_id, const base::StringView& key, Variadic value) {
    return HasArg(set_id, key, key, value);
  }

  uint32_t arg_set_id_ = 1;
  std::unique_ptr<PacketSequenceState> sequence_state_;
  TraceProcessorContext context_;
  TraceStorage* storage_;
};

TEST_F(ArgsTableUtilsTest, EnsureTestMessageProtoParses) {
  ProtoToArgsTable helper(&context_);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  EXPECT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();
}

TEST_F(ArgsTableUtilsTest, BasicSingleLayerProto) {
  using namespace protozero::test::protos::pbzero;
  protozero::HeapBuffered<EveryField> msg{kChunkSize, kChunkSize};
  msg->set_field_int32(-1);
  msg->set_field_int64(-333123456789ll);
  msg->set_field_uint32(600);
  msg->set_field_uint64(333123456789ll);
  msg->set_field_sint32(-5);
  msg->set_field_sint64(-9000);
  msg->set_field_fixed32(12345);
  msg->set_field_fixed64(444123450000ll);
  msg->set_field_sfixed32(-69999);
  msg->set_field_sfixed64(-200);
  msg->set_field_double(0.5555);
  msg->set_field_bool(true);
  msg->set_small_enum(SmallEnum::TO_BE);
  msg->set_signed_enum(SignedEnum::NEGATIVE);
  msg->set_big_enum(BigEnum::BEGIN);
  msg->set_nested_enum(EveryField::PONG);
  msg->set_field_float(3.14f);
  msg->set_field_string("FizzBuzz");
  msg->add_repeated_int32(1);
  msg->add_repeated_int32(-1);
  msg->add_repeated_int32(100);
  msg->add_repeated_int32(2000000);

  auto binary_proto = msg.SerializeAsArray();

  storage_->mutable_track_table()->Insert({});
  ProtoToArgsTable helper(&context_);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  auto inserter = context_.args_tracker->AddArgsTo(TrackId(0));
  status = helper.InternProtoFieldsIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.EveryField", nullptr, &inserter,
      sequence_state_->current_generation().get());

  EXPECT_TRUE(status.ok())
      << "InternProtoFieldsIntoArgsTable failed with error: "
      << status.message();

  context_.args_tracker->Flush();
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_int32", Variadic::Integer(-1)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_int64",
                     Variadic::Integer(-333123456789ll)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_uint32",
                     Variadic::UnsignedInteger(600)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_uint64",
                     Variadic::UnsignedInteger(333123456789ll)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_sint32", Variadic::Integer(-5)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_sint64", Variadic::Integer(-9000)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_fixed32", Variadic::Integer(12345)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_fixed64",
                     Variadic::Integer(444123450000ll)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_sfixed32",
                     Variadic::Integer(-69999)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_sfixed64", Variadic::Integer(-200)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_double", Variadic::Real(0.5555)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_bool", Variadic::Boolean(true)));
  EXPECT_TRUE(HasArg(
      ArgSetId(arg_set_id_), "small_enum",
      Variadic::String(*context_.storage->string_pool().GetId("TO_BE"))));
  EXPECT_TRUE(HasArg(
      ArgSetId(arg_set_id_), "signed_enum",
      Variadic::String(*context_.storage->string_pool().GetId("NEGATIVE"))));
  EXPECT_TRUE(HasArg(
      ArgSetId(arg_set_id_), "big_enum",
      Variadic::String(*context_.storage->string_pool().GetId("BEGIN"))));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "nested_enum",
             Variadic::String(*context_.storage->string_pool().GetId("PONG"))));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_float",
                     Variadic::Real(static_cast<double>(3.14f))));
  ASSERT_TRUE(context_.storage->string_pool().GetId("FizzBuzz"));
  EXPECT_TRUE(HasArg(
      ArgSetId(arg_set_id_), "field_string",
      Variadic::String(*context_.storage->string_pool().GetId("FizzBuzz"))));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "repeated_int32",
                     "repeated_int32[0]", Variadic::Integer(1)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "repeated_int32",
                     "repeated_int32[1]", Variadic::Integer(-1)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "repeated_int32",
                     "repeated_int32[2]", Variadic::Integer(100)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "repeated_int32",
                     "repeated_int32[3]", Variadic::Integer(2000000)));
}

TEST_F(ArgsTableUtilsTest, NestedProto) {
  using namespace protozero::test::protos::pbzero;
  protozero::HeapBuffered<NestedA> msg{kChunkSize, kChunkSize};
  msg->set_super_nested()->set_value_c(3);

  auto binary_proto = msg.SerializeAsArray();

  storage_->mutable_track_table()->Insert({});
  ProtoToArgsTable helper(&context_);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  auto inserter = context_.args_tracker->AddArgsTo(TrackId(0));
  status = helper.InternProtoFieldsIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.NestedA", nullptr, &inserter,
      sequence_state_->current_generation().get());
  EXPECT_TRUE(status.ok())
      << "InternProtoFieldsIntoArgsTable failed with error: "
      << status.message();
  context_.args_tracker->Flush();
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "super_nested.value_c",
                     Variadic::Integer(3)));
}

TEST_F(ArgsTableUtilsTest, CamelCaseFieldsProto) {
  using namespace protozero::test::protos::pbzero;
  protozero::HeapBuffered<CamelCaseFields> msg{kChunkSize, kChunkSize};
  msg->set_barbaz(true);
  msg->set_moomoo(true);
  msg->set___bigbang(true);

  auto binary_proto = msg.SerializeAsArray();

  storage_->mutable_track_table()->Insert({});
  ProtoToArgsTable helper(&context_);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  auto inserter = context_.args_tracker->AddArgsTo(TrackId(0));
  status = helper.InternProtoFieldsIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.CamelCaseFields", nullptr, &inserter,
      sequence_state_->current_generation().get());
  EXPECT_TRUE(status.ok())
      << "InternProtoFieldsIntoArgsTable failed with error: "
      << status.message();
  context_.args_tracker->Flush();
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "barBaz", Variadic::Boolean(true)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "MooMoo", Variadic::Boolean(true)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "__bigBang", Variadic::Boolean(true)));
}

TEST_F(ArgsTableUtilsTest, NestedProtoParsingOverrideHandled) {
  using namespace protozero::test::protos::pbzero;
  protozero::HeapBuffered<NestedA> msg{kChunkSize, kChunkSize};
  msg->set_super_nested()->set_value_c(3);

  auto binary_proto = msg.SerializeAsArray();

  storage_->mutable_track_table()->Insert({});
  ProtoToArgsTable helper(&context_);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  helper.AddParsingOverride(
      "super_nested.value_c",
      [](const ProtoToArgsTable::ParsingOverrideState& state,
         const protozero::Field& field, ArgsTracker::BoundInserter* inserter) {
        EXPECT_EQ(field.type(), protozero::proto_utils::ProtoWireType::kVarInt);
        EXPECT_TRUE(state.context);
        EXPECT_TRUE(state.sequence_state);
        auto id = state.context->storage->InternString(
            "super_nested.value_b.replaced");
        int32_t val = field.as_int32();
        inserter->AddArg(id, id, Variadic::Integer(val + 1));
        // We've handled this field by adding the desired args.
        return true;
      });

  auto inserter = context_.args_tracker->AddArgsTo(TrackId(0));
  status = helper.InternProtoFieldsIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.NestedA", nullptr, &inserter,
      sequence_state_->current_generation().get());
  EXPECT_TRUE(status.ok())
      << "InternProtoFieldsIntoArgsTable failed with error: "
      << status.message();
  context_.args_tracker->Flush();
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "super_nested.value_b.replaced",
                     Variadic::Integer(4)));
}

TEST_F(ArgsTableUtilsTest, NestedProtoParsingOverrideSkipped) {
  using namespace protozero::test::protos::pbzero;
  protozero::HeapBuffered<NestedA> msg{kChunkSize, kChunkSize};
  msg->set_super_nested()->set_value_c(3);

  auto binary_proto = msg.SerializeAsArray();

  storage_->mutable_track_table()->Insert({});
  ProtoToArgsTable helper(&context_);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  helper.AddParsingOverride(
      "super_nested.value_c",
      [](const ProtoToArgsTable::ParsingOverrideState& state,
         const protozero::Field& field, ArgsTracker::BoundInserter*) {
        static int val = 0;
        ++val;
        EXPECT_EQ(1, val);
        EXPECT_EQ(field.type(), protozero::proto_utils::ProtoWireType::kVarInt);
        EXPECT_TRUE(state.sequence_state);
        EXPECT_TRUE(state.context);
        // By returning false we expect this field to be handled like regular.
        return false;
      });

  auto inserter = context_.args_tracker->AddArgsTo(TrackId(0));
  status = helper.InternProtoFieldsIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.NestedA", nullptr, &inserter,
      sequence_state_->current_generation().get());
  EXPECT_TRUE(status.ok())
      << "InternProtoFieldsIntoArgsTable failed with error: "
      << status.message();
  context_.args_tracker->Flush();
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "super_nested.value_c",
                     Variadic::Integer(3)));
}

TEST_F(ArgsTableUtilsTest, LookingUpInternedStateParsingOverride) {
  using namespace protozero::test::protos::pbzero;
  // The test proto, we will use |value_c| as the source_location iid.
  protozero::HeapBuffered<NestedA> msg{kChunkSize, kChunkSize};
  msg->set_super_nested()->set_value_c(3);
  auto binary_proto = msg.SerializeAsArray();

  // The interned source location.
  protozero::HeapBuffered<protos::pbzero::SourceLocation> src_loc{kChunkSize,
                                                                  kChunkSize};
  src_loc->set_iid(3);
  src_loc->set_file_name("test_file_name");
  // We need to update sequence_state to point to it.
  auto binary_data = src_loc.SerializeAsArray();
  std::unique_ptr<uint8_t[]> buffer(new uint8_t[binary_data.size()]);
  for (size_t i = 0; i < binary_data.size(); ++i) {
    buffer.get()[i] = binary_data[i];
  }
  TraceBlobView blob(std::move(buffer), 0, binary_data.size());
  sequence_state_->InternMessage(
      protos::pbzero::InternedData::kSourceLocationsFieldNumber,
      std::move(blob));

  ProtoToArgsTable helper(&context_);
  // Now we override the behaviour of |value_c| so we can expand the iid into
  // multiple args rows.
  helper.AddParsingOverride(
      "super_nested.value_c",
      [](const ProtoToArgsTable::ParsingOverrideState& state,
         const protozero::Field& field, ArgsTracker::BoundInserter* inserter) {
        EXPECT_EQ(field.type(), protozero::proto_utils::ProtoWireType::kVarInt);
        auto* decoder = state.sequence_state->LookupInternedMessage<
            protos::pbzero::InternedData::kSourceLocationsFieldNumber,
            protos::pbzero::SourceLocation>(field.as_uint64());
        if (!decoder) {
          // Lookup failed fall back on default behaviour.
          return false;
        }
        TraceStorage* storage = state.context->storage.get();
        auto file_name_id = storage->InternString("file_name");
        auto line_number_id = storage->InternString("line_number");
        auto file_id = storage->InternString(decoder->file_name());
        inserter->AddArg(file_name_id, file_name_id, Variadic::String(file_id));
        inserter->AddArg(line_number_id, line_number_id, Variadic::Integer(2));
        return true;
      });

  storage_->mutable_track_table()->Insert({});
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  auto inserter = context_.args_tracker->AddArgsTo(TrackId(0));
  status = helper.InternProtoFieldsIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.NestedA", nullptr, &inserter,
      sequence_state_->current_generation().get());
  EXPECT_TRUE(status.ok())
      << "InternProtoFieldsIntoArgsTable failed with error: "
      << status.message();
  auto file_name_id = storage_->string_pool().GetId("test_file_name");
  ASSERT_TRUE(file_name_id);
  context_.args_tracker->Flush();
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "file_name",
                     Variadic::String(*file_name_id)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "line_number", Variadic::Integer(2)));
}

}  // namespace
}  // namespace trace_processor
}  // namespace perfetto