// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: 2019-2024 Second State INC

//===-- wasmedge/vm/vm.h - VM execution flow class definition -------------===//
//
// Part of the WasmEdge Project.
//
//===----------------------------------------------------------------------===//
///
/// \file
/// This file is the definition class of VM class.
///
//===----------------------------------------------------------------------===//
#pragma once

#include "common/async.h"
#include "common/configure.h"
#include "common/errcode.h"
#include "common/filesystem.h"
#include "common/types.h"

#include "executor/executor.h"
#include "loader/loader.h"
#include "validator/validator.h"

#include "runtime/instance/module.h"
#include "runtime/storemgr.h"

#include <cstdint>
#include <memory>
#include <shared_mutex>
#include <string>
#include <string_view>
#include <unordered_map>
#include <utility>
#include <vector>

namespace WasmEdge {
namespace VM {

/// VM execution flow class
class VM {
public:
  VM() = delete;
  VM(const Configure &Conf);
  VM(const Configure &Conf, Runtime::StoreManager &S);
  ~VM() = default;

  /// ======= Functions can be called before instantiated stage. =======
  /// Register wasm modules and host modules.
  Expect<void> registerModule(std::string_view Name,
                              const std::filesystem::path &Path) {
    std::unique_lock Lock(Mutex);
    return unsafeRegisterModule(Name, Path);
  }
  Expect<void> registerModule(std::string_view Name, Span<const Byte> Code) {
    std::unique_lock Lock(Mutex);
    return unsafeRegisterModule(Name, Code);
  }
  Expect<void> registerModule(std::string_view Name,
                              const AST::Module &Module) {
    std::unique_lock Lock(Mutex);
    return unsafeRegisterModule(Name, Module);
  }
  Expect<void>
  registerModule(const Runtime::Instance::ModuleInstance &ModInst) {
    std::unique_lock Lock(Mutex);
    return unsafeRegisterModule(ModInst);
  }

  /// Rapidly load, validate, instantiate, and run wasm function.
  Expect<std::vector<std::pair<ValVariant, ValType>>>
  runWasmFile(const std::filesystem::path &Path, std::string_view Func,
              Span<const ValVariant> Params = {},
              Span<const ValType> ParamTypes = {}) {
    std::unique_lock Lock(Mutex);
    return unsafeRunWasmFile(Path, Func, Params, ParamTypes);
  }
  Expect<std::vector<std::pair<ValVariant, ValType>>>
  runWasmFile(Span<const Byte> Code, std::string_view Func,
              Span<const ValVariant> Params = {},
              Span<const ValType> ParamTypes = {}) {
    std::unique_lock Lock(Mutex);
    return unsafeRunWasmFile(Code, Func, Params, ParamTypes);
  }
  Expect<std::vector<std::pair<ValVariant, ValType>>>
  runWasmFile(const AST::Module &Module, std::string_view Func,
              Span<const ValVariant> Params = {},
              Span<const ValType> ParamTypes = {}) {
    std::unique_lock Lock(Mutex);
    return unsafeRunWasmFile(Module, Func, Params, ParamTypes);
  }

  Async<Expect<std::vector<std::pair<ValVariant, ValType>>>>
  asyncRunWasmFile(const std::filesystem::path &Path, std::string_view Func,
                   Span<const ValVariant> Params = {},
                   Span<const ValType> ParamTypes = {});
  Async<Expect<std::vector<std::pair<ValVariant, ValType>>>>
  asyncRunWasmFile(Span<const Byte> Code, std::string_view Func,
                   Span<const ValVariant> Params = {},
                   Span<const ValType> ParamTypes = {});
  Async<Expect<std::vector<std::pair<ValVariant, ValType>>>>
  asyncRunWasmFile(const AST::Module &Module, std::string_view Func,
                   Span<const ValVariant> Params = {},
                   Span<const ValType> ParamTypes = {});

