#include "./stacktrace.hpp"

#include <optional>
#include <string>
#include "../detect.hpp"
#include "../util/string.hpp"

#if MIJIN_TARGET_OS != MIJIN_OS_WINDOWS && (MIJIN_COMPILER == MIJIN_COMPILER_CLANG || MIJIN_COMPILER == MIJIN_COMPILER_GCC)
#define MIJIN_USE_LIBBACKTRACE 1
#else
#define MIJIN_USE_LIBBACKTRACE 0
#endif

#if MIJIN_USE_LIBBACKTRACE
    #include <backtrace.h>
#elif MIJIN_TARGET_OS == MIJIN_OS_WINDOWS
    #include <array>
    #include <cstddef>
    #include <mutex>
    #include <Windows.h>
    #include <DbgHelp.h>
    #include "../util/winundef.hpp"
    #pragma comment(lib, "dbghelp")
#endif


namespace mijin
{
namespace
{
//
// internal defines
//

//
// internal constants
//

//
// internal types
//

#if MIJIN_USE_LIBBACKTRACE
struct BacktraceData
{
    std::optional<std::string> error;
    std::vector<Stackframe> stackframes;
};
#elif MIJIN_TARGET_OS == MIJIN_OS_WINDOWS
HANDLE gProcessHandle = nullptr;
#endif

//
// internal variables
//

thread_local Optional<Stacktrace> gCurrentExceptionStackTrace;

#if MIJIN_TARGET_OS == MIJIN_OS_WINDOWS
std::mutex gDbgHelpMutex;
bool gDbgHelpInitCalled = false;
#endif

//
// internal functions
//

#if MIJIN_USE_LIBBACKTRACE
int backtraceFullCallback(void* data, std::uintptr_t programCounter, const char* filename, int lineno, const char* function)
{
    BacktraceData& btData = *static_cast<BacktraceData*>(data);
    btData.stackframes.push_back({
        .address = reinterpret_cast<void*>(programCounter), // NOLINT(performance-no-int-to-ptr)
        .filename = filename ? filename : "<unknown>",
        .function = function ? function : "<unknown>",
        .lineNumber = lineno
    });
    return 0;
}

void backtraceErrorCallback(void* data, const char* msg, int /* errnum */)
{
    BacktraceData& btData = *static_cast<BacktraceData*>(data);
    btData.error = msg;
}

thread_local backtrace_state* gBacktraceState = nullptr;
#elif MIJIN_TARGET_OS == MIJIN_OS_WINDOWS
void cleanupDbgHelp() MIJIN_NOEXCEPT
{
    if (!SymCleanup(gProcessHandle))
    {
        [[maybe_unused]] const DWORD error = GetLastError();
        MIJIN_ERROR("Error cleaning up DbgHelp.");
    }
}

[[nodiscard]]
bool initDbgHelp() MIJIN_NOEXCEPT
{
    if (gDbgHelpInitCalled)
    {
        return gProcessHandle != nullptr; // if init was successful, process handle is not null
    }
    gDbgHelpInitCalled = true;

    SymSetOptions(SYMOPT_UNDNAME | SYMOPT_DEFERRED_LOADS);

    HANDLE hCurrentProcess = GetCurrentProcess();
    HANDLE hCopy = nullptr;
    if (!DuplicateHandle(hCurrentProcess, hCurrentProcess, hCurrentProcess, &hCopy, 0, FALSE, DUPLICATE_SAME_ACCESS))
    {
        [[maybe_unused]] const DWORD error = GetLastError();
        MIJIN_ERROR("Error duplicating process handle.");
        return false;
    }
    if (!SymInitialize(hCopy, nullptr, true))
    {
        [[maybe_unused]] const DWORD error = GetLastError();
        MIJIN_ERROR("Error initializing DbHelp.");
        return false;
    }

    [[maybe_unused]] const int result = std::atexit(&cleanupDbgHelp);
    MIJIN_ASSERT(result == 0, "Error registering DbgHelp cleanup handler.");

    // only copy in the end so we can still figure out if initialization was successful
    gProcessHandle = hCopy;
    return true;
}
#endif // MIJIN_USE_LIBBACKTRACE
} // namespace

//
// public functions
//

Result<Stacktrace> captureStacktrace(unsigned skipFrames) MIJIN_NOEXCEPT
{
#if MIJIN_USE_LIBBACKTRACE
    BacktraceData btData;
    if (gBacktraceState == nullptr)
    {
        gBacktraceState = backtrace_create_state(/*filename = */ nullptr, /* threaded = */ false, &backtraceErrorCallback, &btData);
    }
    if (btData.error.has_value())
    {
        return ResultError(std::move(*btData.error));
    }
    if (gBacktraceState == nullptr)
    {
        return ResultError("Error initializing libbacktrace.");
    }
    backtrace_full(gBacktraceState, static_cast<int>(skipFrames) + 1, &backtraceFullCallback, &backtraceErrorCallback, &btData);
    if (btData.error.has_value())
    {
        return ResultError(std::move(*btData.error));
    }

    return Stacktrace(std::move(btData.stackframes));
#elif MIJIN_TARGET_OS == MIJIN_OS_WINDOWS
    if (!initDbgHelp())
    {
        return ResultError("error initializing DbgHelp");
    }

    const HANDLE hThread = GetCurrentThread();

    CONTEXT context;
    RtlCaptureContext(&context);

    STACKFRAME64 stackFrame = {
        .AddrPC = {
            .Offset = context.Rip,
            .Mode = AddrModeFlat
        },
        .AddrFrame = {
            .Offset = context.Rbp,
            .Mode = AddrModeFlat
        },
        .AddrStack = {
            .Offset = context.Rsp,
            .Mode = AddrModeFlat
        }
    };

    ++skipFrames; // always skip the first frame (the current function)

    // for symbol info
    DWORD64 displacement64 = 0;
    static constexpr std::size_t SYMBOL_BUFFER_SIZE = sizeof(SYMBOL_INFO) + (MAX_SYM_NAME * sizeof(char));
    std::array<std::byte, SYMBOL_BUFFER_SIZE> symbolBuffer alignas(SYMBOL_INFO);
    SYMBOL_INFO& symbolInfo = *std::bit_cast<SYMBOL_INFO*>(symbolBuffer.data());
    symbolInfo.SizeOfStruct = sizeof(SYMBOL_BUFFER_SIZE);
    symbolInfo.MaxNameLen = MAX_SYM_NAME;

    // for file and line info
    DWORD displacement = 0;
    IMAGEHLP_LINE64 line64;
    line64.SizeOfStruct = sizeof(IMAGEHLP_LINE64);

    std::vector<Stackframe> stackframes;
    while (StackWalk64(
            /* MachineType                = */ IMAGE_FILE_MACHINE_AMD64,
            /* hProcess                   = */ gProcessHandle,
            /* hThread                    = */ hThread,
            /* StackFrame                 = */ &stackFrame,
            /* ContextRecord              = */ &context,
            /* ReadMemoryRoutine          = */ nullptr,
            /* FunctionTableAccessRoutine = */ &SymFunctionTableAccess64,
            /* GetModuleBaseRoutine       = */ &SymGetModuleBase64,
            /* TranslateAddress           = */ nullptr
        ))
    {
        if (skipFrames > 0)
        {
            --skipFrames;
            continue;
        }
        Stackframe& frame = stackframes.emplace_back();
        const DWORD64 baseAddress = SymGetModuleBase64(gProcessHandle, stackFrame.AddrPC.Offset);
        const DWORD64 relativeAddress = stackFrame.AddrPC.Offset - baseAddress;

        frame.address = std::bit_cast<void*>(relativeAddress);

        if (SymFromAddr(gProcessHandle, stackFrame.AddrPC.Offset, &displacement64, &symbolInfo))
        {
            frame.function = symbolInfo.Name;
        }

        if (SymGetLineFromAddr64(gProcessHandle, stackFrame.AddrPC.Offset, &displacement, &line64))
        {
            frame.filename = line64.FileName;
            frame.lineNumber = static_cast<int>(line64.LineNumber);
        }
    }

    return Stacktrace(std::move(stackframes));
#else // MIJIN_USE_LIBBACKTRACE || (MIJIN_TARGET_OS == MIJIN_OS_WINDOWS)
	(void) skipFrames;
	return ResultError("not implemented");
#endif // MIJIN_USE_LIBBACKTRACE
}

const Optional<Stacktrace>& getExceptionStacktrace() MIJIN_NOEXCEPT
{
    return gCurrentExceptionStackTrace;
}
} // namespace mijin

