#include "Diagnostics/Win32CrashTrace.h"

#include <XCEditor/Foundation/UIEditorRuntimeTrace.h>

#include <dbghelp.h>

#include <array>
#include <cstddef>
#include <cstdint>
#include <string>
#include <string_view>
#include <utility>

#pragma comment(lib, "Dbghelp.lib")

namespace XCEngine::UI::Editor::App {

namespace {

bool g_symbolHandlerInitialized = false;

std::string NarrowWideString(std::wstring_view text) {
    if (text.empty()) {
        return {};
    }

    const int sizeNeeded = WideCharToMultiByte(
        CP_UTF8,
        0,
        text.data(),
        static_cast<int>(text.size()),
        nullptr,
        0,
        nullptr,
        nullptr);
    if (sizeNeeded <= 0) {
        return {};
    }

    std::string result(static_cast<std::size_t>(sizeNeeded), '\0');
    WideCharToMultiByte(
        CP_UTF8,
        0,
        text.data(),
        static_cast<int>(text.size()),
        result.data(),
        sizeNeeded,
        nullptr,
        nullptr);
    return result;
}

std::string ResolveModulePathFromAddress(DWORD64 address) {
    if (address == 0u) {
        return {};
    }

    HMODULE moduleHandle = nullptr;
    if (!GetModuleHandleExW(
            GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
                GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
            reinterpret_cast<LPCWSTR>(address),
            &moduleHandle) ||
        moduleHandle == nullptr) {
        return {};
    }

    wchar_t modulePath[MAX_PATH] = {};
    const DWORD length = GetModuleFileNameW(moduleHandle, modulePath, MAX_PATH);
    if (length == 0u) {
        return {};
    }

    return NarrowWideString(std::wstring_view(modulePath, length));
}

void EnsureSymbolHandlerInitialized(HANDLE process) {
    if (g_symbolHandlerInitialized) {
        return;
    }

    SymSetOptions(SYMOPT_DEFERRED_LOADS | SYMOPT_UNDNAME | SYMOPT_LOAD_LINES);
    g_symbolHandlerInitialized = SymInitialize(process, nullptr, TRUE) == TRUE;
}

std::vector<UIEditorCrashTraceFrame> CaptureExceptionStack(
    const EXCEPTION_POINTERS* exceptionPointers) {
    std::vector<UIEditorCrashTraceFrame> frames = {};
    if (exceptionPointers == nullptr || exceptionPointers->ContextRecord == nullptr) {
        return frames;
    }

    HANDLE process = GetCurrentProcess();
    HANDLE thread = GetCurrentThread();
    EnsureSymbolHandlerInitialized(process);

    CONTEXT context = *exceptionPointers->ContextRecord;
    STACKFRAME64 stackFrame = {};
    DWORD machineType = 0u;
#if defined(_M_X64)
    machineType = IMAGE_FILE_MACHINE_AMD64;
    stackFrame.AddrPC.Offset = context.Rip;
    stackFrame.AddrFrame.Offset = context.Rbp;
    stackFrame.AddrStack.Offset = context.Rsp;
#elif defined(_M_IX86)
    machineType = IMAGE_FILE_MACHINE_I386;
    stackFrame.AddrPC.Offset = context.Eip;
    stackFrame.AddrFrame.Offset = context.Ebp;
    stackFrame.AddrStack.Offset = context.Esp;
#else
    return frames;
#endif
    stackFrame.AddrPC.Mode = AddrModeFlat;
    stackFrame.AddrFrame.Mode = AddrModeFlat;
    stackFrame.AddrStack.Mode = AddrModeFlat;

    constexpr std::size_t kMaxFrames = 32u;
    frames.reserve(kMaxFrames);
    for (std::size_t frameIndex = 0u; frameIndex < kMaxFrames; ++frameIndex) {
        const BOOL advanced = StackWalk64(
            machineType,
            process,
            thread,
            &stackFrame,
            &context,
            nullptr,
            SymFunctionTableAccess64,
            SymGetModuleBase64,
            nullptr);
        if (advanced == FALSE || stackFrame.AddrPC.Offset == 0u) {
            break;
        }

        UIEditorCrashTraceFrame frame = {};
        frame.address = static_cast<std::uintptr_t>(stackFrame.AddrPC.Offset);
        frame.moduleName = ResolveModulePathFromAddress(stackFrame.AddrPC.Offset);

        std::array<std::byte, sizeof(SYMBOL_INFO) + MAX_SYM_NAME> symbolBuffer = {};
        auto* symbol = reinterpret_cast<SYMBOL_INFO*>(symbolBuffer.data());
        symbol->SizeOfStruct = sizeof(SYMBOL_INFO);
        symbol->MaxNameLen = MAX_SYM_NAME;
        DWORD64 displacement = 0u;
        if (g_symbolHandlerInitialized &&
            SymFromAddr(process, stackFrame.AddrPC.Offset, &displacement, symbol) == TRUE) {
            frame.displacement = static_cast<std::uintptr_t>(displacement);
            frame.symbolName = symbol->Name;
        }

        IMAGEHLP_LINE64 line = {};
        line.SizeOfStruct = sizeof(line);
        DWORD lineDisplacement = 0u;
        if (g_symbolHandlerInitialized &&
            SymGetLineFromAddr64(
                process,
                stackFrame.AddrPC.Offset,
                &lineDisplacement,
                &line) == TRUE) {
            frame.line = line.LineNumber;
            frame.lineDisplacement = lineDisplacement;
            if (line.FileName != nullptr) {
                frame.sourceFile = line.FileName;
            }
        }

        frames.push_back(std::move(frame));
    }

    return frames;
}

} // namespace

void AppendWin32UnhandledExceptionCrashTrace(EXCEPTION_POINTERS* exceptionInfo) {
    UIEditorCrashTrace trace = {};
    if (exceptionInfo != nullptr && exceptionInfo->ExceptionRecord != nullptr) {
        trace.exceptionCode = exceptionInfo->ExceptionRecord->ExceptionCode;
        trace.exceptionAddress = exceptionInfo->ExceptionRecord->ExceptionAddress;
        trace.exceptionModule = ResolveModulePathFromAddress(
            reinterpret_cast<DWORD64>(trace.exceptionAddress));
    }
    trace.threadId = GetCurrentThreadId();
    trace.stackFrames = CaptureExceptionStack(exceptionInfo);
    AppendUIEditorCrashTrace(trace);
}

} // namespace XCEngine::UI::Editor::App