/***
    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 <Source/Process/AuProcessMap.hpp>
#include <Source/Process/AuProcessMap.Win32.hpp>

namespace Aurora::Debug
{
    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)
    {
        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;

            auto result = StackWalk64
            (
            #if defined(AURORA_ARCH_X64)
                IMAGE_FILE_MACHINE_AMD64,
            #else
                IMAGE_FILE_MACHINE_I386,
            #endif
                INVALID_HANDLE_VALUE,
                INVALID_HANDLE_VALUE,
                &stack,
                &cpy,
                NULL,
                SymFunctionTableAccess64,
                SymGetModuleBase64,
                NULL
            );

            if (!result)
            {
                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 (SymGetLineFromAddr64(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 {};
        ctx.ContextFlags = CONTEXT_ALL;
        if (!GetThreadContext(GetCurrentThread(), &ctx))
        {
            return {};
        }

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