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

    File: ExceptionWatcher.Win32.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "Debug.hpp"
#include "ExceptionWatcher.Win32.hpp"

#include <Source/Process/ProcessMap.Win32.hpp>
#include <Source/Telemetry/Telemetry.hpp>

#include <Windows.h>
#include <Dbghelp.h>
#include <codecvt>

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

#include <vcruntime_exception.h>
#include <ehdata.h>

static thread_local int     gDebugLocked = 0;

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

            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);
                }
            }

            if (SymFromAddr(process, (ULONG64)stack.AddrPC.Offset, &displacement, pSymbol))
            {
                frameCurrent.label = pSymbol->Name;
            }

        #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 = std::make_tuple(line.FileName, line.LineNumber, 0);
            }
        #endif 

            backTrace.push_back(frameCurrent);
        }
    }

    #define EXCEPTION_ENTRY(n) {n, #n}
    static const AuHashMap<DWORD, AuString> kExceptionTable
    {
        EXCEPTION_ENTRY(STILL_ACTIVE),
        EXCEPTION_ENTRY(EXCEPTION_ACCESS_VIOLATION),
        EXCEPTION_ENTRY(EXCEPTION_DATATYPE_MISALIGNMENT),
        EXCEPTION_ENTRY(EXCEPTION_BREAKPOINT),
        EXCEPTION_ENTRY(EXCEPTION_SINGLE_STEP),
        EXCEPTION_ENTRY(EXCEPTION_ARRAY_BOUNDS_EXCEEDED),
        EXCEPTION_ENTRY(EXCEPTION_FLT_DENORMAL_OPERAND),
        EXCEPTION_ENTRY(EXCEPTION_FLT_DIVIDE_BY_ZERO),
        EXCEPTION_ENTRY(EXCEPTION_FLT_INEXACT_RESULT),
        EXCEPTION_ENTRY(EXCEPTION_FLT_INVALID_OPERATION),
        EXCEPTION_ENTRY(EXCEPTION_FLT_OVERFLOW),
        EXCEPTION_ENTRY(EXCEPTION_FLT_STACK_CHECK),
        EXCEPTION_ENTRY(EXCEPTION_FLT_UNDERFLOW),
        EXCEPTION_ENTRY(EXCEPTION_INT_DIVIDE_BY_ZERO),
        EXCEPTION_ENTRY(EXCEPTION_INT_OVERFLOW),
        EXCEPTION_ENTRY(EXCEPTION_PRIV_INSTRUCTION),
        EXCEPTION_ENTRY(EXCEPTION_IN_PAGE_ERROR),
        EXCEPTION_ENTRY(EXCEPTION_ILLEGAL_INSTRUCTION),
        EXCEPTION_ENTRY(EXCEPTION_NONCONTINUABLE_EXCEPTION),
        EXCEPTION_ENTRY(EXCEPTION_STACK_OVERFLOW),
        EXCEPTION_ENTRY(EXCEPTION_INVALID_DISPOSITION),
        EXCEPTION_ENTRY(EXCEPTION_GUARD_PAGE),
        EXCEPTION_ENTRY(EXCEPTION_INVALID_HANDLE),
        //EXCEPTION_ENTRY(EXCEPTION_POSSIBLE_DEADLOCK),
        EXCEPTION_ENTRY(CONTROL_C_EXIT)
    };
