#include <d3d9.h>

#include "avs/game.h"
#include "external/nvenc/nvEncodeAPI.h"
#include "hooks/libraryhook.h"
#include "hooks/graphics/backends/d3d9/d3d9_device.h"
#include "util/detour.h"
#include "util/libutils.h"
#include "util/logging.h"

#include "nvenc_hook.h"

#ifdef SPICE64

typedef NVENCSTATUS(NVENCAPI *NvEncodeAPICreateInstance_Type)(NV_ENCODE_API_FUNCTION_LIST*);
static NvEncodeAPICreateInstance_Type NvEncodeAPICreateInstance_orig = nullptr;
static PNVENCOPENENCODESESSIONEX nvEncOpenEncodeSessionEx_orig = nullptr;
static BOOL initialized = false;

namespace nvenc_hook {

    NVENCSTATUS nvEncOpenEncodeSessionEx_hook(
        NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS *openSessionExParams, 
        void **encoder
    ) {
        WrappedIDirect3DDevice9 *wrappedDevice;
        try {
            wrappedDevice = (WrappedIDirect3DDevice9*)openSessionExParams->device;
            // log_misc("nvenc_hook", 
            //     "nvEncOpenEncodeSessionEx hook hit (wrapped: {}) (real: {})", 
            //     fmt::ptr(wrappedDevice), 
            //     fmt::ptr(wrappedDevice->pReal)
            // );
            openSessionExParams->device = wrappedDevice->pReal;
        } catch (const std::exception &ex) {
            // log_misc("nvenc_hook", "Cannot cast to WrappedIDirect3DDevice9. D3D9 hooks might be disabled.");
        }
        return nvEncOpenEncodeSessionEx_orig(openSessionExParams, encoder);
    }

    NVENCSTATUS NvEncodeAPICreateInstance_hook(NV_ENCODE_API_FUNCTION_LIST *pFunctionList) {        
        // log_misc("nvenc_hook", "NvEncodeAPICreateInstance hook hit");
        auto status = NvEncodeAPICreateInstance_orig(pFunctionList);

        // The game will call NvEncodeAPICreateInstance multiple times
        // Using a flag to avoid creating trampoline repeatedly
        if (!initialized) {
            // hook functions
            detour::trampoline_try(
                pFunctionList->nvEncOpenEncodeSessionEx,
                nvEncOpenEncodeSessionEx_hook,
                &nvEncOpenEncodeSessionEx_orig);
            // log_misc("nvenc_hook", "Created hook for nvEncOpenEncodeSessionEx");
            initialized = true;
        }

        return status;
    }

    void initialize() {
        HMODULE nvenc = libutils::try_library("nvEncodeAPI64.dll");
        if (nvenc == nullptr) {
            log_warning("nvenc_hook", "Failed to find nvEncodeAPI64.dll");
            return;
        }
        bool success = detour::trampoline_try(
            (NvEncodeAPICreateInstance_Type)libutils::try_proc(nvenc, "NvEncodeAPICreateInstance"), 
            NvEncodeAPICreateInstance_hook, 
            &NvEncodeAPICreateInstance_orig
            );
        if (success) {
            log_misc("nvenc_hook", "Created hook for NvEncodeAPICreateInstance");
        } else {
            log_warning("nvenc_hook", "Failed to hook NvEncodeAPICreateInstance");
        }
    }
}

#else

namespace nvenc_hook {
    void initialize() {
        return;
    }
}

#endif