  /// Let runtimeTool can check which arguments should be prepared.
  bool holdsModule() {
    if (ActiveModInst) {
      return true;
    }
    return false;
  }
  bool holdsComponent() {
    if (ActiveCompInst) {
      return true;
    }
    return false;
  }

  /// Load given wasm file, wasm bytecode, or wasm module.
  Expect<void> loadWasm(const std::filesystem::path &Path) {
    std::unique_lock Lock(Mutex);
    return unsafeLoadWasm(Path);
  }
  Expect<void> loadWasm(Span<const Byte> Code) {
    std::unique_lock Lock(Mutex);
    return unsafeLoadWasm(Code);
  }
  Expect<void> loadWasm(const AST::Module &Module) {
    std::unique_lock Lock(Mutex);
    return unsafeLoadWasm(Module);
  }

  /// ======= Functions can be called after loaded stage. =======
  /// Validate loaded wasm module.
  Expect<void> validate() {
    std::unique_lock Lock(Mutex);
    return unsafeValidate();
  }

  /// ======= Functions can be called after validated stage. =======
  /// Instantiate validated wasm module.
  Expect<void> instantiate() {
    std::unique_lock Lock(Mutex);
    return unsafeInstantiate();
  }

  /// ======= Functions can be called after instantiated stage. =======
  /// Execute wasm with given input.
  Expect<std::vector<std::pair<ValVariant, ValType>>>
  execute(std::string_view Func, Span<const ValVariant> Params = {},
          Span<const ValType> ParamTypes = {}) {
    std::shared_lock Lock(Mutex);
    return unsafeExecute(Func, Params, ParamTypes);
  }

  /// Execute function of registered module with given input.
  Expect<std::vector<std::pair<ValVariant, ValType>>>
  execute(std::string_view ModName, std::string_view Func,
          Span<const ValVariant> Params = {},
          Span<const ValType> ParamTypes = {}) {
    std::shared_lock Lock(Mutex);
    return unsafeExecute(ModName, Func, Params, ParamTypes);
  }

  /// Execute component function with given input.
  Expect<std::vector<std::pair<ComponentValVariant, ComponentValType>>>
  executeComponent(std::string_view Func,
                   Span<const ComponentValVariant> Params = {},
                   Span<const ComponentValType> ParamTypes = {}) {
    std::shared_lock Lock(Mutex);
    return unsafeExecuteComponent(Func, Params, ParamTypes);
  }

  /// Execute function of registered component with given input.
  Expect<std::vector<std::pair<ComponentValVariant, ComponentValType>>>
  executeComponent(std::string_view CompName, std::string_view Func,
                   Span<const ComponentValVariant> Params = {},
                   Span<const ComponentValType> ParamTypes = {}) {
    std::shared_lock Lock(Mutex);
    return unsafeExecuteComponent(CompName, Func, Params, ParamTypes);
  }

  /// Asynchronous execute wasm with given input.
  Async<Expect<std::vector<std::pair<ValVariant, ValType>>>>
  asyncExecute(std::string_view Func, Span<const ValVariant> Params = {},
               Span<const ValType> ParamTypes = {});

  /// Asynchronous execute function of registered module with given input.
  Async<Expect<std::vector<std::pair<ValVariant, ValType>>>>
  asyncExecute(std::string_view ModName, std::string_view Func,
               Span<const ValVariant> Params = {},
               Span<const ValType> ParamTypes = {});

  /// Asynchronous execute component function with given input.
  Async<Expect<std::vector<std::pair<ComponentValVariant, ComponentValType>>>>
  asyncExecuteComponent(std::string_view Func,
                        Span<const ComponentValVariant> Params = {},
                        Span<const ComponentValType> ParamTypes = {});

  /// Asynchronous execute function of registered component with given input.
  Async<Expect<std::vector<std::pair<ComponentValVariant, ComponentValType>>>>
  asyncExecuteComponent(std::string_view ModName, std::string_view Func,
                        Span<const ComponentValVariant> Params = {},
                        Span<const ComponentValType> ParamTypes = {});

