// Copyright 2018 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef SRC_LIB_CALLBACK_TRACE_CALLBACK_H_
#define SRC_LIB_CALLBACK_TRACE_CALLBACK_H_

#include <lib/fit/defer.h>
#include <lib/fit/function.h>
#include <lib/syslog/cpp/macros.h>
#include <lib/trace/event.h>

#include <cstring>
#include <functional>
#include <utility>

namespace callback {

namespace internal {

template <typename C>
class TracingLambda {
 public:
  TracingLambda(C callback, uint64_t id, const char* category, const char* name,
                const char* callback_name)
      : id_(id),
        category_(category),
        name_(name),
        callback_name_(callback_name),
        callback_(std::move(callback)),
        did_run_or_moved_out_(false),
        trace_enabled_(true) {}

  explicit TracingLambda(C callback)
      : id_(0u),
        category_(nullptr),
        name_(nullptr),
        callback_name_(nullptr),
        callback_(std::move(callback)),
        did_run_or_moved_out_(false),
        trace_enabled_(false) {}

  // Copy constructor so that the resulting callback can be used as an
  // fit::function, but acts as a move constructor. Only the last copy is valid.
  TracingLambda(TracingLambda&& other) noexcept
      : id_(other.id_),
        category_(other.category_),
        name_(other.name_),
        callback_name_(other.callback_name_),
        callback_(std::move(other.callback_)),
        did_run_or_moved_out_(other.did_run_or_moved_out_),
        trace_enabled_(other.trace_enabled_) {
    FX_DCHECK(!other.did_run_or_moved_out_);
    other.did_run_or_moved_out_ = true;
  }

  ~TracingLambda() {
    if (!did_run_or_moved_out_ && trace_enabled_) {
      TRACE_ASYNC_END(category_, name_, id_, "NotRun", true);
    }
  }

  template <typename... ArgType>
  auto operator()(ArgType&&... args) const {
    FX_DCHECK(!did_run_or_moved_out_);
    did_run_or_moved_out_ = true;
    if (trace_enabled_) {
      TRACE_ASYNC_END(category_, name_, id_);
      TRACE_ASYNC_BEGIN(category_, callback_name_, id_);
    }
    auto guard = fit::defer([trace_enabled = trace_enabled_, id = id_, category = category_,
                             callback_name = callback_name_] {
      if (trace_enabled) {
        TRACE_ASYNC_END(category, callback_name, id);
      }
    });

    return callback_(std::forward<ArgType>(args)...);
  }

 private:
  const uint64_t id_;
  const char* const category_;
  const char* const name_;
  const char* const callback_name_;
  mutable C callback_;
  mutable bool did_run_or_moved_out_;
  bool trace_enabled_;

  FXL_DISALLOW_COPY_AND_ASSIGN(TracingLambda);
};

// NOLINT suppresses check about 'args' being unused. It might be used in the
// future if TRACE_CALLBACK (see below) is used providing additional arguments.
template <typename C, typename... ArgType>
auto TraceCallback(C callback, const char* category, const char* name, const char* callback_name,
                   ArgType... args) {  // NOLINT
  uint64_t id = TRACE_NONCE();
  TRACE_ASYNC_BEGIN(category, name, id, std::forward<ArgType>(args)...);
  return TracingLambda<C>(std::move(callback), id, category, name, callback_name);
}

template <typename C>
auto TraceCallback(C callback) {
  return TracingLambda<C>(std::move(callback));
}

// Identity functions. This is used to ensure that a C string is a compile time
// constant in conjunction with __builtin_strlen.
template <size_t S>
constexpr const char* CheckConstantCString(const char* value) {
  return value;
}

}  // namespace internal
}  // namespace callback

// Wraps the given callback so that it's traced using async tracing from the
// time it's wrapped to the time it completes. Can be used only for callbacks
// that will be called at most once.
#define TRACE_CALLBACK(cb, category, name, args...)                                             \
  ([&]() mutable {                                                                              \
    auto trace_local_cb__ = cb;                                                                 \
    return (TRACE_ENABLED()                                                                     \
                ? ::callback::internal::TraceCallback(                                          \
                      std::move(trace_local_cb__),                                              \
                      ::callback::internal::CheckConstantCString<__builtin_strlen(category)>(   \
                          category),                                                            \
                      ::callback::internal::CheckConstantCString<__builtin_strlen(name)>(name), \
                      name "_callback", ##args)                                                 \
                : ::callback::internal::TraceCallback(std::move(trace_local_cb__)));            \
  }())

#endif  // SRC_LIB_CALLBACK_TRACE_CALLBACK_H_