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

#include "common/configure.h"
#include "common/filesystem.h"
#include "common/spdlog.h"
#include "common/types.h"
#include "common/version.h"
#include "driver/tool.h"
#include "host/wasi/wasimodule.h"
#include "vm/vm.h"

#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <optional>
#include <string>
#include <string_view>
#include <vector>

using namespace std::literals;

namespace WasmEdge {
namespace Driver {

static int
ToolOnModule(WasmEdge::VM::VM &VM, const std::string &FuncName,
             std::optional<std::chrono::system_clock::time_point> Timeout,
             struct DriverToolOptions &Opt,
             const AST::FunctionType &FuncType) noexcept {
  std::vector<ValVariant> FuncArgs;
  std::vector<ValType> FuncArgTypes;

  for (size_t I = 0;
       I < FuncType.getParamTypes().size() && I + 1 < Opt.Args.value().size();
       ++I) {
    const auto TCode = FuncType.getParamTypes()[I].getCode();
    switch (TCode) {
    case TypeCode::I32: {
      const int32_t Value =
          static_cast<int32_t>(std::stol(Opt.Args.value()[I + 1]));
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(TCode);
      break;
    }
    case TypeCode::I64: {
      const int64_t Value =
          static_cast<int64_t>(std::stoll(Opt.Args.value()[I + 1]));
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(TCode);
      break;
    }
    case TypeCode::F32: {
      const float Value = std::stof(Opt.Args.value()[I + 1]);
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(TCode);
      break;
    }
    case TypeCode::F64: {
      const double Value = std::stod(Opt.Args.value()[I + 1]);
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(TCode);
      break;
    }
    // TODO: FuncRef and ExternRef
    default:
      break;
    }
  }
  if (FuncType.getParamTypes().size() + 1 < Opt.Args.value().size()) {
    for (size_t I = FuncType.getParamTypes().size() + 1;
         I < Opt.Args.value().size(); ++I) {
      const uint64_t Value =
          static_cast<uint64_t>(std::stoll(Opt.Args.value()[I]));
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(TypeCode::I64);
    }
  }

  auto AsyncResult = VM.asyncExecute(FuncName, FuncArgs, FuncArgTypes);
  if (Timeout.has_value()) {
    if (!AsyncResult.waitUntil(*Timeout)) {
      AsyncResult.cancel();
    }
  }
  if (auto Result = AsyncResult.get()) {
    // Print results.
    for (size_t I = 0; I < Result->size(); ++I) {
      switch ((*Result)[I].second.getCode()) {
      case TypeCode::I32:
        fmt::print("{}\n"sv, (*Result)[I].first.get<int32_t>());
        break;
      case TypeCode::I64:
        fmt::print("{}\n"sv, (*Result)[I].first.get<int64_t>());
        break;
      case TypeCode::F32:
        fmt::print("{}\n"sv, (*Result)[I].first.get<float>());
        break;
      case TypeCode::F64:
        fmt::print("{}\n"sv, (*Result)[I].first.get<double>());
        break;
      case TypeCode::V128:
        fmt::print("{}\n"sv, uint128((*Result)[I].first.get<uint128_t>()));
        break;
      case TypeCode::Ref: {
        if ((*Result)[I].second.isFuncRefType()) {
          fmt::print("<funcref>\n"sv);
        } else if ((*Result)[I].second.isExternRefType()) {
          fmt::print("<externref>\n"sv);
        } else {
          fmt::print("<anyref>\n"sv);
        }
        break;
      }
      case TypeCode::RefNull: {
        if ((*Result)[I].second.isFuncRefType()) {
          fmt::print("<null funcref>\n"sv);
        } else if ((*Result)[I].second.isExternRefType()) {
          fmt::print("<null externref>\n"sv);
        } else {
          fmt::print("<null anyref>\n"sv);
        }
        break;
      }
      default:
        break;
      }
    }
    return EXIT_SUCCESS;
  } else {
    // It indicates that the execution of wasm has been aborted
    return 128 + SIGABRT;
  }
}

static int
ToolOnComponent(WasmEdge::VM::VM &VM, const std::string &FuncName,
                std::optional<std::chrono::system_clock::time_point> Timeout,
                struct DriverToolOptions &Opt,
                const AST::Component::FuncType &FuncType) noexcept {
  std::vector<ComponentValVariant> FuncArgs;
  std::vector<ComponentValType> FuncArgTypes;

  for (size_t I = 0;
       I < FuncType.getParamList().size() && I + 1 < Opt.Args.value().size();
       ++I) {
    const auto TCode = FuncType.getParamList()[I].getValType().getCode();
    switch (TCode) {
    case ComponentTypeCode::S32:
    case ComponentTypeCode::U32: {
      const uint32_t Value =
          static_cast<uint32_t>(std::stol(Opt.Args.value()[I + 1]));
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(TCode);
      break;
    }
    case ComponentTypeCode::S64:
    case ComponentTypeCode::U64: {
      const uint64_t Value =
          static_cast<uint64_t>(std::stoll(Opt.Args.value()[I + 1]));
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(TCode);
      break;
    }
    case ComponentTypeCode::F32: {
      const float Value = std::stof(Opt.Args.value()[I + 1]);
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(TCode);
      break;
    }
    case ComponentTypeCode::F64: {
      const double Value = std::stod(Opt.Args.value()[I + 1]);
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(TCode);
      break;
    }
    case ComponentTypeCode::String: {
      const std::string Value = Opt.Args.value()[I + 1];
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(TCode);
      break;
    }
    // TODO: COMPONENT - other types.
    default:
      break;
    }
  }
  if (FuncType.getParamList().size() + 1 < Opt.Args.value().size()) {
    for (size_t I = FuncType.getParamList().size() + 1;
         I < Opt.Args.value().size(); ++I) {
      const uint64_t Value =
          static_cast<uint64_t>(std::stoll(Opt.Args.value()[I]));
      FuncArgs.emplace_back(Value);
      FuncArgTypes.emplace_back(ComponentTypeCode::U64);
    }
  }

  auto AsyncResult = VM.asyncExecuteComponent(FuncName, FuncArgs, FuncArgTypes);
  if (Timeout.has_value()) {
    if (!AsyncResult.waitUntil(*Timeout)) {
      AsyncResult.cancel();
    }
  }
  if (auto Result = AsyncResult.get()) {
    // Print results.
    for (auto &&Val : *Result) {
      switch (Val.second.getCode()) {
      case ComponentTypeCode::S32:
        fmt::print("{}\n"sv, std::get<ValVariant>(Val.first).get<int32_t>());
        break;
      case ComponentTypeCode::U32:
        fmt::print("{}\n"sv, std::get<ValVariant>(Val.first).get<uint32_t>());
        break;
      case ComponentTypeCode::S64:
        fmt::print("{}\n"sv, std::get<ValVariant>(Val.first).get<int64_t>());
        break;
      case ComponentTypeCode::U64:
        fmt::print("{}\n"sv, std::get<ValVariant>(Val.first).get<uint64_t>());
        break;
      case ComponentTypeCode::F32:
        fmt::print("{}\n"sv, std::get<ValVariant>(Val.first).get<float>());
        break;
      case ComponentTypeCode::F64:
        fmt::print("{}\n"sv, std::get<ValVariant>(Val.first).get<double>());
        break;
      case ComponentTypeCode::String:
        fmt::print("{}\n"sv, std::get<std::string>(Val.first));
        break;
      default:
        break;
      }
    }

    return EXIT_SUCCESS;
  } else {
    // It indicates that the execution of wasm has been aborted
    return 128 + SIGABRT;
  }
}

int Tool(struct DriverToolOptions &Opt) noexcept {
  std::ios::sync_with_stdio(false);
  Log::setInfoLoggingLevel();

  Configure Conf;
  // WASM standard configuration has the highest priority.
  if (Opt.PropWASM1.value()) {
    Conf.setWASMStandard(Standard::WASM_1);
  }
  if (Opt.PropWASM2.value()) {
    Conf.setWASMStandard(Standard::WASM_2);
  }
  if (Opt.PropWASM3.value()) {
    Conf.setWASMStandard(Standard::WASM_3);
  }

  // Proposals adjustment.
  if (Opt.PropMutGlobals.value()) {
    Conf.removeProposal(Proposal::ImportExportMutGlobals);
  }
  if (Opt.PropNonTrapF2IConvs.value()) {
    Conf.removeProposal(Proposal::NonTrapFloatToIntConversions);
  }
  if (Opt.PropSignExtendOps.value()) {
    Conf.removeProposal(Proposal::SignExtensionOperators);
  }
  if (Opt.PropMultiValue.value()) {
    Conf.removeProposal(Proposal::MultiValue);
  }
  if (Opt.PropBulkMemOps.value()) {
    Conf.removeProposal(Proposal::BulkMemoryOperations);
  }
  if (Opt.PropSIMD.value()) {
    Conf.removeProposal(Proposal::SIMD);
  }
  if (Opt.PropTailCall.value()) {
    Conf.removeProposal(Proposal::TailCall);
  }
  if (Opt.PropExtendConst.value()) {
    Conf.removeProposal(Proposal::ExtendedConst);
  }
  if (Opt.PropMultiMem.value()) {
    Conf.removeProposal(Proposal::MultiMemories);
  }
  if (Opt.PropRelaxedSIMD.value()) {
    Conf.removeProposal(Proposal::RelaxSIMD);
  }
  if (Opt.PropExceptionHandling.value()) {
    Conf.removeProposal(Proposal::ExceptionHandling);
  }
  // TODO: MEMORY64 - enable the option.
  // if (Opt.PropMemory64.value()) {
  //   Conf.removeProposal(Proposal::Memory64);
  // }

  // Handle the proposal removal which has dependency.
  // The GC proposal depends on the func-ref proposal, and the func-ref proposal
  // depends on the ref-types proposal.
  if (Opt.PropGC.value()) {
    Conf.removeProposal(Proposal::GC);
  }
  if (Opt.PropFunctionReference.value()) {
    // This will automatically not work if the GC proposal not disabled.
    Conf.removeProposal(Proposal::FunctionReferences);
  }
  if (Opt.PropRefTypes.value()) {
    // This will automatically not work if the GC or func-ref proposal not
    // disabled.
    Conf.removeProposal(Proposal::ReferenceTypes);
  }

  if (Opt.PropThreads.value()) {
    Conf.addProposal(Proposal::Threads);
  }
  if (Opt.PropComponent.value()) {
    Conf.addProposal(Proposal::Component);
    spdlog::warn("component model is enabled, this is experimental."sv);
  }
  if (Opt.PropAll.value()) {
    Conf.setWASMStandard(Standard::WASM_3);
    Conf.addProposal(Proposal::Threads);
    spdlog::warn("component model is enabled, this is experimental."sv);
    Conf.addProposal(Proposal::Component);
  }

  std::optional<std::chrono::system_clock::time_point> Timeout;
  if (Opt.TimeLim.value() > 0) {
    Timeout = std::chrono::system_clock::now() +
              std::chrono::milliseconds(Opt.TimeLim.value());
  }
  if (Opt.GasLim.value().size() > 0) {
    Conf.getStatisticsConfigure().setCostMeasuring(true);
    Conf.getStatisticsConfigure().setCostLimit(
        static_cast<uint32_t>(Opt.GasLim.value().back()));
  }
  if (Opt.MemLim.value().size() > 0) {
    Conf.getRuntimeConfigure().setMaxMemoryPage(
        static_cast<uint32_t>(Opt.MemLim.value().back()));
  }
  if (Opt.ConfEnableAllStatistics.value()) {
    Conf.getStatisticsConfigure().setInstructionCounting(true);
    Conf.getStatisticsConfigure().setCostMeasuring(true);
    Conf.getStatisticsConfigure().setTimeMeasuring(true);
  } else {
    if (Opt.ConfEnableInstructionCounting.value()) {
      Conf.getStatisticsConfigure().setInstructionCounting(true);
    }
    if (Opt.ConfEnableGasMeasuring.value()) {
      Conf.getStatisticsConfigure().setCostMeasuring(true);
    }
    if (Opt.ConfEnableTimeMeasuring.value()) {
      Conf.getStatisticsConfigure().setTimeMeasuring(true);
    }
  }
  if (Opt.ConfEnableJIT.value()) {
    Conf.getRuntimeConfigure().setEnableJIT(true);
    Conf.getCompilerConfigure().setOptimizationLevel(
        WasmEdge::CompilerConfigure::OptimizationLevel::O1);
  }
  if (Opt.ConfEnableCoredump.value()) {
    Conf.getRuntimeConfigure().setEnableCoredump(true);
  }
  if (Opt.ConfCoredumpWasmgdb.value()) {
    Conf.getRuntimeConfigure().setCoredumpWasmgdb(true);
  }
  if (Opt.ConfForceInterpreter.value()) {
    Conf.getRuntimeConfigure().setForceInterpreter(true);
  }
  if (Opt.ConfAFUNIX.value()) {
    Conf.getRuntimeConfigure().setAllowAFUNIX(true);
  }

  for (const auto &Name : Opt.ForbiddenPlugins.value()) {
    Conf.addForbiddenPlugins(Name);
  }

  Conf.addHostRegistration(HostRegistration::Wasi);
  const auto InputPath =
      std::filesystem::absolute(std::filesystem::u8path(Opt.SoName.value()));

  // Create VM and get WASI module instance.
  VM::VM VM(Conf);
  Host::WasiModule *WasiMod = dynamic_cast<Host::WasiModule *>(
      VM.getImportModule(HostRegistration::Wasi));

  // Load, validate, and instantiate WASM or Component.
  if (auto Result = VM.loadWasm(InputPath.u8string()); !Result) {
    return EXIT_FAILURE;
  }
  if (auto Result = VM.validate(); !Result) {
    return EXIT_FAILURE;
  }
  if (auto Result = VM.instantiate(); !Result) {
    return EXIT_FAILURE;
  }

  auto HasValidCommandModStartFunc = [&]() {
    bool HasStart = false;
    bool Valid = false;

    auto Functions = VM.getFunctionList();
    for (auto &[FuncName, Type] : Functions) {
      if (FuncName == "_start") {
        HasStart = true;
        if (Type.getReturnTypes().size() == 0 &&
            Type.getParamTypes().size() == 0) {
          Valid = true;
          break;
        }
      }
    }

    // if HasStart but not Valid, insert _start to enter reactor mode
    if (HasStart && !Valid) {
      Opt.Args.value().insert(Opt.Args.value().begin(), "_start");
    }

    return HasStart && Valid;
  };

  // TODO: COMPONENT - does component start function named as "_start"?
  bool EnterCommandMode = !Opt.Reactor.value() && HasValidCommandModStartFunc();

  // Initialize WASI module.
  WasiMod->init(Opt.Dir.value(),
                InputPath.filename()
                    .replace_extension(std::filesystem::u8path("wasm"sv))
                    .u8string(),
                Opt.Args.value(), Opt.Env.value());

  if (EnterCommandMode) {
    // command mode

    // TODO: COMPONENT - currently not supported.
    auto AsyncResult = VM.asyncExecute("_start"sv);
    if (Timeout.has_value()) {
      if (!AsyncResult.waitUntil(*Timeout)) {
        AsyncResult.cancel();
      }
    }
    if (auto Result = AsyncResult.get();
        Result || Result.error() == ErrCode::Value::Terminated) {
      return static_cast<int>(WasiMod->getExitCode());
    } else {
      // It indicates that the execution of wasm has been aborted
      return 128 + SIGABRT;
    }
  } else {
    // reactor mode

    // Get the function name to invoke.
    if (Opt.Args.value().empty()) {
      fmt::print(
          stderr,
          "A function name is required when reactor mode is enabled.\n"sv);
      return EXIT_FAILURE;
    }
    const auto &FuncName = Opt.Args.value().front();

    if (VM.holdsModule()) {
      // WASM case.

      // Check the exported function name and function type first.
      const auto InitFunc = "_initialize"s;
      bool HasInit = false;
      const AST::FunctionType *FuncType = nullptr;
      for (const auto &Func : VM.getFunctionList()) {
        if (Func.first == InitFunc) {
          // Found the init function.
          HasInit = true;
        } else if (Func.first == FuncName) {
          // Found the function to invoke.
          FuncType = &Func.second;
        }
      }

      // If found initialize function, invoke it first.
      if (HasInit) {
        auto AsyncResult = VM.asyncExecute(InitFunc);
        if (Timeout.has_value()) {
          if (!AsyncResult.waitUntil(*Timeout)) {
            AsyncResult.cancel();
          }
        }
        if (auto Result = AsyncResult.get(); unlikely(!Result)) {
          // It indicates that the execution of wasm has been aborted.
          return 128 + SIGABRT;
        }
      }
      return ToolOnModule(VM, FuncName, Timeout, Opt, *FuncType);
    } else if (VM.holdsComponent()) {
      // Component case.

      // Check the exported function name and function type first.
      const AST::Component::FuncType *FuncType = nullptr;
      for (const auto &Func : VM.getComponentFunctionList()) {
        if (Func.first == FuncName) {
          // Found the function to invoke.
          FuncType = &Func.second;
        }
      }
      // TODO: COMPONENT - Check the exported function name and function type
      // first.
      return ToolOnComponent(VM, FuncName, Timeout, Opt, *FuncType);
    } else {
      // which means VM has neither instantiated module nor instantiated
      // component
      return 128 + SIGABRT;
    }
  }
}

} // namespace Driver
} // namespace WasmEdge
