/***
    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
{
    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, 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(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)
        
    #else
        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()
    {
        static AuUInt uCounter {};
        ULONG ullActualResolution {};

        if (gShouldResPathDoNothing)
        {
            return;
        }

        if (gRuntimeConfig.threadingConfig.bEnableAgrSchedulingRatelimit)
        {
            AuUInt uCount {};
            if (!(((uCount = uCounter++)) & (1ull << (gRuntimeConfig.threadingConfig.bWinXpThrough7BlazeOptimizerPower + 5ull))))
            {
                if (uCount != 0)
                {
                    return;
                }
            }
            else
            {

                uCounter = 0;
            }

        }

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