  /// Stop execution
  void stop() noexcept { ExecutorEngine.stop(); }

  /// ======= Functions which are stageless. =======
  /// Clean up VM status
  void cleanup() {
    std::unique_lock Lock(Mutex);
    return unsafeCleanup();
  }

  /// Get list of callable functions and corresponding function types.
  std::vector<std::pair<std::string, const AST::FunctionType &>>
  getFunctionList() const {
    std::shared_lock Lock(Mutex);
    return unsafeGetFunctionList();
  }

  /// Get list of callable component functions and corresponding function types.
  std::vector<std::pair<std::string, const AST::Component::FuncType &>>
  getComponentFunctionList() const {
    std::shared_lock Lock(Mutex);
    return unsafeGetComponentFunctionList();
  }

  /// Get pre-registered module instance by configuration.
  Runtime::Instance::ModuleInstance *
  getImportModule(const HostRegistration Type) const {
    std::shared_lock Lock(Mutex);
    return unsafeGetImportModule(Type);
  }

  /// Get current instantiated module instance.
  const Runtime::Instance::ModuleInstance *getActiveModule() const {
    std::shared_lock Lock(Mutex);
    return unsafeGetActiveModule();
  }

  /// Getter of store set in VM.
  Runtime::StoreManager &getStoreManager() noexcept { return StoreRef; }
  const Runtime::StoreManager &getStoreManager() const noexcept {
    return StoreRef;
  }

  /// Getter of loader in VM.
  Loader::Loader &getLoader() noexcept { return LoaderEngine; }

  /// Getter of validator in VM.
  Validator::Validator &getValidator() noexcept { return ValidatorEngine; }

  /// Getter of executor in VM.
  Executor::Executor &getExecutor() noexcept { return ExecutorEngine; }

  /// Getter of statistics.
  Statistics::Statistics &getStatistics() noexcept { return Stat; }

private:
  Expect<void> unsafeRegisterModule(std::string_view Name,
                                    const std::filesystem::path &Path);
  Expect<void> unsafeRegisterModule(std::string_view Name,
                                    Span<const Byte> Code);
  Expect<void> unsafeRegisterModule(std::string_view Name,
                                    const AST::Module &Module);
  Expect<void>
  unsafeRegisterModule(const Runtime::Instance::ModuleInstance &ModInst);

  Expect<std::vector<std::pair<ValVariant, ValType>>>
  unsafeRunWasmFile(const std::filesystem::path &Path, std::string_view Func,
                    Span<const ValVariant> Params = {},
                    Span<const ValType> ParamTypes = {});
  Expect<std::vector<std::pair<ValVariant, ValType>>>
  unsafeRunWasmFile(Span<const Byte> Code, std::string_view Func,
                    Span<const ValVariant> Params = {},
                    Span<const ValType> ParamTypes = {});
  Expect<std::vector<std::pair<ValVariant, ValType>>>
  unsafeRunWasmFile(const AST::Module &Module, std::string_view Func,
                    Span<const ValVariant> Params = {},
                    Span<const ValType> ParamTypes = {});
  Expect<std::vector<std::pair<ValVariant, ValType>>>
  unsafeRunWasmFile(const AST::Component::Component &Component,
                    std::string_view Func, Span<const ValVariant> Params = {},
                    Span<const ValType> ParamTypes = {});

  Expect<void> unsafeLoadWasm(const std::filesystem::path &Path);
  Expect<void> unsafeLoadWasm(Span<const Byte> Code);
  Expect<void> unsafeLoadWasm(const AST::Module &Module);

  Expect<void> unsafeValidate();

  Expect<void> unsafeInstantiate();

