/*
 * Copyright (C) 2023 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/plugins/stack_functions/stack_functions.h"

#include <stdlib.h>
#include <cstdint>
#include <cstring>
#include <deque>
#include <iterator>
#include <optional>
#include <type_traits>
#include <vector>

#include <memory>

#include "perfetto/base/compiler.h"
#include "perfetto/base/logging.h"
#include "perfetto/base/status.h"
#include "perfetto/ext/base/status_macros.h"
#include "perfetto/ext/base/utils.h"
#include "perfetto/protozero/scattered_heap_buffer.h"
#include "perfetto/trace_processor/basic_types.h"
#include "perfetto/trace_processor/status.h"
#include "protos/perfetto/trace_processor/stack.pbzero.h"
#include "src/trace_processor/core/plugin/plugin.h"
#include "src/trace_processor/perfetto_sql/engine/perfetto_sql_connection.h"
#include "src/trace_processor/sqlite/bindings/sqlite_function.h"
#include "src/trace_processor/sqlite/bindings/sqlite_result.h"
#include "src/trace_processor/sqlite/bindings/sqlite_type.h"
#include "src/trace_processor/sqlite/bindings/sqlite_value.h"
#include "src/trace_processor/sqlite/sqlite_utils.h"
#include "src/trace_processor/storage/trace_storage.h"
#include "src/trace_processor/tables/profiler_tables_py.h"
#include "src/trace_processor/types/trace_processor_context.h"

namespace perfetto {
namespace trace_processor {
namespace {

using protos::pbzero::Stack;

void SetBytesResult(sqlite3_context* ctx, const std::vector<uint8_t>& src) {
  return sqlite::result::TransientBytes(ctx, src.data(),
                                        static_cast<int>(src.size()));
}

// CAT_STACKS(root BLOB/STRING, level_1 BLOB/STRING, …, leaf BLOB/STRING)
// Creates a Stack by concatenating other Stacks. Also accepts strings for which
// it generates a fake Frame
struct CatStacksFunction : public sqlite::Function<CatStacksFunction> {
  static constexpr char kName[] = "CAT_STACKS";
  static constexpr int kArgCount = -1;

  static void Step(sqlite3_context* ctx, int argc, sqlite3_value** argv) {
    PERFETTO_DCHECK(argc >= 0);
    protozero::HeapBuffered<Stack> stack;

    // Note, this SQL function expects the root frame to be the first argument.
    // Stack expects the opposite, thus iterates the args in reverse order.
    for (int i = argc; i > 0; --i) {
      int arg_index = i - 1;
      switch (sqlite::value::Type(argv[arg_index])) {
        case sqlite::Type::kBlob: {
          const void* blob = sqlite::value::Blob(argv[arg_index]);
          int size = sqlite::value::Bytes(argv[arg_index]);
          stack->AppendRawProtoBytes(blob, static_cast<size_t>(size));
          break;
        }
        case sqlite::Type::kText: {
          stack->add_entries()->set_name(sqlite::value::Text(argv[arg_index]));
          break;
        }
        case sqlite::Type::kNull:
          break;
        case sqlite::Type::kInteger:
        case sqlite::Type::kFloat:
          return sqlite::utils::SetError(
              ctx, base::ErrStatus(
                       "CAT_STACKS: entry %d must be BLOB, STRING, or NULL",
                       arg_index));
      }
    }

    return SetBytesResult(ctx, stack.SerializeAsArray());
  }
};

// STACK_FROM_STACK_PROFILE_CALLSITE(callsite_id LONG, [annotate BOOLEAN])
// Creates a stack by taking a callsite_id (reference to the
// stack_profile_callsite table) and generating a list of frames (by walking the
// stack_profile_callsite table)
// Optionally annotates frames (annotate param has a default value of false)
//
// Important: Annotations might interfere with certain aggregations, as we
// will could have a frame that is annotated with different annotations. That
// will lead to multiple functions being generated (same name, line etc, but
// different annotation).
struct StackFromStackProfileCallsiteFunction
    : public sqlite::Function<StackFromStackProfileCallsiteFunction> {
  static constexpr char kName[] = "STACK_FROM_STACK_PROFILE_CALLSITE";
  static constexpr int kArgCount = -1;

  using UserData = TraceStorage;

  static void Step(sqlite3_context* ctx, int argc, sqlite3_value** argv) {
    TraceStorage* storage = GetUserData(ctx);
    PERFETTO_DCHECK(argc == 1 || argc == 2);

    int64_t callsite_id_long = 0;
    switch (sqlite::value::Type(argv[0])) {
      case sqlite::Type::kInteger:
        callsite_id_long = sqlite::value::Int64(argv[0]);
        break;
      case sqlite::Type::kNull:
        return sqlite::utils::ReturnNullFromFunction(ctx);
      case sqlite::Type::kFloat:
      case sqlite::Type::kText:
      case sqlite::Type::kBlob:
        return sqlite::utils::SetError(
            ctx,
            "STACK_FROM_STACK_PROFILE_CALLSITE: callsite_id must be integer");
    }

    auto callsite_id =
        storage->stack_profile_callsite_table().TryCastId(callsite_id_long);
    if (!callsite_id) {
      return sqlite::utils::SetError(
          ctx, base::ErrStatus("STACK_FROM_STACK_PROFILE_CALLSITE: callsite_id "
                               "does not exist: %" PRId64,
                               callsite_id_long));
    }

    bool annotate = false;
    if (argc == 2) {
      switch (sqlite::value::Type(argv[1])) {
        case sqlite::Type::kInteger:
          // true = 1 and false = 0 in SQL
          annotate = (sqlite::value::Int64(argv[1]) != 0);
          break;
        case sqlite::Type::kNull:
          return sqlite::utils::ReturnNullFromFunction(ctx);
        case sqlite::Type::kFloat:
        case sqlite::Type::kText:
        case sqlite::Type::kBlob:
          return sqlite::utils::SetError(
              ctx,
              "STACK_FROM_STACK_PROFILE_CALLSITE: annotate must be integer");
      }
    }

    protozero::HeapBuffered<Stack> stack;
    if (annotate) {
      stack->add_entries()->set_annotated_callsite_id(callsite_id->value);
    } else {
      stack->add_entries()->set_callsite_id(callsite_id->value);
    }
    return SetBytesResult(ctx, stack.SerializeAsArray());
  }
};

// STACK_FROM_STACK_PROFILE_FRAME(frame_id LONG)
// Creates a stack with just the frame referenced by frame_id (reference to the
// stack_profile_frame table)
struct StackFromStackProfileFrameFunction
    : public sqlite::Function<StackFromStackProfileFrameFunction> {
  static constexpr char kName[] = "STACK_FROM_STACK_PROFILE_FRAME";
  static constexpr int kArgCount = 1;

  using UserData = TraceStorage;

  static void Step(sqlite3_context* ctx, int argc, sqlite3_value** argv) {
    TraceStorage* storage = GetUserData(ctx);
    PERFETTO_DCHECK(argc == 1);

    int64_t frame_id_long = 0;
    switch (sqlite::value::Type(argv[0])) {
      case sqlite::Type::kInteger:
        frame_id_long = sqlite::value::Int64(argv[0]);
        break;
      case sqlite::Type::kNull:
        return sqlite::utils::ReturnNullFromFunction(ctx);
      case sqlite::Type::kFloat:
      case sqlite::Type::kText:
      case sqlite::Type::kBlob:
        return sqlite::utils::SetError(
            ctx, "STACK_FROM_STACK_PROFILE_FRAME: frame_id must be integer");
    }

    auto frame_id =
        storage->stack_profile_frame_table().TryCastId(frame_id_long);
    if (!frame_id) {
      return sqlite::utils::SetError(
          ctx, base::ErrStatus("STACK_FROM_STACK_PROFILE_FRAME: frame_id does "
                               "not exist: %" PRId64,
                               frame_id_long));
    }

    protozero::HeapBuffered<Stack> stack;
    stack->add_entries()->set_frame_id(frame_id->value);
    return SetBytesResult(ctx, stack.SerializeAsArray());
  }
};

}  // namespace

namespace stack_functions {
namespace {

class StackFunctionsPlugin : public Plugin<StackFunctionsPlugin> {
 public:
  ~StackFunctionsPlugin() override;

  void RegisterFunctions(PerfettoSqlConnection*,
                         std::vector<FunctionRegistration>& out) override {
    TraceStorage* storage = trace_context_->storage.get();
    out.push_back(MakeFunctionRegistration<CatStacksFunction>(nullptr));
    out.push_back(
        MakeFunctionRegistration<StackFromStackProfileFrameFunction>(storage));
    out.push_back(
        MakeFunctionRegistration<StackFromStackProfileCallsiteFunction>(
            storage));
  }
};

StackFunctionsPlugin::~StackFunctionsPlugin() = default;

}  // namespace

void RegisterPlugin() {
  static PluginRegistration reg(
      []() -> std::unique_ptr<PluginBase> {
        return std::make_unique<StackFunctionsPlugin>();
      },
      StackFunctionsPlugin::kPluginId, StackFunctionsPlugin::kDepIds.data(),
      StackFunctionsPlugin::kDepIds.size());
  base::ignore_result(reg);
}

}  // namespace stack_functions

}  // namespace trace_processor
}  // namespace perfetto
