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

    File: AuProcAddresses.NT.cpp
    Date: 2023-2-16
    Author: Reece
***/
#include <RuntimeInternal.hpp>

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

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

        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_BI(Kernel32, KernelBase, VirtualAlloc2)
        ADD_GET_PROC_BI(Kernel32, KernelBase, MapViewOfFile3)
        ADD_GET_PROC_BI(Kernel32, KernelBase, UnmapViewOfFile2)
        
        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(Sync, WaitOnAddress)
        ADD_GET_PROC(Sync, WakeByAddressSingle)
        ADD_GET_PROC(Sync, WakeByAddressAll)
            
        ADD_GET_PROC(WS2, GetAddrInfoExCancel)

        ADD_GET_PROC(AdvancedApi, CryptAcquireContextW)
        ADD_GET_PROC(AdvancedApi, CryptReleaseContext)
        ADD_GET_PROC(AdvancedApi, CryptGenRandom)
        
        ADD_GET_PROC(BCrypt, BCryptGenRandom)

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

        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.bPerferNt51XpCondvarsOver8 &&
                                (pRtlWaitOnAddress || AuBuild::kCurrentPlatform != AuBuild::EPlatform::ePlatformWin32)) ||
                               !pNtWaitForKeyedEvent;

        gUseNativeWaitSemapahore = bool(pWaitOnAddress);
    }
}