#include "FSR3Upscaler_DX12.h"

#include <d3d12.h>
#include <dxgi.h>

// TODO: find a better place to put this, including the UnityLog and FFX API Loader includes
#include <sstream>
#include "ffx_api/ffx_api_loader.h"
#include "UnityPluginAPI/IUnityLog.h"
static IUnityLog* s_Log = nullptr;
static HMODULE s_ffxModule = nullptr;
static ffxFunctions s_ffxFunctions{};
static bool LoadFidelityFXLibrary(_In_ LPCWSTR lpLibFileName, void* device)
{
    s_ffxModule = LoadLibrary(lpLibFileName);
    if (s_ffxModule == nullptr)
    {
        UNITY_LOG_ERROR(s_Log, "Failed to load FidelityFX library!");
        return false;
    }

    ffxLoadFunctions(&s_ffxFunctions, s_ffxModule);
    if (s_ffxFunctions.CreateContext == nullptr)
    {
        UNITY_LOG_ERROR(s_Log, "Failed to load FidelityFX library functions!");
        return false;
    }

    // Check that we can actually create an upscaler with this library
    ffx::QueryDescGetVersions versionQuery{};
    versionQuery.createDescType = FFX_API_CREATE_CONTEXT_DESC_TYPE_UPSCALE;
    versionQuery.device = device;
    uint64_t versionCount = 0;
    versionQuery.outputCount = &versionCount;
    s_ffxFunctions.Query(nullptr, &versionQuery.header);

    if (versionCount == 0)
    {
        UNITY_LOG_ERROR(s_Log, "Failed to load FidelityFX upscaler versions!");
        return false;
    }

    // Obtain the default upscaler version and log its name
    std::vector<const char*> versionNames;
    std::vector<uint64_t> versionIds;
    versionIds.resize(versionCount);
    versionNames.resize(versionCount);
    versionQuery.versionIds = versionIds.data();
    versionQuery.versionNames = versionNames.data();
    s_ffxFunctions.Query(nullptr, &versionQuery.header);

    std::stringstream ss;
    ss << "Loaded FidelityFX upscaler: FSR " << versionNames[0];
    UNITY_LOG(s_Log, ss.str().c_str());

    return true;
}

bool FSR3Upscaler_DX12::Init()
{
    ID3D12Device* device = m_GraphicsDevice->GetDevice();
    if (device == nullptr)
        return false;

    m_DX12BackendDesc.device = device;
    m_FrameFenceEventHandle = CreateEventEx(nullptr, nullptr, 0, EVENT_ALL_ACCESS);

    return LoadFidelityFXLibrary(TEXT("amd_fidelityfx_dx12.dll"), device);
}

void FSR3Upscaler_DX12::Shutdown()
{
    UpscalerGraphicsDevice::Shutdown();     // TODO: change to FSR3Upscaler_FFXBase::Shutdown

    std::lock_guard<std::mutex> lock(m_Mutex);  // TODO: see if we can rearrange this to only require mutexes in UpscalerGraphicsDevice

    if (s_ffxModule != nullptr) // TODO: move to base FSR3Upscaler_FFXBase class
    {
        FreeLibrary(s_ffxModule);
        s_ffxModule = nullptr;
    }

    if (m_FrameFenceEventHandle != nullptr)
    {
        CloseHandle(m_FrameFenceEventHandle);
        m_FrameFenceEventHandle = nullptr;
    }

    m_DX12BackendDesc.device = nullptr;
}

bool FSR3Upscaler_DX12::InitFeature(FSR3Feature_DX12& feature, const FSR3CommandInitializationData* initData)
{
    ffx::CreateContextDescUpscale createUpscaling;
    createUpscaling.maxUpscaleSize = { initData->displaySizeWidth, initData->displaySizeHeight };
    createUpscaling.maxRenderSize = { initData->maxRenderSizeWidth, initData->maxRenderSizeHeight };
    createUpscaling.flags = initData->flags;

    return FFX_API_RETURN_OK == s_ffxFunctions.CreateContext(&feature.upscalingContext, ffx::LinkHeaders(createUpscaling.header, m_DX12BackendDesc.header), nullptr);
}