  Expect<std::vector<std::pair<ValVariant, ValType>>>
  unsafeExecute(std::string_view Func, Span<const ValVariant> Params = {},
                Span<const ValType> ParamTypes = {});
  Expect<std::vector<std::pair<ValVariant, ValType>>>

  unsafeExecute(std::string_view Mod, std::string_view Func,
                Span<const ValVariant> Params = {},
                Span<const ValType> ParamTypes = {});

  Expect<std::vector<std::pair<ComponentValVariant, ComponentValType>>>
  unsafeExecuteComponent(std::string_view Func,
                         Span<const ComponentValVariant> Params = {},
                         Span<const ComponentValType> ParamTypes = {});

  Expect<std::vector<std::pair<ComponentValVariant, ComponentValType>>>
  unsafeExecuteComponent(std::string_view Comp, std::string_view Func,
                         Span<const ComponentValVariant> Params = {},
                         Span<const ComponentValType> ParamTypes = {});

  void unsafeCleanup();

  std::vector<std::pair<std::string, const AST::FunctionType &>>
  unsafeGetFunctionList() const;

  std::vector<std::pair<std::string, const AST::Component::FuncType &>>
  unsafeGetComponentFunctionList() const;

  Runtime::Instance::ModuleInstance *
  unsafeGetImportModule(const HostRegistration Type) const;

  const Runtime::Instance::ModuleInstance *unsafeGetActiveModule() const;

  enum class VMStage : uint8_t { Inited, Loaded, Validated, Instantiated };

  void unsafeInitVM();
  void unsafeLoadBuiltInHosts();
  void unsafeLoadPlugInHosts();
  void unsafeRegisterBuiltInHosts();
  void unsafeRegisterPlugInHosts();

  /// Helper function for execution.
  Expect<std::vector<std::pair<ValVariant, ValType>>>
  unsafeExecute(const Runtime::Instance::ModuleInstance *ModInst,
                std::string_view Func, Span<const ValVariant> Params = {},
                Span<const ValType> ParamTypes = {});

  Expect<std::vector<std::pair<ComponentValVariant, ComponentValType>>>
  unsafeExecuteComponent(const Runtime::Instance::ComponentInstance *CompInst,
                         std::string_view Func,
                         Span<const ComponentValVariant> Params = {},
                         Span<const ComponentValType> ParamTypes = {});

  /// \name VM environment.
  /// @{
  const Configure Conf;
  Statistics::Statistics Stat;
  VMStage Stage;
  mutable std::shared_mutex Mutex;
  /// @}

  /// \name VM components.
  /// @{
  Loader::Loader LoaderEngine;
  Validator::Validator ValidatorEngine;
  Executor::Executor ExecutorEngine;
  /// @}

  /// \name VM Storage.
  /// @{
  /// Loaded AST module.
  std::unique_ptr<AST::Module> Mod;
  std::unique_ptr<AST::Component::Component> Comp;
  /// Active module instance.
  std::unique_ptr<Runtime::Instance::ModuleInstance> ActiveModInst;
  std::unique_ptr<Runtime::Instance::ComponentInstance> ActiveCompInst;
  /// Registered module instances by user.
  std::vector<std::unique_ptr<Runtime::Instance::ModuleInstance>> RegModInsts;
  /// Built-in module instances mapped to the configurations. For WASI.
  std::unordered_map<HostRegistration,
                     std::unique_ptr<Runtime::Instance::ModuleInstance>>
      BuiltInModInsts;
  /// Loaded module instances from plug-ins.
  std::vector<std::unique_ptr<Runtime::Instance::ModuleInstance>>
      PlugInModInsts;
  std::vector<std::unique_ptr<Runtime::Instance::ComponentInstance>>
      PlugInCompInsts;
  /// Self-owned store (nullptr if an outside store is assigned in constructor).
  std::unique_ptr<Runtime::StoreManager> Store;
  /// Reference to the store.
  Runtime::StoreManager &StoreRef;
  /// @}
};

} // namespace VM
} // namespace WasmEdge
