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

#include "system/stacktrace.h"
#include "common/spdlog.h"
#include <fmt/ranges.h>

#if WASMEDGE_OS_WINDOWS
#include "system/winapi.h"
#elif WASMEDGE_OS_LINUX
#include <unwind.h>
#elif WASMEDGE_OS_MACOS
#include <execinfo.h>
#endif

namespace WasmEdge {

using namespace std::literals;

Span<void *const> stackTrace(Span<void *> Buffer) noexcept {
#if WASMEDGE_OS_WINDOWS
  struct DbgHelp {
    DbgHelp() noexcept : Process(winapi::GetCurrentProcess()) {
      winapi::SymSetOptions(winapi::SYMOPT_DEFERRED_LOADS_);
      winapi::SymInitializeW(Process, nullptr, true);
      SelfBase = winapi::SymGetModuleBase64(
          Process, reinterpret_cast<winapi::DWORD64_>(&stackTrace));
      NtDllBase = winapi::SymGetModuleBase64(
          Process, reinterpret_cast<winapi::DWORD64_>(
                       &winapi::RtlCaptureStackBackTrace));
      Kernel32Base = winapi::SymGetModuleBase64(
          Process, reinterpret_cast<winapi::DWORD64_>(&winapi::CloseHandle));
    }
    ~DbgHelp() noexcept { winapi::SymCleanup(Process); }
    void refresh() noexcept { winapi::SymRefreshModuleList(Process); }
    winapi::HANDLE_ Process;
    winapi::DWORD64_ SelfBase, NtDllBase, Kernel32Base;
  };
  static DbgHelp Helper;
  Helper.refresh();
  auto Depth = static_cast<size_t>(winapi::RtlCaptureStackBackTrace(
      1u, static_cast<winapi::ULONG_>(Buffer.size()), Buffer.data(), nullptr));
  size_t NewDepth = 0;
  for (size_t I = 0; I < Depth; ++I) {
    auto Base = winapi::SymGetModuleBase64(
        Helper.Process, reinterpret_cast<winapi::DWORD64_>(Buffer[I]));
    if (Base == 0 || (Base != Helper.SelfBase && Base != Helper.NtDllBase &&
                      Base != Helper.Kernel32Base)) {
      Buffer[NewDepth++] = Buffer[I];
    }
  }
  return Buffer.first(static_cast<size_t>(NewDepth));
#elif WASMEDGE_OS_LINUX
  struct BacktraceState {
    Span<void *> Buffer;
    size_t Index;
  };
  BacktraceState State{Buffer, 0};
  _Unwind_Backtrace(
      [](struct _Unwind_Context *Ctx, void *Arg) noexcept {
        auto &State = *static_cast<BacktraceState *>(Arg);
        if (State.Index >= State.Buffer.size()) {
          return _URC_END_OF_STACK;
        }
        State.Buffer[State.Index++] =
            reinterpret_cast<void *>(_Unwind_GetIP(Ctx));
        return _URC_NO_REASON;
      },
      &State);
  return Buffer.first(State.Index);
#elif WASMEDGE_OS_MACOS
  const auto Depth = backtrace(Buffer.data(), Buffer.size());
  return Buffer.first(Depth);
#endif
}

Span<const uint32_t>
interpreterStackTrace(const Runtime::StackManager &StackMgr,
                      Span<uint32_t> Buffer) noexcept {
  size_t Index = 0;
  if (auto Module = StackMgr.getModule()) {
    const auto FuncInsts = Module->getFunctionInstances();
    std::map<AST::InstrView::iterator, int64_t> Funcs;
    for (size_t I = 0; I < FuncInsts.size(); ++I) {
      const auto &Func = FuncInsts[I];
      if (Func && Func->isWasmFunction()) {
        const auto &Instrs = Func->getInstrs();
        Funcs.emplace(Instrs.end(), INT64_C(-1));
        Funcs.emplace(Instrs.begin(), I);
      }
    }
    for (const auto &Frame : StackMgr.getFramesSpan()) {
      auto Entry = Frame.From;
      auto Iter = Funcs.lower_bound(Entry);
      if ((Iter == Funcs.end() || Iter->first > Entry) &&
          Iter != Funcs.begin()) {
        --Iter;
      }
      if (Iter != Funcs.end() && Iter->first < Entry &&
          Iter->second >= INT64_C(0) && Index < Buffer.size()) {
        Buffer[Index++] = static_cast<uint32_t>(Iter->second);
      }
    }
  }
  return Buffer.first(Index);
}

Span<const uint32_t> compiledStackTrace(const Runtime::StackManager &StackMgr,
                                        Span<uint32_t> Buffer) noexcept {
  std::array<void *, 256> StackTraceBuffer;
  return compiledStackTrace(StackMgr, stackTrace(StackTraceBuffer), Buffer);
}

Span<const uint32_t> compiledStackTrace(const Runtime::StackManager &StackMgr,
                                        Span<void *const> Stack,
                                        Span<uint32_t> Buffer) noexcept {
  std::map<void *, int64_t> Funcs;
  size_t Index = 0;
  if (auto Module = StackMgr.getModule()) {
    const auto FuncInsts = Module->getFunctionInstances();
    for (size_t I = 0; I < FuncInsts.size(); ++I) {
      const auto &Func = FuncInsts[I];
      if (Func && Func->isCompiledFunction()) {
        Funcs.emplace(
            reinterpret_cast<void *>(Func->getFuncType().getSymbol().get()),
            INT64_C(-1));
        Funcs.emplace(Func->getSymbol().get(), I);
      }
    }
    for (auto Entry : Stack) {
      auto Iter = Funcs.lower_bound(Entry);
      if ((Iter == Funcs.end() || Iter->first > Entry) &&
          Iter != Funcs.begin()) {
        --Iter;
      }
      if (Iter != Funcs.end() && Iter->first < Entry &&
          Iter->second >= INT64_C(0) && Index < Buffer.size()) {
        Buffer[Index++] = static_cast<uint32_t>(Iter->second);
      }
    }
  }
  return Buffer.first(Index);
}

void dumpStackTrace(Span<const uint32_t> Stack) noexcept {
  spdlog::error("calling stack:{}"sv, fmt::join(Stack, ", "sv));
}

} // namespace WasmEdge
