#include "winunwanted3dpartydllsblocker.h"
#include "vmsunwanted3dpartydllsblocker.h"

#define USE_DETOURS

#ifdef USE_DETOURS
#include "detoursunwanted3dpartydllsblocker.h"
#else
#include "../vmscl/windows/process/api_hook/vmsPeFnHook.h"
#endif

#ifdef USE_DETOURS
    std::shared_ptr<DetoursUnwanted3dPartyDllsBlocker> g_blocker;
#else
    std::shared_ptr <vmsPeFnHook> g_hooker;
std::unique_ptr <vmsUnwanted3dPartyDllsBlocker> g_hooksFurther;
#endif

    void HookAll()
{
    // Originally we had a known problem with nahimicosd.dll only.
    // (https://bugreports.qt.io/browse/QTBUG-86329)
    // But I've decided to take the list from here:
    // https://github.com/moonlight-stream/moonlight-qt/blob/master/AntiHooking/antihookingprotection.cpp
    // Just in case. I don't think this can has unwanted side effects for us.

        std::vector<std::string> unwantedDlls = {
            // These A-Volute DLLs shipped with various audio driver packages improperly handle
            // D3D9 exclusive fullscreen in a way that causes CreateDeviceEx() to deadlock.
            // https://github.com/moonlight-stream/moonlight-qt/issues/102
            "nahimicosd.dll", // ASUS Sonic Radar 3
            "ssaudioosd.dll", // SteelSeries headsets
            "ss2osd.dll", // ASUS Sonic Studio 2
            "nahimic2osd.dll",
            "nahimicmsiosd.dll",
            "nhasusphoebusosd.dll" // ASUS Phoebus

            // This DLL has been seen in several crash reports. Some Googling
            // suggests it's highly unstable and causes issues in many games.
            "ezfrd32.dll",
            "ezfrd64.dll",

            // These are the newer dList DLLs for Optimus hybrid graphics DDI.
            // https://docs.microsoft.com/en-us/windows-hardware/drivers/display/hybrid-system-ddi
            //
            // We forcefully block them from loading because Optimus has a bug that
            // deadlocks DXVA2 when we present with D3DPRESENT_DONOTWAIT. This will prevent
            // Optimus from ever using the dGPU even if the user has requested it.
            // https://github.com/moonlight-stream/moonlight-qt/issues/240
            // https://github.com/moonlight-stream/moonlight-qt/issues/235
            "nvdlist.dll",
            "nvdlistx.dll",

            // These are the older dList/AppInit DLLs for Optimus hybrid graphics DDI.
            // https://docs.microsoft.com/en-us/windows-hardware/drivers/display/hybrid-system-ddi
            //
            // These seem to cause a crash in PresentEx() in full-screen exclusive mode.
            // This block will prevent Optimus from ever using the dGPU even if the user has requested it.
            // https://github.com/moonlight-stream/moonlight-qt/issues/386
            //
            // d3d9!CSwapChain::BltToHybridPrimary+0x200:
            // 00007ffa`23f37e58 488b01          mov     rax,qword ptr [rcx] ds:00000000`00000038=????????????????
            // 00 0000004e`496ff4e0 00007ffa`23f39e2c d3d9!CSwapChain::BltToHybridPrimary+0x200
            // 01 0000004e`496ff880 00007ffa`23ee39ce d3d9!CSwapChain::FlipToSurface+0x15c
            // 02 0000004e`496ff900 00007ffa`23f4dd75 d3d9!CSwapChain::PresentMain+0x3e13e
            // 03 0000004e`496ffab0 00007ffa`23f4dccd d3d9!CBaseDevice::PresentMain+0x9d
            // 04 0000004e`496ffb00 00007ff7`8e31016f d3d9!CBaseDevice::PresentEx+0xbd
            // 05 0000004e`496ffb50 00007ff7`8e30df1e Moonlight!DXVA2Renderer::renderFrame+0x61f [C:\moonlight-qt\app\streaming\video\ffmpeg-renderers\dxva2.cpp @ 1035]
            // 06 0000004e`496ffd50 00007ff7`8e30e46a Moonlight!Pacer::renderFrame+0x3e [C:\moonlight-qt\app\streaming\video\ffmpeg-renderers\pacer\pacer.cpp @ 265]
            // 07 (Inline Function) --------`-------- Moonlight!Pacer::renderLastFrameAndUnlock+0x197 [C:\moonlight-qt\app\streaming\video\ffmpeg-renderers\pacer\pacer.cpp @ 156]
            // 08 0000004e`496ffda0 00007ffa`14476978 Moonlight!Pacer::renderThread+0x1ca [C:\moonlight-qt\app\streaming\video\ffmpeg-renderers\pacer\pacer.cpp @ 88]
            // 09 0000004e`496ffde0 00007ffa`14476ee2 SDL2!SDL_RunThread+0x38 [C:\Users\camer\SDL\src\thread\SDL_thread.c @ 276]
            // 0a 0000004e`496ffe10 00007ffa`2aae0e82 SDL2!RunThread+0x12 [C:\Users\camer\SDL\src\thread\windows\SDL_systhread.c @ 83]
            // 0b 0000004e`496ffe40 00007ffa`2d627bd4 ucrtbase!thread_start<unsigned int (__cdecl*)(void *),1>+0x42
            // 0c 0000004e`496ffe70 00007ffa`2da4ce51 kernel32!BaseThreadInitThunk+0x14
            // 0d 0000004e`496ffea0 00000000`00000000 ntdll!RtlUserThreadStart+0x21
            //
            "nvinit.dll",
            "nvinitx.dll",

            // In some unknown circumstances, RTSS tries to hook in the middle of an instruction, leaving garbage
            // code inside d3d9.dll that causes a crash when executed:
            //
            // 0:000> u
            // d3d9!D3D9GetCurrentOwnershipMode+0x5d:
            // 00007ff8`95b95861 9b              wait
            // 00007ff8`95b95862 a7              cmps    dword ptr [rsi],dword ptr [rdi]   <--- crash happens here
            // 00007ff8`95b95863 ff              ???
            // 00007ff8`95b95864 bfe8ca8a00      mov     edi,8ACAE8h
            // 00007ff8`95b95869 00eb            add     bl,ch
            // 00007ff8`95b9586b f1              ???
            // 00007ff8`95b9586c b808000000      mov     eax,8
            // 00007ff8`95b95871 ebe6            jmp     d3d9!D3D9GetCurrentOwnershipMode+0x55 (00007ff8`95b95859)
            //
            // Disassembling starting at the exact address of the attempted hook yields the intended jmp instruction
            //
            // 0:000> u d3d9!D3D9GetCurrentOwnershipMode+0x5c:
            // 00007ff8`95b95860 e99ba7ffbf      jmp     00007ff8`55b90000
            //
            // Since the RTSS OSD doesn't even work with DXVA2, we'll just block the hooks entirely.
            "rtsshooks.dll",
            "rtsshooks64.dll",
            };

    std::sort(unwantedDlls.begin(), unwantedDlls.end());

#ifdef USE_DETOURS
    g_blocker = std::make_shared<DetoursUnwanted3dPartyDllsBlocker>(
        std::move(unwantedDlls));
#else
    g_hooker = std::make_shared <vmsPeFnHook> ();

    g_hooksFurther = std::make_unique <vmsUnwanted3dPartyDllsBlocker> (
        g_hooker,
        std::move(unwantedDlls));

    g_hooksFurther->hook(true);
#endif
}

void UnhookAll()
{
#ifdef USE_DETOURS
    g_blocker.reset();
#else
    g_hooker->RemoveAllHooks ();
    g_hooksFurther.reset();
    g_hooker.reset();
#endif
}

BOOL APIENTRY DllMain( HMODULE hModule,
                      DWORD  ul_reason_for_call,
                      LPVOID lpReserved
                      )
{
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
        HookAll();
        DisableThreadLibraryCalls(hModule);
        break;

    case DLL_PROCESS_DETACH:
#ifndef USE_DETOURS
        //////////////////////////////////////////////////////////////////////////
        // VC2013 bug workaround
        // (https://connect.microsoft.com/VisualStudio/feedback/details/1282596)
        if (lpReserved)
            new std::shared_ptr <vmsPeFnHook> (g_hooker);
            //////////////////////////////////////////////////////////////////////////
#endif
        UnhookAll();
        break;
    }
    return TRUE;
}