/*
 * Copyright (C) 2021 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.
 */

#ifndef INCLUDE_PERFETTO_TRACING_TRACED_PROTO_H_
#define INCLUDE_PERFETTO_TRACING_TRACED_PROTO_H_

#include "perfetto/base/template_util.h"
#include "perfetto/protozero/field_writer.h"
#include "perfetto/protozero/proto_utils.h"
#include "perfetto/tracing/traced_value.h"

namespace perfetto {
class EventContext;
namespace internal {
template <typename FieldMetadata,
          bool is_message,
          protozero::proto_utils::RepetitionType repetition_type>
struct TypedProtoWriterImpl;
}

// A Wrapper around a protozero message to allow C++ classes to specify how it
// should be serialised into the trace:
//
// class Foo {
//  public:
//   void WriteIntoTrace(perfetto::TracedProto<pbzero::Foo> message) {
//     message->set_int_field(int_field_);
//   }
// };
//
// This class also exposes EventContext, e.g. to enable data interning.
//
// NOTE: the functionality below is not ready yet.
// TODO(altimin): Make the interop below possible.
// TracedProto also provides a seamless integration with writing untyped
// values via TracedValue / TracedDictionary / TracedArray:
//
// - TracedValue can be converted to a TracedProto, either by calling
//   TracedValue::WriteProto<T>() or implicitly.
// - If a proto message has a repeating DebugAnnotation debug_annotations
//   field, it can be filled using the TracedDictionary obtained from
//   TracedProto::AddDebugAnnotations.
template <typename MessageType>
class TracedProto {
 public:
  // implicit
  TracedProto(TracedValue&& value)
      : TracedProto(std::move(value).WriteProto<MessageType>()) {}
  ~TracedProto() = default;

  TracedProto(const TracedProto&) = delete;
  TracedProto& operator=(const TracedProto&) = delete;
  TracedProto& operator=(TracedProto&&) = delete;
  TracedProto(TracedProto&&) = default;

  MessageType* operator->() const { return message_; }

  MessageType* message() { return message_; }

  // Write additional untyped values into the same context, which is useful
  // when a given C++ class has a typed representation, but also either has
  // members which can only be written into an untyped context (e.g. they are
  // autogenerated) or it's desirable to have a way to quickly extend the
  // trace representation of this class (e.g. for debugging).
  //
  // The usage of the returned TracedDictionary should not be interleaved with
  // writing into |message| as this results in an inefficient proto layout. To
  // enforce this, AddDebugAnnotations should be called on TracedProto&&, i.e.
  // std::move(message).AddDebugAnnotations().
  //
  // This requires a 'repeated DebugAnnotations debug_annotations' field in
  // MessageType.
  template <typename Check = void>
  TracedDictionary AddDebugAnnotations() && {
    static_assert(
        std::is_base_of<
            protozero::proto_utils::FieldMetadataBase,
            typename MessageType::FieldMetadata_DebugAnnotations>::value,
        "This message does not have a |debug_annotations| field. Please add a"
        "'repeated perfetto.protos.DebugAnnotation debug_annnotations = N;' "
        "field to your message.");
    return TracedDictionary(message_, MessageType::kDebugAnnotations, context_,
                            nullptr);
  }

  // Start writing a single entry corresponding to the given |field| and return
  // TracedProto should be used to populate this further.
  // This method requires |field|'s type to be a nested message, but both
  // repeated and non-repeated complex fields are supported.
  template <typename FieldMetadata>
  TracedProto<typename FieldMetadata::cpp_field_type> WriteNestedMessage(
      FieldMetadata) {
    static_assert(std::is_base_of<MessageType,
                                  typename FieldMetadata::message_type>::value,
                  "Field should belong to the current message");
    static_assert(
        FieldMetadata::kProtoFieldType ==
            protozero::proto_utils::ProtoSchemaType::kMessage,
        "AddItem() can be used only for nested message fields. To write a "
        "primitive field, use traced_proto->set_field() or traced_proto.Set()");
    return Wrap(
        message_->template BeginNestedMessage<
            typename FieldMetadata::cpp_field_type>(FieldMetadata::kFieldId));
  }

