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

    File: AuProcAddresses.NT.cpp
    Date: 2023-2-16
    Author: Reece
***/
#define _WIN32_WINNT 0x0602
#include <RuntimeInternal.hpp>
#include "Source/Threading/Primitives/AuConditionMutex.NT.hpp"

#define AURORA_HAS_LOAD_PGKD
#define AURORA_HAS_GET_PROC_NONWIN32

namespace Aurora
{
    static bool gShouldResPathDoNothing {};

#if !defined(AURORA_PLATFORM_WIN32)
    HMODULE UWPLibraryW(LPCWSTR lpLibFileName);

    FARPROC UWPProcAddress(HMODULE hModule,
                           LPCSTR lpProcName);
#endif

#if !defined(AURORA_DLL_BLACKLIST)
    static AuUInt32 gBlockedDLLTable[64] {};
#else
    static AuUInt32 gBlockedDLLTable[64] { AURORA_DLL_BLACKLIST };
#endif

    AUKN_SYM bool AddBlockedDLL(const char *pString)
    {
        if (!pString)
        {
            return false;
        }

        auto uHash = AuFnv1a32Runtime(pString, strlen(pString));

        for (AU_ITERATE_N(i, AuArraySize(gBlockedDLLTable)))
        {
            if (gBlockedDLLTable[i])
            {
                continue;
            }

            gBlockedDLLTable[i] = uHash;
            return true;
        }

        return false;
    }

    static bool IsBlocked(const char *pString)
    {
        auto uHash = AuFnv1a32Runtime(pString, strlen(pString));

        for (AU_ITERATE_N(i, AuArraySize(gBlockedDLLTable)))
        {
            auto uCurrent = gBlockedDLLTable[i];

            if (uCurrent == uHash)
            {
                return true;
            }
            else if (uCurrent == 0)
            {
                return false;
            }
        }

        return false;
    }