#undef EXCEPTION_ENTRY
#define EXCEPTION_ENTRY(n) {n, true}
    static const AuHashMap<DWORD, bool> kExceptionFatalTable
    {
        EXCEPTION_ENTRY(EXCEPTION_ACCESS_VIOLATION),
        EXCEPTION_ENTRY(EXCEPTION_DATATYPE_MISALIGNMENT),
        EXCEPTION_ENTRY(EXCEPTION_ARRAY_BOUNDS_EXCEEDED),
        EXCEPTION_ENTRY(EXCEPTION_FLT_DENORMAL_OPERAND),
        EXCEPTION_ENTRY(EXCEPTION_FLT_DIVIDE_BY_ZERO),
        EXCEPTION_ENTRY(EXCEPTION_FLT_INEXACT_RESULT),
        EXCEPTION_ENTRY(EXCEPTION_FLT_INVALID_OPERATION),
        EXCEPTION_ENTRY(EXCEPTION_FLT_OVERFLOW),
        EXCEPTION_ENTRY(EXCEPTION_FLT_STACK_CHECK),
        EXCEPTION_ENTRY(EXCEPTION_FLT_UNDERFLOW),
        EXCEPTION_ENTRY(EXCEPTION_INT_DIVIDE_BY_ZERO),
        EXCEPTION_ENTRY(EXCEPTION_INT_OVERFLOW),
        EXCEPTION_ENTRY(EXCEPTION_PRIV_INSTRUCTION),
        EXCEPTION_ENTRY(EXCEPTION_IN_PAGE_ERROR),
        EXCEPTION_ENTRY(EXCEPTION_ILLEGAL_INSTRUCTION),
        EXCEPTION_ENTRY(EXCEPTION_NONCONTINUABLE_EXCEPTION),
        EXCEPTION_ENTRY(EXCEPTION_STACK_OVERFLOW),
        EXCEPTION_ENTRY(EXCEPTION_INVALID_DISPOSITION),
        EXCEPTION_ENTRY(EXCEPTION_GUARD_PAGE)
    };