  // Write a given |value| into proto  as a new |field| of the current message.
  // This method supports both nested messages and primitive types (i.e. int or
  // string), but requires the |field| to be non-repeateable (i.e. optional).
  // For repeatable fields, AppendValue or AppendFrom should be used.
  template <typename FieldMetadata, typename ValueType>
  void Set(FieldMetadata, ValueType&& value) {
    static_assert(std::is_base_of<MessageType,
                                  typename FieldMetadata::message_type>::value,
                  "Field should belong to the current message");
    static_assert(
        FieldMetadata::kRepetitionType ==
            protozero::proto_utils::RepetitionType::kNotRepeated,
        "Set() can't be used with repeated fields due to ambiguity between "
        "writing |value| as a single entry or treating |value| as a container "
        "and writing all contained items as multiple entries. Please use "
        "dedicated AppendValue() or AppendFrom() methods to differentiate "
        "between "
        "these two situations");

    internal::TypedProtoWriterImpl<
        FieldMetadata,
        FieldMetadata::kProtoFieldType ==
            protozero::proto_utils::ProtoSchemaType::kMessage,
        protozero::proto_utils::RepetitionType::kNotRepeated>::
        Write(*this, std::forward<ValueType>(value));
  }

  // Write a given |value| a single entry into the repeated |field| of the
  // current message. If the field is not repeated, Set() should be used
  // instead.
  template <typename FieldMetadata, typename ValueType>
  void AppendValue(FieldMetadata, ValueType&& value) {
    static_assert(std::is_base_of<MessageType,
                                  typename FieldMetadata::message_type>::value,
                  "Field should belong to the current message");
    static_assert(
        FieldMetadata::kRepetitionType ==
            protozero::proto_utils::RepetitionType::kRepeatedNotPacked,
        "Append*() methods can be used only with repeated fields. "
        "Please use Set() for non-repeated");

    // Write a single value into a given repeated field by explicitly passing
    // "kNotRepeated" to the TypedProtoWriterImpl.
    internal::TypedProtoWriterImpl<
        FieldMetadata,
        FieldMetadata::kProtoFieldType ==
            protozero::proto_utils::ProtoSchemaType::kMessage,
        protozero::proto_utils::RepetitionType::kNotRepeated>::
        Write(*this, std::forward<ValueType>(value));
  }

  // Write a given |value| as a set of entries into the repeated |field| of the
  // current message. If the field is not repeated, Set() should be used
  // instead.
  template <typename FieldMetadata, typename ValueType>
  void AppendFrom(FieldMetadata, ValueType&& value) {
    static_assert(std::is_base_of<MessageType,
                                  typename FieldMetadata::message_type>::value,
                  "Field should belong to the current message");
    static_assert(
        FieldMetadata::kRepetitionType ==
            protozero::proto_utils::RepetitionType::kRepeatedNotPacked,
        "Append*() methods can be used only with repeated fields. "
        "Please use Set() for non-repeated");

    internal::TypedProtoWriterImpl<
        FieldMetadata,
        FieldMetadata::kProtoFieldType ==
            protozero::proto_utils::ProtoSchemaType::kMessage,
        protozero::proto_utils::RepetitionType::kRepeatedNotPacked>::
        Write(*this, std::forward<ValueType>(value));
  }

  // Write a nested message into a field according to the provided metadata.
  // TODO(altimin): Replace the current usages in Chrome with the functions
  // above and make these methods private.
  template <typename FieldMetadata>
  TracedProto<typename FieldMetadata::cpp_field_type> WriteNestedMessage() {
    return WriteNestedMessage(FieldMetadata());
  }

 private:
  friend class EventContext;
  friend class TracedValue;
  // Allow TracedProto<Foo> to create TracedProto<Bar>.
  template <typename T>
  friend class TracedProto;