    void InitNTAddresses()
    {
    #if defined(AURORA_PLATFORM_WIN32)
        pLoadLibraryW = LoadLibraryW;
        pGetProcAddress = GetProcAddress;

    #define ADD_LOAD_LIB(name)                                          \
                                                                        \
        HMODULE h ## name {};                                           \
        if (!IsBlocked(#name))                                          \
        {                                                               \
            h ## name = GetModuleHandleW(k## name ## DllName);          \
            if (!h ## name)                                             \
            {                                                           \
                h ## name        = pLoadLibraryW(k## name ## DllName);  \
            }                                                           \
        }

        ADD_LOAD_LIB(Kernel32);
        ADD_LOAD_LIB(Nt);
        ADD_LOAD_LIB(KernelBase);
        ADD_LOAD_LIB(Sync);
        ADD_LOAD_LIB(WS2);
        ADD_LOAD_LIB(AdvancedApi);
        ADD_LOAD_LIB(BCrypt);
        ADD_LOAD_LIB(Theme);
        ADD_LOAD_LIB(Shell);
        ADD_LOAD_LIB(PSAPILegacy);
        ADD_LOAD_LIB(DbgHelper);
        ADD_LOAD_LIB(WinTrust);
        ADD_LOAD_LIB(IPHelper);
        ADD_LOAD_LIB(COM);
        ADD_LOAD_LIB(User32);

    #define ADD_GET_PROC(name, proc)                                                                       \
        if (!IsBlocked(#proc))                                                                             \
        {                                                                                                  \
            if (h ## name)                                                                                 \
            {                                                                                              \
                p ## proc = AuReinterpretCast<decltype(p ## proc)>(pGetProcAddress(h ## name, #proc));     \
            }                                                                                              \
        }

    #define ADD_GET_PROC_BI(name, name2, proc)                                                             \
        if (!IsBlocked(#proc))                                                                             \
        {                                                                                                  \
            p ## proc = nullptr;                                                                           \
            if (h ## name)                                                                                 \
            {                                                                                              \
                p ## proc = AuReinterpretCast<decltype(p ## proc)>(pGetProcAddress(h ## name, #proc));     \
            }                                                                                              \
            if (!p ## proc)                                                                                \
            {                                                                                              \
                if (h ## name2)                                                                            \
                {                                                                                          \
                    p ## proc = AuReinterpretCast<decltype(p ## proc)>(pGetProcAddress(h ## name2, #proc));\
                }                                                                                          \
            }                                                                                              \
        }

    #define ADD_GET_PROC_BI2(name, name2, proc, proc2)                                                     \
        if (!IsBlocked(#proc))                                                                             \
        {                                                                                                  \
            p ## proc2 = nullptr;                                                                          \
            if (h ## name)                                                                                 \
            {                                                                                              \
                p ## proc2 = AuReinterpretCast<decltype(p ## proc2)>(pGetProcAddress(h ## name, #proc));   \
            }                                                                                              \
            if (!p ## proc2)                                                                               \
            {                                                                                              \
                if (h ## name2)                                                                            \
                {                                                                                          \
                    p ## proc2 = AuReinterpretCast<decltype(p ## proc2)>(pGetProcAddress(h ## name2, #proc2)); \
                }                                                                                          \
            }                                                                                              \
        }

    #define ADD_GET_PROC_INTERNAL_MAP(name, proc, symbol)                                                  \
        if (!IsBlocked(#proc))                                                                             \
        {                                                                                                  \
            if (h ## name)                                                                                 \
            {                                                                                              \
                p ## proc = AuReinterpretCast<decltype(p ## proc)>(pGetProcAddress(h ## name, #symbol));   \
            }                                                                                              \
        }

        if (pRtlGetVersion)
        {
            return;
        }

        ADD_GET_PROC(Nt, RtlGetVersion)
        ADD_GET_PROC(Nt, NtDelayExecution)
        ADD_GET_PROC(Nt, NtWaitForKeyedEvent)
        ADD_GET_PROC(Nt, NtReleaseKeyedEvent)
        ADD_GET_PROC(Nt, NtOpenKeyedEvent)
        ADD_GET_PROC(Nt, NtCreateKeyedEvent)
        ADD_GET_PROC(Nt, RtlWaitOnAddress)
        ADD_GET_PROC(Nt, RtlWakeByAddressAll)
        ADD_GET_PROC(Nt, RtlWakeAddressSingle)
        ADD_GET_PROC(Nt, ZwSetTimerResolution)
        ADD_GET_PROC(Nt, NtQueryInformationProcess)
        ADD_GET_PROC(Nt, NtNotifyChangeDirectoryFile)
        ADD_GET_PROC(Nt, NtTerminateProcess)
            
        ADD_GET_PROC_BI(Kernel32, KernelBase, VirtualAlloc2)
        ADD_GET_PROC_BI(Kernel32, KernelBase, MapViewOfFile3)
        ADD_GET_PROC_BI(Kernel32, KernelBase, UnmapViewOfFile2)
        ADD_GET_PROC_BI(Kernel32, KernelBase, CreateFileW)
        ADD_GET_PROC_BI(Kernel32, KernelBase, CreateFile2W)

        ADD_GET_PROC(Kernel32, GetSystemCpuSetInformation)
        ADD_GET_PROC(Kernel32, GetLogicalProcessorInformation)
        ADD_GET_PROC(Kernel32, SetThreadDescription)
        ADD_GET_PROC(Kernel32, SetThreadInformation)
        ADD_GET_PROC(Kernel32, SetThreadSelectedCpuSets)
        ADD_GET_PROC(Kernel32, PrefetchVirtualMemory)
        ADD_GET_PROC(Kernel32, SetThreadGroupAffinity)
        ADD_GET_PROC(Kernel32, FindFirstStreamW)
        ADD_GET_PROC(Kernel32, FindNextStreamW)
        ADD_GET_PROC(Kernel32, FindClose)
        ADD_GET_PROC(Kernel32, CancelIoEx)
        ADD_GET_PROC(Kernel32, CancelSynchronousIo)
        ADD_GET_PROC(Kernel32, SetFileInformationByHandle)
        ADD_GET_PROC(Kernel32, GetLocaleInfoEx)
        ADD_GET_PROC(Kernel32, LCIDToLocaleName)
        ADD_GET_PROC(Kernel32, GetLocaleInfoW)
        ADD_GET_PROC(Kernel32, GetThreadId)
        ADD_GET_PROC(Kernel32, VerifyVersionInfoW)
        ADD_GET_PROC(Kernel32, VerSetConditionMask)
        ADD_GET_PROC(Kernel32, QueryPerformanceCounter)
        ADD_GET_PROC(Kernel32, QueryPerformanceFrequency)
        ADD_GET_PROC(Kernel32, RemoveDllDirectory)
        ADD_GET_PROC(Kernel32, AddDllDirectory)
            
        ADD_GET_PROC_BI2(Kernel32, PSAPILegacy, K32GetProcessMemoryInfo, GetProcessMemoryInfo)
            
        ADD_GET_PROC(Sync, WaitOnAddress)
        ADD_GET_PROC(Sync, WakeByAddressSingle)
        ADD_GET_PROC(Sync, WakeByAddressAll)
          
        ADD_GET_PROC(DbgHelper, UnDecorateSymbolName)
        ADD_GET_PROC(DbgHelper, MiniDumpWriteDump)
        ADD_GET_PROC(DbgHelper, SymInitialize)
        ADD_GET_PROC(DbgHelper, SymGetModuleBase64)
        ADD_GET_PROC(DbgHelper, SymGetLineFromAddr64)
        ADD_GET_PROC(DbgHelper, SymFunctionTableAccess64)
        ADD_GET_PROC(DbgHelper, StackWalk64)

        ADD_GET_PROC(WS2, GetAddrInfoExW)
        ADD_GET_PROC(WS2, GetAddrInfoExCancel)
        ADD_GET_PROC(WS2, FreeAddrInfoExW)
        ADD_GET_PROC(WS2, getaddrinfo)
        ADD_GET_PROC(WS2, freeaddrinfo)

        ADD_GET_PROC(WinTrust, WinVerifyTrust)
            
        ADD_GET_PROC(IPHelper, GetAdaptersAddresses)
        ADD_GET_PROC(IPHelper, GetAdaptersInfo)
        
        ADD_GET_PROC(AdvancedApi, CryptAcquireContextW)
        ADD_GET_PROC(AdvancedApi, CryptReleaseContext)
        ADD_GET_PROC(AdvancedApi, CryptGenRandom)

        ADD_GET_PROC(AdvancedApi, RegSetValueExW)
        ADD_GET_PROC(AdvancedApi, RegCloseKey)
        ADD_GET_PROC(AdvancedApi, RegOpenKeyExW)
        ADD_GET_PROC(AdvancedApi, RegQueryValueExW)
            
        ADD_GET_PROC(AdvancedApi, DeregisterEventSource)
        ADD_GET_PROC(AdvancedApi, RegisterEventSourceW)
        ADD_GET_PROC(AdvancedApi, ReportEventW)
            
        ADD_GET_PROC(AdvancedApi, SetEntriesInAclA)
        ADD_GET_PROC(AdvancedApi, AllocateAndInitializeSid)
        ADD_GET_PROC(AdvancedApi, SetNamedSecurityInfoW)
        ADD_GET_PROC(AdvancedApi, FreeSid)
            
        ADD_GET_PROC_INTERNAL_MAP(AdvancedApi, RtlGenRandom, SystemFunction036)
            
        ADD_GET_PROC(BCrypt, BCryptGenRandom)

        ADD_GET_PROC(Theme, SetWindowTheme)
        
        ADD_GET_PROC(COM, CoInitializeEx)
        ADD_GET_PROC(COM, CoUninitialize)
        ADD_GET_PROC(COM, CoTaskMemFree)

        ADD_GET_PROC(Shell, SHGetKnownFolderPath)
        ADD_GET_PROC(Shell, SHGetFolderPathA)
        ADD_GET_PROC(Shell, CommandLineToArgvW)
        ADD_GET_PROC(Shell, ShellExecuteW)

        ADD_GET_PROC(User32, GetClipboardData);
        ADD_GET_PROC(User32, MapVirtualKeyA);
        ADD_GET_PROC(User32, CloseClipboard);
        ADD_GET_PROC(User32, GetWindowThreadProcessId);
        ADD_GET_PROC(User32, SendMessageA);
        ADD_GET_PROC(User32, EnumThreadWindows);
        ADD_GET_PROC(User32, DispatchMessageW);
        ADD_GET_PROC(User32, TranslateMessage);
        ADD_GET_PROC(User32, MsgWaitForMultipleObjectsEx);
        ADD_GET_PROC(User32, MsgWaitForMultipleObjects);
        ADD_GET_PROC(User32, PeekMessageW);
        ADD_GET_PROC(User32, SetPropW);
        ADD_GET_PROC(User32, OpenClipboard);
        
        if (pNtCreateKeyedEvent &&
            Threading::Primitives::gKeyedEventHandle == INVALID_HANDLE_VALUE)
        {
            if (!gUseNativeWaitCondvar)
            {
                SysAssert(pNtCreateKeyedEvent);
                pNtCreateKeyedEvent(&Threading::Primitives::gKeyedEventHandle, -1, NULL, 0);
            }
        }

    #else
        pLoadLibraryW        = UWPLibraryW;
        pGetProcAddress      = UWPProcAddress;

        pCreateFile2W        = CreateFile2FromAppW;

        pWaitOnAddress       = WaitOnAddress;
        pWakeByAddressSingle = WakeByAddressSingle;
        pWakeByAddressAll    = WakeByAddressAll;
        pVirtualAlloc2       = VirtualAlloc2FromApp;
        pMapViewOfFile3      = MapViewOfFile3FromApp;

        pCancelIoEx          = CancelIoEx;
        pCancelSynchronousIo = CancelSynchronousIo;

        pGetLocaleInfoEx     = GetLocaleInfoEx;

        pLCIDToLocaleName    = LCIDToLocaleName;

        pSetFileInformationByHandle = SetFileInformationByHandle;

        pFindClose           = FindClose;

        pGetSystemCpuSetInformation = GetSystemCpuSetInformation;
        pGetLogicalProcessorInformation = GetLogicalProcessorInformation;

        pSetThreadInformation = SetThreadInformation;
        pSetThreadDescription = SetThreadDescription;
        pSetThreadSelectedCpuSets = SetThreadSelectedCpuSets;

        pGetAddrInfoExCancel = GetAddrInfoExCancel;

        pPrefetchVirtualMemory = PrefetchVirtualMemory;

        pQueryPerformanceCounter   = decltype(pQueryPerformanceCounter)(QueryPerformanceCounter);
        pQueryPerformanceFrequency = decltype(pQueryPerformanceFrequency)(QueryPerformanceFrequency);

        // https://github.com/LWJGL/lwjgl3/blob/master/modules/lwjgl/remotery/src/main/c/Remotery.c#L1188
        // Xbox main SDK has a better API we should use
        // So... 
        // TODO: Xbox One and later: https://github.com/microsoft/Xbox-ATG-Samples/blob/main/XDKSamples/Graphics/AdvancedESRAM12/PageAllocator.cpp#L193-L206
        // Require AuProcess for that given target

        pUnmapViewOfFile2    = UnmapViewOfFile2; // < isn't portable but
        // "This topic lists the Win32 APIs that are part of the Universal Windows Platform (UWP) and that are implemented by all Windows 10 devices."
        // UnmapViewOfFile2 -> Introduced into api-ms-win-core-memory-l1-1-5.dll in 10.0.17134.

        pNtDelayExecution    = nullptr /* ... (you dont need it, but it'll help a ton) */;
    #endif

        gUseNativeWaitMutex = (pWaitOnAddress &&
                               !gRuntimeConfig.threadingConfig.bPreferNt51XpMutexesOver8 &&
                               (pRtlWaitOnAddress || AuBuild::kCurrentPlatform != AuBuild::EPlatform::ePlatformWin32)) || 
                              !pNtWaitForKeyedEvent;

        gUseNativeWaitCondvar = (pWaitOnAddress &&
                                 !gRuntimeConfig.threadingConfig.bPreferNt51XpCondvarsOver8 &&
                                 (pRtlWaitOnAddress || AuBuild::kCurrentPlatform != AuBuild::EPlatform::ePlatformWin32)) ||
                                !pNtWaitForKeyedEvent;

        gUseNativeWaitSemapahore = bool(pWaitOnAddress);
    }

    void Win32DropInit()
    {
        gShouldResPathDoNothing = 
            (AuBuild::kCurrentPlatform != AuBuild::EPlatform::ePlatformWin32) ||
            (!gRuntimeConfig.threadingConfig.bEnableAggressiveScheduling &&
            AuSwInfo::IsWindows10OrGreater());

        gUseFastFail = AuSwInfo::IsWindows8Point1OrGreater()
        #if defined(PF_FASTFAIL_AVAILABLE)
            && ::IsProcessorFeaturePresent(PF_FASTFAIL_AVAILABLE)
        #endif
            ;
    }

    void Win32DropSchedulerResolution()
    {
        ULONG ullActualResolution {};

        if (gShouldResPathDoNothing)
        {
            return;
        }

        if (gRuntimeConfig.threadingConfig.bEnableAgrSchedulingRatelimit)
        {
            static Aurora::Utility::RateLimiter limiter;
            if (!limiter.nsTimeStep)
            {
                limiter.SetNextStep(AuMSToNS<AuUInt64>(gRuntimeConfig.threadingConfig.bWinXpThrough7BlazeOptimizerPower));
            }
            else if (!limiter.CheckExchangePass())
            {
                return;
            }
        }

        if (pZwSetTimerResolution)
        {
            auto uRet = pZwSetTimerResolution(1, true, &ullActualResolution);

            if (uRet == 0)
            {
                return;
            }
            else if (uRet == 0xC0000245)
            {
                if ((uRet = pZwSetTimerResolution(5'000, true, &ullActualResolution)) == 0)
                {
                    return;
                }
            }
        }

        // ...SetProcessInformation?
    }

    void Win32Terminate()
    {
        if (gUseFastFail)
        {
            __fastfail('fokd');
        }
        else
        {
            if (pNtTerminateProcess)
            {
                pNtTerminateProcess((HANDLE)-1, 0x0);
            }

            ::TerminateProcess(::GetCurrentProcess(), 0);
        }
    }

    HANDLE Win32Open(LPCWSTR lpFileName,
                     DWORD   dwDesiredAccess,
                     DWORD   dwShareMode,
                     bool    bInherit,
                     DWORD   dwCreationDisposition,
                     DWORD   dwFlags,
                     DWORD   dwAttributes
    )
    {
        SECURITY_ATTRIBUTES attrs {};
        attrs.nLength        = sizeof(attrs);
        attrs.bInheritHandle = bInherit ? TRUE : FALSE;

        bool bSpecialFlags = (dwFlags & 0x4ffff7);
        bool bSpecialAttrs = (dwAttributes & 0xFFB00008);

        if (!pCreateFile2W)
        {
            if (bSpecialFlags)
            {
                SysPushErrorFeatureMissing("Do not use Windows8+ attributes/flags under Win32Open");
                return INVALID_HANDLE_VALUE;
            }

            if (bSpecialAttrs)
            {
                SysPushErrorFeatureMissing("Do not use Windows8+ attributes/flags under Win32Open");
                return INVALID_HANDLE_VALUE;
            }
        }

        if (!dwAttributes)
        {
            dwAttributes = FILE_ATTRIBUTE_NORMAL;
        }

        if (pCreateFileW &&
            !bSpecialFlags &&
            !bSpecialAttrs)
        {
            return pCreateFileW(lpFileName,
                                dwDesiredAccess,
                                dwShareMode,
                                &attrs,
                                dwCreationDisposition,
                                dwFlags | dwAttributes,
                                NULL);
        }

        if (pCreateFile2W)
        {
            _CREATEFILE2_EXTENDED_PARAMETERS params {};
            bool bRead {};
            HANDLE hHandle {};

            params.dwSize = sizeof(_CREATEFILE2_EXTENDED_PARAMETERS);
            params.dwFileFlags = dwFlags;
            params.dwFileAttributes = dwAttributes;
            params.lpSecurityAttributes = &attrs;

            if ((bRead = (::wcscmp(lpFileName, L"CONIN$") == 0)) ||
                (::wcscmp(lpFileName, L"CONOUT$") == 0) ||
                (::wcscmp(lpFileName, L"CONERR$") == 0))
            {
                if (!bRead)
                {
                    dwDesiredAccess = GENERIC_READ | GENERIC_WRITE;
                }
                else
                {
                    dwDesiredAccess = GENERIC_READ;
                }

                if ((hHandle = pCreateFile2W(lpFileName,
                                             dwDesiredAccess,
                                             dwShareMode,
                                             dwCreationDisposition,
                                             &params)) != INVALID_HANDLE_VALUE)
                {
                    return hHandle;
                }

                lpFileName = L"CON";
            }

            return pCreateFile2W(lpFileName,
                                 dwDesiredAccess,
                                 dwShareMode,
                                 dwCreationDisposition,
                                 &params);
        }

        return INVALID_HANDLE_VALUE;
    }

#if !defined(AURORA_PLATFORM_WIN32)
    HMODULE UWPLibraryW(LPCWSTR lpLibFileName)
    {
        auto pos = std::basic_string_view<wchar_t>(lpLibFileName).find_last_of('/');
        if (pos == std::basic_string_view<wchar_t>::npos)
        {
        #if defined(AURORA_HAS_LOAD_PGKD)
            return (HMODULE)::LoadPackagedLibrary(lpLibFileName, 0);
        #else
            return (HMODULE)lpLibFileName;
        #endif
        }
        else
        {
            auto pString = &lpLibFileName[pos];
        #if defined(AURORA_HAS_LOAD_PGKD)
            return (HMODULE)::LoadPackagedLibrary(pString, 0);
        #else
            return (HMODULE)pString;
        #endif
        }
    }

    FARPROC UWPProcAddress(HMODULE hModule,
                           LPCSTR lpProcName)
    {
        if (!hModule)
        {
            return nullptr;
        }
    
    #if defined(AURORA_HAS_GET_PROC_NONWIN32)
        return GetProcAddress(hModule, lpProcName);
    #endif

        return nullptr;
    }
#endif
}