#undef EXCEPTION_ENTRY
   
    static bool IsReadable(const void *address)
    {
        MEMORY_BASIC_INFORMATION info;

        if (!VirtualQuery(address, &info, sizeof(info)))
        {
            return false;
        }

        if (!info.BaseAddress)
        {
            return false;
        }

        return (info.Protect & (PAGE_READONLY | PAGE_READWRITE)) != 0;
    }

    void InitWin32()
    {
    #if defined(DEBUG) || defined(STAGING)
        SymInitialize(GetCurrentProcess(), NULL, TRUE);
    #endif 

        AddVectoredExceptionHandler(1,
        [](_EXCEPTION_POINTERS *ExceptionInfo) -> LONG
        {
            Telemetry::NewBlockboxEntry entry;
            entry.type = Telemetry::ENewBlackBoxEntry::eWinCxxException;

            if (ExceptionInfo->ExceptionRecord->ExceptionCode < STATUS_GUARD_PAGE_VIOLATION)
            {
                return EXCEPTION_CONTINUE_SEARCH;
            }

            if (ExceptionInfo->ExceptionRecord->ExceptionCode == EXCEPTION_BREAKPOINT)
            {
                return EXCEPTION_CONTINUE_SEARCH;
            }

            auto minimal = gDebugLocked++;
            
            if (minimal >= 5)
            {
                SysPanic("Nested Exception");
            }

            bool cxxThrow     = ExceptionInfo->ExceptionRecord->ExceptionInformation[0] == EH_MAGIC_NUMBER1;
            bool cxxThrowPure = ExceptionInfo->ExceptionRecord->ExceptionInformation[0] == EH_PURE_MAGIC_NUMBER1;

            if ((ExceptionInfo->ExceptionRecord->ExceptionCode == EH_EXCEPTION_NUMBER) &&
                (ExceptionInfo->ExceptionRecord->NumberParameters >= 4) &&
                 (cxxThrow) ||
                 (cxxThrowPure)
                )
            {
                HMODULE handle {};
                auto *throwInfo = reinterpret_cast<ThrowInfo *>(ExceptionInfo->ExceptionRecord->ExceptionInformation[2]);

                if (throwInfo)
                {
                    auto attribs = throwInfo->attributes;

                    if (_EH_RELATIVE_TYPEINFO)
                    {
                        handle = reinterpret_cast<HMODULE>(ExceptionInfo->ExceptionRecord->ExceptionInformation[3]);
                    }
                    
                    const auto catchableTypeArray = reinterpret_cast<const CatchableTypeArray*>(reinterpret_cast<AuUInt>(handle) + static_cast<AuUInt>(throwInfo->pCatchableTypeArray)); 
                    const auto type               = reinterpret_cast<CatchableType*>           (reinterpret_cast<AuUInt>(handle) + static_cast<AuUInt>(catchableTypeArray->arrayOfCatchableTypes[0]));
                    const auto descriptor         = reinterpret_cast<TypeDescriptor*>          (reinterpret_cast<AuUInt>(handle) + static_cast<AuUInt>(type->pType));
               
                    // .?AVException and .?AVexception follow the same abi of const char * at pos 0 (__std_exception_data) after a vtable ptr of i dont care to define how many functions
                    // odds are, these assumptions about the ABI will never break. the QOL improvement of knowing of who is farting in my address space far exceeds an issue under win 720 targets 10 years in the future
                    if (strnicmp(descriptor->name, ".?AVException@", AuArraySize(".?AVException@") - 1) == 0)
                    {
                        auto exception = reinterpret_cast<std::exception *>(ExceptionInfo->ExceptionRecord->ExceptionInformation[1]);
                        if (IsReadable(exception))
                        {
                           entry.wincxx.str = exception->what();
                        }
                    }
                    else if (((type->properties & CT_IsSimpleType) != 0) && 
                             ((type->properties & CT_IsWinRTHandle) == 0))
                    {
                        // assert descriptor->name == ".PEAD"? `DEAP.`?
                        auto possibleStringPointer = reinterpret_cast<const char **>(ExceptionInfo->ExceptionRecord->ExceptionInformation[1]);
                        if (IsReadable(possibleStringPointer))
                        {
                            auto string = *possibleStringPointer;
                            // Not a string? I don't care, it has a null byte under 1k in a readable page
                            if (IsReadable(string) && (strnlen(string, 4096) < 1024))
                            {
                                entry.wincxx.str = string;
                            }
                        }
                    }
                    else if (strnicmp(descriptor->name, ".?AV", AuArraySize(".?AV") - 1) == 0)
                    {
                        /* Annoying
                                  https://blog.quarkslab.com/visual-c-rtti-inspection.html
                                  This structure is very important to identify an object since it contains its VFT (field pVFTable) and its mangled name. That's why it usually starts with ".?AV", which means "a C++ class". These structures are stored in the section ".data".
                                  We decided to do pattern matching on ".?AV" to get the field name of _TypeInformation and thus retrieves the RTTICompleteObjectLocator.
                           ... we would then have to traverse the hierarchy to determine the root most signature (std::exceptions vtable hash) 
                           TODO(Reece): fix me, this is evil and shouldn't make it into the wild
                           Fix before 1.0
                        */

                        auto exception = reinterpret_cast<std::exception *>(ExceptionInfo->ExceptionRecord->ExceptionInformation[1]);
                        if (IsReadable(exception))
                        {
                           entry.wincxx.str = exception->what();
                        }
                    }
                    else if (strlen(descriptor->name) == 0)
                    {
                        auto exception = reinterpret_cast<std::exception *>(ExceptionInfo->ExceptionRecord->ExceptionInformation[1]);
                        if (IsReadable(exception))
                        {
                           entry.wincxx.str = exception->what();
                        }
                    }
                }
            }

            bool isCritical = AuTryFind(kExceptionFatalTable, ExceptionInfo->ExceptionRecord->ExceptionCode);
            if (entry.wincxx.str.empty())
            {
                const AuString *msg;
                if (AuTryFind(kExceptionTable, ExceptionInfo->ExceptionRecord->ExceptionCode, msg))
                {
                    entry.wincxx.str = *msg;
                }
                else
                {
                    entry.wincxx.str = std::to_string(ExceptionInfo->ExceptionRecord->ExceptionCode);
                }
            }
            else
            {
                if (entry.wincxx.str.find("invalid sto") != std::string::npos)
                {
                    gDebugLocked = 0;
                    return EXCEPTION_CONTINUE_SEARCH;
                }
            }

            try
            {
                if (minimal < 2)
                {
                    ParseStack(ExceptionInfo->ContextRecord, entry.wincxx.stack.backtrace);
                }
            }
            catch (...)
            {

            }

        #if defined(STAGING) || defined(DEBUG)
            bool isInternal = true;
        #else
            bool isInternal = false;
        #endif

            if ((isCritical || isInternal) && (minimal == 0))
            {
                LogWarn("NT Exception: 0x{:x}, {}", ExceptionInfo->ExceptionRecord->ExceptionCode, entry.wincxx.str);
                LogWarn("{}", StringifyStackTrace(entry.wincxx.stack.backtrace));
            }

            try
            {
                Telemetry::Report(entry);
                
                if (isCritical)
                {
                    Telemetry::Mayday();
                }
            }
            catch (...)
            {

            }

            gDebugLocked = 0;
            return EXCEPTION_CONTINUE_SEARCH;
        });
    }
}