  // Wraps a raw protozero message using the same context as the current object.
  template <typename ChildMessageType>
  TracedProto<ChildMessageType> Wrap(ChildMessageType* message) {
    return TracedProto<ChildMessageType>(message, context_);
  }

  // Context might be null here when writing typed message which is
  // nested into untyped legacy trace event macro argument.
  // TODO(altimin): Turn this into EventContext& when this case is eliminated
  // and expose it in public API.
  EventContext* context() const { return context_; }

  TracedProto(MessageType* message, EventContext* context)
      : message_(message), context_(context) {}

  MessageType* const message_;
  EventContext* context_;
};

template <typename MessageType, typename ValueType>
void WriteIntoTracedProto(TracedProto<MessageType> message, ValueType&& value);

namespace internal {

template <typename FieldMetadata,
          bool is_message,
          protozero::proto_utils::RepetitionType repetition_type>
struct TypedProtoWriterImpl;

// Simple non-repeated field.
template <typename FieldMetadata>
struct TypedProtoWriterImpl<
    FieldMetadata,
    /*is_message=*/false,
    protozero::proto_utils::RepetitionType::kNotRepeated> {
  template <typename Proto, typename ValueType>
  static void Write(TracedProto<Proto>& context, ValueType&& value) {
    protozero::internal::FieldWriter<FieldMetadata::kProtoFieldType>::Append(
        *context.message(), FieldMetadata::kFieldId, value);
  }
};

// Simple repeated non-packed field.
template <typename FieldMetadata>
struct TypedProtoWriterImpl<
    FieldMetadata,
    /*is_message=*/false,
    protozero::proto_utils::RepetitionType::kRepeatedNotPacked> {
  template <typename Proto, typename ValueType>
  static void Write(TracedProto<Proto>& context, ValueType&& value) {
    for (auto&& item : value) {
      protozero::internal::FieldWriter<FieldMetadata::kProtoFieldType>::Append(
          *context.message(), FieldMetadata::kFieldId, item);
    }
  }
};

// Nested repeated non-packed field.
template <typename FieldMetadata>
struct TypedProtoWriterImpl<
    FieldMetadata,
    /*is_message=*/true,
    protozero::proto_utils::RepetitionType::kNotRepeated> {
  template <typename Proto, typename ValueType>
  static void Write(TracedProto<Proto>& context, ValueType&& value) {
    WriteIntoTracedProto(context.template WriteNestedMessage<FieldMetadata>(),
                         std::forward<ValueType>(value));
  }
};

// Nested repeated non-packed field.
template <typename FieldMetadata>
struct TypedProtoWriterImpl<
    FieldMetadata,
    /*is_message=*/true,
    protozero::proto_utils::RepetitionType::kRepeatedNotPacked> {
  template <typename Proto, typename ValueType>
  static void Write(TracedProto<Proto>& context, ValueType&& value) {
    for (auto&& item : value) {
      WriteIntoTracedProto(context.template WriteNestedMessage<FieldMetadata>(),
                           item);
    }
  }
};

constexpr int kMaxWriteTracedProtoImplPriority = 1;

// If perfetto::TraceFormatTraits<T>::WriteIntoTrace(TracedProto<MessageType>,
// T) is available, use it.
template <typename MessageType, typename T>
decltype(TraceFormatTraits<base::remove_cvref_t<T>>::WriteIntoTrace(
             std::declval<TracedProto<MessageType>>(),
             std::declval<T>()),
         void())
WriteIntoTracedProtoImpl(base::priority_tag<1>,
                         TracedProto<MessageType> message,
                         T&& value) {
  TraceFormatTraits<base::remove_cvref_t<T>>::WriteIntoTrace(
      std::move(message), std::forward<T>(value));
}

// If T has WriteIntoTrace(TracedProto<MessageType>) method, use it.
template <typename MessageType, typename T>
decltype(
    std::declval<T>().WriteIntoTrace(std::declval<TracedProto<MessageType>>()),
    void())
WriteIntoTracedProtoImpl(base::priority_tag<0>,
                         TracedProto<MessageType> message,
                         T&& value) {
  value.WriteIntoTrace(std::move(message));
}

// TypedProtoWriter takes the protozero message (TracedProto<MessageType>),
// field description (FieldMetadata) and value and writes the given value
// into the given field of the given protozero message.
//
// This is primarily used for inline writing of typed messages:
// TRACE_EVENT(..., pbzero::Message:kField, value);
//
// Ideally we would use a function here and not a struct, but passing template
// arguments directly to the function (e.g. foo<void>()) isn't supported until
// C++20, so we have to use a helper struct here.
template <typename FieldMetadata>
struct TypedProtoWriter {
 private:
  using ProtoSchemaType = protozero::proto_utils::ProtoSchemaType;
  using RepetitionType = protozero::proto_utils::RepetitionType;