void FSR3Upscaler_DX12::SetTexture(FSR3TextureDesc* textureDesc, UnityTextureID textureID)
{
    textureDesc->image = (intptr_t)m_GraphicsDevice->TextureFromNativeTexture(textureID);
}

void FSR3Upscaler_DX12::Execute(FSR3Feature_DX12& feature, const FSR3CommandExecutionData* execData)
{
    ffx::DispatchDescUpscale dispatchUpscale{};

    UnityGraphicsD3D12RecordingState state;
    m_GraphicsDevice->CommandRecordingState(&state);
    dispatchUpscale.commandList = state.commandList;

    // Keep track of which frame this dispatch is happening on, so we can verify when it's finished
    feature.dispatchFrameValue = m_GraphicsDevice->GetNextFrameFenceValue();

    dispatchUpscale.color = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.colorInput.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
    dispatchUpscale.depth = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.depth.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
    dispatchUpscale.motionVectors = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.motionVectors.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
    dispatchUpscale.exposure = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.exposureTexture.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
    dispatchUpscale.reactive = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.reactiveMask.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
    dispatchUpscale.transparencyAndComposition = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.transparencyMask.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
    dispatchUpscale.output = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.colorOutput.image, FFX_API_RESOURCE_STATE_UNORDERED_ACCESS);

    if (dispatchUpscale.reactive.resource == nullptr)
        dispatchUpscale.reactive = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.biasColorMask.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);

    dispatchUpscale.jitterOffset.x = execData->jitterOffsetX;
    dispatchUpscale.jitterOffset.y = execData->jitterOffsetY;
    dispatchUpscale.motionVectorScale.x = execData->MVScaleX;
    dispatchUpscale.motionVectorScale.y = execData->MVScaleY;
    dispatchUpscale.reset = execData->reset;
    dispatchUpscale.enableSharpening = execData->enableSharpening;
    dispatchUpscale.sharpness = execData->sharpness;
    dispatchUpscale.frameTimeDelta = execData->frameTimeDelta;
    dispatchUpscale.preExposure = execData->preExposure;
    dispatchUpscale.renderSize.width = execData->renderSizeWidth;
    dispatchUpscale.renderSize.height = execData->renderSizeHeight;
    dispatchUpscale.upscaleSize.width = feature.upscaleSizeWidth;
    dispatchUpscale.upscaleSize.height = feature.upscaleSizeHeight;
    dispatchUpscale.cameraFovAngleVertical = execData->cameraFovAngleVertical;

    if (feature.flags & FFX_UPSCALE_ENABLE_DEPTH_INVERTED)
    {
        dispatchUpscale.cameraFar = execData->cameraNear;
        dispatchUpscale.cameraNear = execData->cameraFar;
    }
    else
    {
        dispatchUpscale.cameraFar = execData->cameraFar;
        dispatchUpscale.cameraNear = execData->cameraNear;
    }

    s_ffxFunctions.Dispatch(&feature.upscalingContext, &dispatchUpscale.header);
}

bool FSR3Upscaler_DX12::IsValidFeature(FSR3Feature_DX12& feature)
{
    return feature.upscalingContext != nullptr;
}

void FSR3Upscaler_DX12::AwaitEndOfFrame(uint64_t frameValue)
{
    // Thanks to https://alextardif.com/D3D11To12P1.html for explaining how fences work and how to make the CPU wait for a command queue to complete.
    ID3D12Fence* frameFence = m_GraphicsDevice->GetFrameFence();
    if (frameValue > frameFence->GetCompletedValue())
    {
        frameFence->SetEventOnCompletion(frameValue, m_FrameFenceEventHandle);
        WaitForSingleObjectEx(m_FrameFenceEventHandle, INFINITE, false);
    }
}

void FSR3Upscaler_DX12::DestroyContext(FSR3Feature_DX12& feature)
{
    s_ffxFunctions.DestroyContext(&feature.upscalingContext, nullptr);
}