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

#include "loader/aot_section.h"
#include "loader/loader.h"
#include "loader/shared_library.h"

#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

using namespace std::literals;

namespace WasmEdge {
namespace Loader {

// Load binary to construct Module node. See "include/loader/loader.h".
Expect<void> Loader::loadModule(AST::Module &Mod,
                                std::optional<uint64_t> Bound) {
  uint64_t StartOffset = FMgr.getOffset();
  auto ReportError = [](auto E) {
    spdlog::error(ErrInfo::InfoAST(ASTNodeAttr::Module));
    return E;
  };

  // Variables to record the loaded section types.
  HasDataSection = false;
  std::vector<uint8_t> Secs = {0x0BU, 0x0AU, 0x0CU, 0x09U, 0x08U, 0x07U, 0x06U,
                               0x0DU, 0x05U, 0x04U, 0x03U, 0x02U, 0x01U};
  uint64_t Offset = FMgr.getOffset();

  // Read Section index and create Section nodes.
  while (!Bound.has_value() || Bound.value() > Offset - StartOffset) {
    uint8_t NewSectionId = 0x00;
    // If not read section ID, seems the end of file and break.
    if (auto Res = FMgr.readByte()) {
      NewSectionId = *Res;
    } else {
      if (Res.error() == ErrCode::Value::UnexpectedEnd) {
        break;
      } else {
        return logLoadError(Res.error(), FMgr.getLastOffset(),
                            ASTNodeAttr::Module);
      }
    }

    // Sections except the custom section should be unique and in order.
    if (NewSectionId > 0x00U && NewSectionId < 0x0EU) {
      while (Secs.size() > 0 && Secs.back() != NewSectionId) {
        Secs.pop_back();
      }
      if (Secs.empty()) {
        return logLoadError(ErrCode::Value::JunkSection, FMgr.getLastOffset(),
                            ASTNodeAttr::Module);
      }
      Secs.pop_back();
    }

    switch (NewSectionId) {
    case 0x00:
      Mod.getCustomSections().emplace_back();
      EXPECTED_TRY(
          loadSection(Mod.getCustomSections().back()).map_error(ReportError));
      break;
    case 0x01:
      EXPECTED_TRY(loadSection(Mod.getTypeSection()).map_error(ReportError));
      break;
    case 0x02:
      EXPECTED_TRY(loadSection(Mod.getImportSection()).map_error(ReportError));
      break;
    case 0x03:
      EXPECTED_TRY(
          loadSection(Mod.getFunctionSection()).map_error(ReportError));
      break;
    case 0x04:
      EXPECTED_TRY(loadSection(Mod.getTableSection()).map_error(ReportError));
      break;
    case 0x05:
      EXPECTED_TRY(loadSection(Mod.getMemorySection()).map_error(ReportError));
      break;
    case 0x06:
      EXPECTED_TRY(loadSection(Mod.getGlobalSection()).map_error(ReportError));
      break;
    case 0x07:
      EXPECTED_TRY(loadSection(Mod.getExportSection()).map_error(ReportError));
      break;
    case 0x08:
      EXPECTED_TRY(loadSection(Mod.getStartSection()).map_error(ReportError));
      break;
    case 0x09:
      EXPECTED_TRY(loadSection(Mod.getElementSection()).map_error(ReportError));
      break;
    case 0x0A:
      EXPECTED_TRY(loadSection(Mod.getCodeSection()).map_error(ReportError));
      break;
    case 0x0B:
      EXPECTED_TRY(loadSection(Mod.getDataSection()).map_error(ReportError));
      break;
    case 0x0C:
      // This section is for BulkMemoryOperations or ReferenceTypes proposal.
      if (!Conf.hasProposal(Proposal::BulkMemoryOperations) &&
          !Conf.hasProposal(Proposal::ReferenceTypes)) {
        return logNeedProposal(ErrCode::Value::MalformedSection,
                               Proposal::BulkMemoryOperations,
                               FMgr.getLastOffset(), ASTNodeAttr::Module);
      }
      EXPECTED_TRY(
          loadSection(Mod.getDataCountSection()).map_error(ReportError));
      HasDataSection = true;
      break;
    case 0x0D:
      // This section is for ExceptionHandling proposal.
      if (!Conf.hasProposal(Proposal::ExceptionHandling)) {
        return logNeedProposal(ErrCode::Value::MalformedSection,
                               Proposal::ExceptionHandling,
                               FMgr.getLastOffset(), ASTNodeAttr::Module);
      }
      EXPECTED_TRY(loadSection(Mod.getTagSection()).map_error(ReportError));
      break;
    default:
      return logLoadError(ErrCode::Value::MalformedSection,
                          FMgr.getLastOffset(), ASTNodeAttr::Module);
    }

    Offset = FMgr.getOffset();
  }

  setTagFunctionType(Mod.getTagSection(), Mod.getImportSection(),
                     Mod.getTypeSection());

  // Verify the function section and code section are matched.
  if (Mod.getFunctionSection().getContent().size() !=
      Mod.getCodeSection().getContent().size()) {
    spdlog::error(ErrCode::Value::IncompatibleFuncCode);
    spdlog::error(ErrInfo::InfoAST(ASTNodeAttr::Module));
    return Unexpect(ErrCode::Value::IncompatibleFuncCode);
  }

  // Verify the data count section and data segments are matched.
  if (Mod.getDataCountSection().getContent()) {
    if (Mod.getDataSection().getContent().size() !=
        *(Mod.getDataCountSection().getContent())) {
      spdlog::error(ErrCode::Value::IncompatibleDataCount);
      spdlog::error(ErrInfo::InfoAST(ASTNodeAttr::Module));
      return Unexpect(ErrCode::Value::IncompatibleDataCount);
    }
  }

  return {};
}

// Setup symbols from loaded binary. See "include/loader/loader.h".
Expect<void> Loader::loadExecutable(AST::Module &Mod,
                                    std::shared_ptr<Executable> Exec) {
  auto &SubTypes = Mod.getTypeSection().getContent();
  size_t Offset = 0;
  for (const auto &ImpDesc : Mod.getImportSection().getContent()) {
    if (ImpDesc.getExternalType() == ExternalType::Function) {
      ++Offset;
    }
  }
  auto &CodeSegs = Mod.getCodeSection().getContent();

  // Check the symbols.
  auto FuncTypeSymbols = Exec->getTypes(SubTypes.size());
  auto CodeSymbols = Exec->getCodes(Offset, CodeSegs.size());
  auto IntrinsicsSymbol = Exec->getIntrinsics();
  if (unlikely(FuncTypeSymbols.size() != SubTypes.size())) {
    spdlog::error("    AOT section -- number of types not matching:{} {}, "
                  "use interpreter mode instead.",
                  FuncTypeSymbols.size(), SubTypes.size());
    return Unexpect(ErrCode::Value::IllegalGrammar);
  }
  if (unlikely(CodeSymbols.size() != CodeSegs.size())) {
    spdlog::error("    AOT section -- number of codes not matching:{} {}, "
                  "use interpreter mode instead.",
                  CodeSymbols.size(), CodeSegs.size());
    return Unexpect(ErrCode::Value::IllegalGrammar);
  }
  if (unlikely(!IntrinsicsSymbol)) {
    spdlog::error("    AOT section -- intrinsics table symbol not found, use "
                  "interpreter mode instead.");
    return Unexpect(ErrCode::Value::IllegalGrammar);
  }

  // Set the symbols into the module.
  uint32_t FuncTypeIdx = 0;
  for (auto &SubType : SubTypes) {
    if (SubType.getCompositeType().isFunc()) {
      SubType.getCompositeType().getFuncType().setSymbol(
          std::move(FuncTypeSymbols[FuncTypeIdx]));
    }
    FuncTypeIdx++;
  }
  for (size_t I = 0; I < CodeSegs.size(); ++I) {
    CodeSegs[I].setSymbol(std::move(CodeSymbols[I]));
  }
  Mod.setSymbol(std::move(IntrinsicsSymbol));
  if (!Conf.getRuntimeConfigure().isForceInterpreter()) {
    // If the configure is set to force interpreter mode, not to set the
    // symbol.
    if (auto &Symbol = Mod.getSymbol()) {
      *Symbol = IntrinsicsTable;
    }
  }

  return {};
}

Expect<void> Loader::loadUniversalWASM(AST::Module &Mod) {
  if (!Conf.getRuntimeConfigure().isForceInterpreter()) {
    auto Exec = std::make_shared<AOTSection>();
    if (auto Res = Exec->load(Mod.getAOTSection()); unlikely(!Res)) {
      spdlog::error("    AOT section -- library load failed:{} , use "
                    "interpreter mode instead.",
                    Res.error());
    } else {
      if (loadExecutable(Mod, Exec)) {
        return {};
      }
    }
  }

  // Fallback to the interpreter mode case: Re-read the code section.
  WASMType = InputType::WASM;
  FMgr.seek(Mod.getCodeSection().getStartOffset());
  return loadSection(Mod.getCodeSection()).map_error([](auto E) {
    spdlog::error(ErrInfo::InfoAST(ASTNodeAttr::Module));
    return E;
  });
}

Expect<void> Loader::loadModuleAOT(AST::AOTSection &AOTSection) {
  // Find and Read the AOT custom section first. Jump the others.
  // This loop is for checking the input is an universal WASM or not.
  // Therefore, if the configure is set as force interpreter mode, skip this.
  while (WASMType != InputType::SharedLibrary) {
    // This loop only overview the custom sections and read the AOT section.
    // For the other general errors, break and handle in the sequentially
    // parsing below.
    uint8_t NewSectionId = 0x00;
    if (auto Res = FMgr.readByte()) {
      NewSectionId = *Res;
    } else {
      break;
    }

    if (NewSectionId == 0x00U) {
      // Load the section size.
      uint32_t ContentSize = 0;
      if (auto Res = FMgr.readU32()) {
        ContentSize = *Res;
      } else {
        break;
      }
      if (ContentSize > FMgr.getRemainSize()) {
        break;
      }

      // Load the section name.
      auto StartOffset = FMgr.getOffset();
      std::string Name;
      if (auto Res = FMgr.readName()) {
        // The UTF-8 failed case will be ignored here.
        Name = std::move(*Res);
      }

      auto ReadSize = FMgr.getOffset() - StartOffset;
      if (ContentSize < ReadSize) {
        // Syntax error of overread. Jump to the next section.
        FMgr.seek(StartOffset + ContentSize);
        continue;
      }

      if (Name == "wasmedge") {
        // Found the AOT section in universal WASM. Load the AOT code.
        // Read the content.
        std::vector<uint8_t> Content;
        if (auto Res = FMgr.readBytes(ContentSize - ReadSize)) {
          Content = std::move(*Res);
        } else {
          break;
        }

        // Load the AOT section.
        FileMgr VecMgr;
        AST::AOTSection NewAOTSection;
        VecMgr.setCode(Content);
        if (auto Res = loadSection(VecMgr, NewAOTSection)) {
          // Also handle the duplicated AOT sections case.
          // If the new AOT section discovered, use the new one.
          WASMType = InputType::UniversalWASM;
          AOTSection = std::move(NewAOTSection);
        } else {
          // If the new AOT section load failed, use the old one or the
          // interpreter mode.
          if (WASMType == InputType::UniversalWASM) {
            spdlog::info(
                "    Load AOT section failed. Use the previous succeeded one.");
          } else {
            spdlog::info(
                "    Load AOT section failed. Use interpreter mode instead.");
          }
        }
      } else {
        // Found other custom sections. Jump to the next section.
        FMgr.seek(StartOffset + ContentSize);
        continue;
      }
    } else {
      if (auto Res = FMgr.jumpContent(); unlikely(!Res)) {
        break;
      }
    }
  }
  return {};
}

} // namespace Loader
} // namespace WasmEdge
