/***
    Copyright (C) 2022 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: Stack.Win32.cpp
    Date: 2022-1-26
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "Stack.hpp"
#include "Stack.Win32.hpp"
#include <Dbghelp.h>
#include <Aux_ulib.h>
#include <Source/Process/AuProcessMap.hpp>
#include <Source/Process/AuProcessMap.Win32.hpp>

namespace Aurora::Debug
{
    // Do not use a futex here!
    static AuMutex gMutex;

    static AuUInt GetImageBase(HMODULE mod)
    {
        if (!mod)
        {
            return 0;
        }

        if (mod == INVALID_HANDLE_VALUE)
        {
            return 0;
        }

        auto cache = Process::GetFromModuleCache(reinterpret_cast<AuUInt>(mod));
        if (!cache.moduleMeta)
        {
            return 0;
        }

        return cache.moduleMeta->origVa;
    }

    void ParseStack(CONTEXT *ctx, StackTrace &backTrace)
    {
        BOOL bLocked {};
    #if defined(AURORA_PLATFORM_WIN32)
        if (AuxUlibIsDLLSynchronizationHeld(&bLocked))
        {
            if (bLocked)
            {
                return;
            }
        }
    #endif


        AU_LOCK_GUARD(gMutex);

        char buffer[sizeof(SYMBOL_INFO) + (MAX_SYM_NAME + 1) * sizeof(char)] = { 0 };
        AuString     backTraceBuffer;
        HMODULE      hModule;
        DWORD64      displacement = 0;
        HANDLE       process = GetCurrentProcess();
        STACKFRAME64 stack = { 0 };
        CONTEXT      cpy = *ctx;
        PSYMBOL_INFO pSymbol = (PSYMBOL_INFO)buffer;

    #if defined(AURORA_ARCH_X64)
        stack.AddrPC.Offset    = ctx->Rip;
        stack.AddrPC.Mode      = AddrModeFlat;
        stack.AddrStack.Offset = ctx->Rsp;
        stack.AddrStack.Mode   = AddrModeFlat;
        stack.AddrFrame.Offset = ctx->Rbp;
        stack.AddrFrame.Mode   = AddrModeFlat;
    #elif defined(AURORA_ARCH_X86)
        stack.AddrPC.Offset    = ctx->Eip;
        stack.AddrPC.Mode      = AddrModeFlat;
        stack.AddrStack.Offset = ctx->Esp;
        stack.AddrStack.Mode   = AddrModeFlat;
        stack.AddrFrame.Offset = ctx->Ebp;
        stack.AddrFrame.Mode   = AddrModeFlat;
    #endif

        for (ULONG frame = 0; ; frame++)
        {
            StackTraceEntry frameCurrent;
            bool bResult {};

            if (pSymGetModuleBase64 && pStackWalk64)
            {
                bResult = pStackWalk64
                (
                #if defined(AURORA_ARCH_X64)
                    IMAGE_FILE_MACHINE_AMD64,
                #else
                    IMAGE_FILE_MACHINE_I386,
                #endif
                    GetCurrentProcess(),
                    INVALID_HANDLE_VALUE,
                    &stack,
                    &cpy,
                    NULL,
                    pSymFunctionTableAccess64,
                    pSymGetModuleBase64,
                    NULL
                );
            }

            if (!bResult)
            {
                break;
            }

            pSymbol->SizeOfStruct = sizeof(SYMBOL_INFO);
            pSymbol->MaxNameLen = MAX_SYM_NAME;

            frameCurrent.address = stack.AddrPC.Offset;
            frameCurrent.relAddress = frameCurrent.address;

            if (GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, (LPCTSTR)(stack.AddrPC.Offset), &hModule))
            {
                if (hModule != NULL)
                {
                    frameCurrent.module = Process::ModuleToPath(hModule);
                    frameCurrent.relAddress = frameCurrent.relAddress - reinterpret_cast<AuUInt>(hModule) + GetImageBase(hModule);
                }
            }

        #if defined(DEBUG) || defined(STAGING)
            IMAGEHLP_LINE64 line;
            DWORD disp;
            
            line.SizeOfStruct = sizeof(IMAGEHLP_LINE64);

            if (pSymGetLineFromAddr64)
            {
                if (pSymGetLineFromAddr64(process, stack.AddrPC.Offset, &disp, &line))
                {
                    frameCurrent.file = AuMakeTuple(line.FileName, line.LineNumber, 0);
                }
            }
        #endif 

            backTrace.push_back(frameCurrent);
        }
    }

    StackTrace PlatformWalkCallStack()
    {
        StackTrace ret;
        CONTEXT ctx {};

    #if defined(AURORA_ARCH_X86)

        ctx.ContextFlags = CONTEXT_CONTROL;

        __asm
        {
        Label:
            mov [ctx.Ebp], ebp;
            mov [ctx.Esp], esp;
            mov eax, [Label];
            mov [ctx.Eip], eax;
        }

    #elif defined(AURORA_ARCH_X64) && defined(AURORA_PLATFORM_WIN32)

        RtlCaptureContext(&ctx);

    #else

        ctx.ContextFlags = CONTEXT_CONTROL;
        if (!GetThreadContext(GetCurrentThread(), &ctx))
        {
            return {};
        }

    #endif

        ParseStack(&ctx, ret);
        return ret;
    }
}