/***
    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"

namespace Aurora
{
    static bool gShouldResPathDoNothing {};

    void InitNTAddresses()
    {
    #if defined(AURORA_PLATFORM_WIN32)
    #define ADD_LOAD_LIB(name)                                          \
        auto h ## name       = GetModuleHandleW(k## name ## DllName);   \
        if (!h ## name)                                                 \
        {                                                               \
            h ## name        = LoadLibraryW(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);

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

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

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

    #define ADD_GET_PROC_INTERNAL_MAP(name, proc, symbol)                                              \
        if (h ## name)                                                                                 \
        {                                                                                              \
            p ## proc = AuReinterpretCast<decltype(p ## proc)>(GetProcAddress(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_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_INTERNAL_MAP(AdvancedApi, RtlGenRandom, SystemFunction036)
            
        ADD_GET_PROC(BCrypt, BCryptGenRandom)

        ADD_GET_PROC(Theme, SetWindowTheme)

        ADD_GET_PROC(Shell, SHGetKnownFolderPath)

        
        if (pNtCreateKeyedEvent &&
            Threading::Primitives::gKeyedEventHandle == INVALID_HANDLE_VALUE)
        {
            if (!gUseNativeWaitCondvar)
            {
                SysAssert(pNtCreateKeyedEvent);
                pNtCreateKeyedEvent(&Threading::Primitives::gKeyedEventHandle, -1, NULL, 0);
            }
        }
        
    #else
        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;

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