  static_assert(FieldMetadata::kRepetitionType !=
                    RepetitionType::kRepeatedPacked,
                "writing packed fields isn't supported yet");

  template <bool is_message, RepetitionType repetition_type>
  struct Writer;

 public:
  template <typename Proto, typename ValueType>
  static void Write(TracedProto<Proto>& context, ValueType&& value) {
    TypedProtoWriterImpl<
        FieldMetadata,
        FieldMetadata::kProtoFieldType == ProtoSchemaType::kMessage,
        FieldMetadata::kRepetitionType>::Write(context,
                                               std::forward<ValueType>(value));
  }
};

}  // namespace internal

// Helper template to determine if a given type can be passed to
// perfetto::WriteIntoTracedProto. These templates will fail to resolve if the
// class does not have necesary support, so they are useful for SFINAE and for
// producing helpful compiler error messages.
template <typename MessageType, typename ValueType, typename Result = void>
using check_traced_proto_support_t =
    decltype(internal::WriteIntoTracedProtoImpl(
        std::declval<
            base::priority_tag<internal::kMaxWriteTracedProtoImplPriority>>(),
        std::declval<TracedProto<MessageType>>(),
        std::declval<ValueType>()));

// check_traced_proto_support<MessageType, T, V>::type is defined (and equal to
// V) iff T supports being passed to WriteIntoTracedProto together with
// TracedProto<MessageType>. See the comment in traced_value_forward.h for more
// details.
template <typename MessageType, typename ValueType, class Result>
struct check_traced_proto_support<
    MessageType,
    ValueType,
    Result,
    check_traced_proto_support_t<MessageType, ValueType, Result>> {
  static constexpr bool value = true;
  using type = Result;
};

template <typename MessageType, typename ValueType>
void WriteIntoTracedProto(TracedProto<MessageType> message, ValueType&& value) {
  // TODO(altimin): Add a URL to the documentation and a list of common failure
  // patterns.
  static_assert(
      std::is_same<check_traced_proto_support_t<MessageType, ValueType>,
                   void>::value,
      "The provided type does not support being serialised into the "
      "provided protozero message. Please see the comment in traced_proto.h "
      "for more details.");

  internal::WriteIntoTracedProtoImpl(
      base::priority_tag<internal::kMaxWriteTracedProtoImplPriority>(),
      std::move(message), std::forward<ValueType>(value));
}

template <typename MessageType, typename FieldMetadataType, typename ValueType>
void WriteTracedProtoField(TracedProto<MessageType>& message,
                           FieldMetadataType,
                           ValueType&& value) {
  static_assert(
      std::is_base_of<protozero::proto_utils::FieldMetadataBase,
                      FieldMetadataType>::value,
      "Field name should be a protozero::internal::FieldMetadata<...>");
  static_assert(
      std::is_base_of<MessageType,
                      typename FieldMetadataType::message_type>::value,
      "Field's parent type should match the context.");

  internal::TypedProtoWriter<FieldMetadataType>::Write(
      message, std::forward<ValueType>(value));
}

}  // namespace perfetto

#endif  // INCLUDE_PERFETTO_TRACING_TRACED_PROTO_H_