#if MIJIN_STDLIB == MIJIN_STDLIB_GLIBC
#include <dlfcn.h>

namespace
{
using cxa_throw_type = void (*)(void*, std::type_info*, void(*)(void*));
using cxa_rethrow_type = void (*)();

cxa_throw_type orig_cxa_throw = nullptr; // Address of the original __cxa_throw
// cxa_rethrow_type orig_cxa_rethrow = nullptr; // Address of the original __cxa_rethrow

extern "C" void __cxa_throw(void* thrown_exception, std::type_info* tinfo, void (*dest)(void*)) // NOLINT(bugprone-reserved-identifier,cert-dcl37-c,cert-dcl51-cpp)
{
    if (!orig_cxa_throw) {
        orig_cxa_throw = reinterpret_cast<cxa_throw_type>(dlsym(RTLD_NEXT, "__cxa_throw"));
    }
    if (mijin::Result<mijin::Stacktrace> stacktrace = mijin::captureStacktrace(); stacktrace.isSuccess())
    {
        mijin::gCurrentExceptionStackTrace = std::move(*stacktrace);
    }
    else
    {
        mijin::gCurrentExceptionStackTrace.reset();
    }
    if (orig_cxa_throw) {
        orig_cxa_throw(thrown_exception, tinfo, dest);
    }
    else {
        std::abort();
    }
}
}
#endif