/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file tvm/runtime/vm/vm.h
 */
#ifndef TVM_RUNTIME_VM_VM_H_
#define TVM_RUNTIME_VM_VM_H_

#include <tvm/ffi/extra/module.h>

#include <memory>
#include <string>
#include <unordered_map>
#include <vector>

#include "../memory/memory_manager.h"
#include "./bytecode.h"
#include "./executable.h"

namespace tvm {
namespace runtime {

using memory::Allocator;
using memory::AllocatorType;
using memory::MemoryManager;
using memory::Storage;
using memory::StorageObj;

namespace vm {

/*!
 * \brief Possible instrument actions.
 */
enum class VMInstrumentReturnKind : int {
  /*! \brief Running as normal. */
  kNoOp = 0,
  /*! \brief Skip the following run, only valid in before. */
  kSkipRun = 1,
};

/*!
 * \brief An object representing a vm closure.
 */
class VMClosureObj : public ffi::Object {
 public:
  /*!
   * \brief The function name. The function could be any
   * function object that is compatible to the VM runtime.
   */
  ffi::String func_name;

  /*!
   * \brief The implementation of the Closure.
   * \note This function takes context pointer(VirtualMachine*)
   *       as the first argument. The rest of arguments follows
   *       the same arguments as the normal function call.
   */
  ffi::Function impl;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.vm.Closure", VMClosureObj, ffi::Object);
};

/*! \brief reference to closure. */
class VMClosure : public ffi::ObjectRef {
 public:
  VMClosure(ffi::String func_name, ffi::Function impl);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VMClosure, ffi::ObjectRef, VMClosureObj);

  /*!
   * \brief Create another ffi::Function with last arguments already bound to last_args.
   *
   * This is a helper function to create captured closures.
   * \param func The input func, can be a VMClosure or ffi::Function.
   * \param last_args The arguments to bound to in the end of the function.
   * \note The new function takes in arguments and append the last_args in the end.
   */
  static ffi::Function BindLastArgs(ffi::Function func, std::vector<ffi::Any> last_args);
};

/*!
 * \brief Represent a VM extension.
 * A VM extension allows the user to extend the VM with target specific functionalities.
 * The VM holds the reference of the extensions to ensure the extensions have the same lifetime
 * as the VM.
 *
 * This is the base class for all VM extensions and should not be used directly.
 */
class VMExtensionNode : public ffi::Object {
 protected:
  TVM_FFI_DECLARE_OBJECT_INFO("runtime.VMExtension", VMExtensionNode, ffi::Object);
};

/*! \brief Managed reference to VM extension. */
class VMExtension : public ffi::ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VMExtension, ffi::ObjectRef, VMExtensionNode);
};

/*!
 * \brief The virtual machine.
 *
 * The virtual machine contains all the current execution state,
 * as well as the executable.
 *
 * The goal is to have a single self-contained object,
 * enabling one to easily pass around VMs, execute them on
 * multiple threads, or serialize them to disk or over the
 * wire.
 */
class VirtualMachine : public ffi::ModuleObj {
 public:
  /*!
   * \brief Initialize the virtual machine for a set of devices.
   * \param devices The set of TVM devices.
   * \param alloc_types The allocator types for each device.
   */
  virtual void Init(const std::vector<Device>& devices,
                    const std::vector<AllocatorType>& alloc_types) = 0;
  /*!
   * \brief Load the executable for the virtual machine.
   * \param exec The executable.
   */
  virtual void LoadExecutable(ffi::ObjectPtr<VMExecutable> exec) = 0;
  /*!
   * \brief Get global function in the VM.
   * \param func_name The name of the function.
   * \return The closure
   */
  virtual VMClosure GetClosure(const ffi::String& func_name) = 0;
  /*!
   * \brief Invoke closure or packed function using ffi::Function convention.
   * \param closure_or_packedfunc A VM closure or a packed_func.
   * \param args The input arguments.
   * \param rv The return value.
   */
  virtual void InvokeClosurePacked(const ffi::ObjectRef& closure_or_packedfunc,
                                   ffi::PackedArgs args, ffi::Any* rv) = 0;
  /*!
   * \brief Set an instrumentation function.
   *
   * If instrument is present, the function will be called
   * before/after each Call instruction.
   *
   * bool instrument(func, func_symbol, before_run, args...)
   *
   * - func: Union[VMClosure, ffi::Function], the function object.
   * - func_symbol: string, the symbol of the function.
   * - before_run: bool, whether it is before or after call.
   * - ret_value: Only valid in after run, otherwise it is null.
   * - args: the arguments being passed to call.
   *
   * instrument can return an int which corresponds to the action value.
   * \sa VMInstrumentAction
   *
   * \param instrument The instrument function.
   */
  virtual void SetInstrument(ffi::Function instrument) = 0;

  /*!
   * \brief Get or create a VM extension. Once created, the extension will be stored in the VM
   * and held until the VM is destructed.
   *
   * \tparam T The type of the extension
   * \return The extension instance
   */
  template <typename T, typename = std::enable_if_t<std::is_base_of<VMExtension, T>::value>>
  T GetOrCreateExtension() {
    using ContainerType = typename T::ContainerType;
    uint32_t key = ContainerType::RuntimeTypeIndex();
    if (auto it = extensions.find(key); it != extensions.end()) {
      ffi::Any value = (*it).second;
      return value.cast<T>();
    }
    auto [it, _] = extensions.emplace(key, T::Create());
    ffi::Any value = (*it).second;
    return value.cast<T>();
  }

  /*!
   * \brief Create a specific instance of VM.
   * \return Created VM
   */
  static ffi::ObjectPtr<VirtualMachine> Create();
  /*!
   * \brief Helper function for vm closure functions to get the context ptr
   * \param arg The argument value.
   */
  static VirtualMachine* GetContextPtr(ffi::AnyView arg) {
    return static_cast<VirtualMachine*>(arg.cast<void*>());
  }

  ~VirtualMachine() {}

  //--------------------------------------------------------------------------
  // The following section contains states that other builtin can depend on
  //--------------------------------------------------------------------------
  /*! \brief The memory allocators. */
  std::vector<Allocator*> allocators;
  /*! \brief Runtime physical device list. */
  std::vector<Device> devices;
  /*! \brief The VM extensions. Mapping from the type index of the extension to the extension
   * instance. */
  std::unordered_map<uint32_t, Any> extensions;
};

}  // namespace vm
}  // namespace runtime
}  // namespace tvm

#endif  // TVM_RUNTIME_VM_VM_H_