/***
    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()
    {
        static AuString kStringRawName = typeid(AuString).raw_name();


    #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]);
                    }
                    
                    auto exception = reinterpret_cast<std::exception *>(ExceptionInfo->ExceptionRecord->ExceptionInformation[1]);
                    const auto catchableTypeArray = reinterpret_cast<const CatchableTypeArray*>(reinterpret_cast<AuUInt>(handle) + static_cast<AuUInt>(throwInfo->pCatchableTypeArray));
                    
                    AuString suffix;
                    bool derivedFromException = {};

                    for (int i = 0; i < catchableTypeArray->nCatchableTypes; i++)
                    {
                        const auto type = reinterpret_cast<CatchableType *>           (reinterpret_cast<AuUInt>(handle) + static_cast<AuUInt>(catchableTypeArray->arrayOfCatchableTypes[i]));
                        const auto descriptor = reinterpret_cast<std::type_info *>    (reinterpret_cast<AuUInt>(handle) + static_cast<AuUInt>(type->pType));

                        entry.wincxx.str += (i == 0 ? "" : AuString(", ")) + descriptor->name();

                        if (strnicmp(descriptor->raw_name(), ".?AVException@", AuArraySize(".?AVException@") - 1) == 0)
                        {
                            derivedFromException = true;
                        }
                        else if (strncmp(descriptor->raw_name(), ".PEAD", AuArraySize(".PEAD") - 1) == 0)
                        {
                            auto possibleStringPointer = reinterpret_cast<const char **>(ExceptionInfo->ExceptionRecord->ExceptionInformation[1]);
                            if (IsReadable(possibleStringPointer))
                            {
                                auto string = *possibleStringPointer;
                                if (IsReadable(string) && (strnlen(string, 4096) < 1024))
                                {
                                    if (suffix.size()) suffix += ", ";
                                    suffix += string;
                                }
                            }
                        }
                        else if (strncmp(descriptor->raw_name(), kStringRawName.data(), kStringRawName.size()) == 0)
                        {
                            auto possibleStdStringPointer = reinterpret_cast<std::string *>(ExceptionInfo->ExceptionRecord->ExceptionInformation[1]);
                            if (IsReadable(possibleStdStringPointer))
                            {
                                auto string = *possibleStdStringPointer;
                                if (suffix.size()) suffix += ", ";
                                suffix += string;
                            }
                        }
                    }

                    if (derivedFromException)
                    {
                        auto fatanthony = exception->what();
                        if (fatanthony)
                        {
                            entry.wincxx.str += AuString("\r\n    ") + fatanthony;
                        }
                    }

                    if (suffix.size())
                    {
                        entry.wincxx.str += AuString("\r\n    ") + suffix;
                    }
                }
            }

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