#include "Upscaler.h"

#include "UnityPluginAPI/IUnityInterface.h"
#include "UnityPluginAPI/IUnityLog.h"
#include "UnityPluginAPI/IUnityGraphics.h"
#include "UnityPluginAPI/IUnityRenderingExtensions.h"

#include "FSR2Upscaler_DX11.h"
#include "FSR3Upscaler_DX12.h"
#include "FSR3Upscaler_Vulkan.h"

static const int32_t BaseEventId = 313;

static IUnityInterfaces* s_UnityInterfaces = nullptr;
static IUnityLog* s_Log = nullptr;
static IUnityGraphics* s_Graphics = nullptr;
static UnityGfxRenderer s_RendererType = kUnityGfxRendererNull;

static void UNITY_INTERFACE_API OnGraphicsDeviceEvent(UnityGfxDeviceEventType eventType);
static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data);
static void UNITY_INTERFACE_API OnSetTextureEvent(int eventID, void* data);

static Upscaler* s_Upscaler = nullptr;

// Unity plugin load event
extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API UnityPluginLoad(IUnityInterfaces* unityInterfaces)
{
    s_UnityInterfaces = unityInterfaces;
    s_Log = unityInterfaces->Get<IUnityLog>();
    s_Graphics = unityInterfaces->Get<IUnityGraphics>();

    s_Graphics->RegisterDeviceEventCallback(OnGraphicsDeviceEvent);

    // Run OnGraphicsDeviceEvent(initialize) manually on plugin load
    // to not miss the event in case the graphics device is already initialized
    OnGraphicsDeviceEvent(kUnityGfxDeviceEventInitialize);
}

// Unity plugin unload event
extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API UnityPluginUnload()
{
    s_Graphics->UnregisterDeviceEventCallback(OnGraphicsDeviceEvent);
}

static void UNITY_INTERFACE_API OnGraphicsDeviceEvent(UnityGfxDeviceEventType eventType)
{
    switch (eventType)
    {
    case kUnityGfxDeviceEventInitialize:
    {
        s_RendererType = s_Graphics->GetRenderer();
        switch (s_RendererType)
        {
        case kUnityGfxRendererD3D11:
        {
            auto* graphicsD3D11 = s_UnityInterfaces->Get<IUnityGraphicsD3D11>();
            if (graphicsD3D11 == nullptr)
            {
                UNITY_LOG_ERROR(s_Log, "Could not obtain D3D11 Graphics interface!");
                return;
            }

            s_Upscaler = new FSR2Upscaler_DX11(graphicsD3D11);
            break;
        }
        case kUnityGfxRendererD3D12:
        {
            auto* graphicsD3D12 = s_UnityInterfaces->Get<IUnityGraphicsD3D12v7>();
            if (graphicsD3D12 == nullptr)
            {
                UNITY_LOG_ERROR(s_Log, "Could not obtain D3D12 Graphics interface!");
                return;
            }

            s_Upscaler = new FSR3Upscaler_DX12(s_Log, graphicsD3D12);

            UnityD3D12PluginEventConfig eventConfig;
            eventConfig.graphicsQueueAccess = kUnityD3D12GraphicsQueueAccess_DontCare;
            eventConfig.flags = kUnityD3D12EventConfigFlag_EnsurePreviousFrameSubmission | kUnityD3D12EventConfigFlag_FlushCommandBuffers | kUnityD3D12EventConfigFlag_SyncWorkerThreads | kUnityD3D12EventConfigFlag_ModifiesCommandBuffersState;
            eventConfig.ensureActiveRenderTextureIsBound = false;
            graphicsD3D12->ConfigureEvent(BaseEventId + FSR3PluginEvent::eExecute, &eventConfig);
            break;
        }
        case kUnityGfxRendererVulkan:
        {
            auto* graphicsVulkan = s_UnityInterfaces->Get<IUnityGraphicsVulkan>();
            if (graphicsVulkan == nullptr)
            {
                UNITY_LOG_ERROR(s_Log, "Could not obtain Vulkan Graphics interface!");
                return;
            }

            s_Upscaler = new FSR3Upscaler_Vulkan(s_Log, graphicsVulkan);

            UnityVulkanPluginEventConfig eventConfig;
            eventConfig.graphicsQueueAccess = kUnityVulkanGraphicsQueueAccess_DontCare;
            eventConfig.flags = kUnityVulkanEventConfigFlag_EnsurePreviousFrameSubmission | kUnityVulkanEventConfigFlag_FlushCommandBuffers | kUnityVulkanEventConfigFlag_SyncWorkerThreads | kUnityVulkanEventConfigFlag_ModifiesCommandBuffersState;
            eventConfig.renderPassPrecondition = kUnityVulkanRenderPass_EnsureInside;
            graphicsVulkan->ConfigureEvent(BaseEventId + FSR3PluginEvent::eExecute, &eventConfig);
            break;
        }
        }

        break;
    }
    case kUnityGfxDeviceEventShutdown:
    {
        if (s_Upscaler != nullptr)
        {
            delete s_Upscaler;
            s_Upscaler = nullptr;
        }

        s_RendererType = kUnityGfxRendererNull;
        break;
    }
    case kUnityGfxDeviceEventBeforeReset:
    {
        break;
    }
    case kUnityGfxDeviceEventAfterReset:
    {
        break;
    }
    };
}

extern "C" bool UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_InitApi()
{
    if (s_Upscaler != nullptr)
    {
        return s_Upscaler->Init();
    }

    return false;
}

extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_ShutdownApi()
{
    if (s_Upscaler != nullptr)
    {
        s_Upscaler->Shutdown();
    }
}

extern "C" uint32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetDeviceVersion()
{
    return 0x0u;
}

extern "C" UnityRenderingEventAndData UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetRenderEventCallback()
{
    return OnRenderEventAndData;
}

extern "C" UnityRenderingEventAndData UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetSetTextureEventCallback()
{
    return OnSetTextureEvent;
}

extern "C" uint32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_CreateFeatureSlot()
{
    if (s_Upscaler != nullptr)
    {
        s_Upscaler->CreateFeatureSlot();
    }

    return UINT32_MAX;
}

extern "C" float UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetUpscaleRatioFromQualityMode(FSR3Quality qualityMode)
{
    switch (qualityMode)
    {
    case FSR3Quality::qQuality:
        return 1.5f;
    case FSR3Quality::qBalanced:
        return 1.7f;
    case FSR3Quality::qPerformance:
        return 2.0f;
    case FSR3Quality::qUltraPerformance:
        return 3.0f;
    default:
        return 1.0f;
    }
}

extern "C" bool UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetRenderResolutionFromQualityMode(FSR3Quality qualityMode, uint32_t displayWidth, uint32_t displayHeight, uint32_t* renderWidth, uint32_t* renderHeight)
{
    if (renderWidth == nullptr || renderHeight == nullptr)
        return false;

    float ratio = AMDUP_GetUpscaleRatioFromQualityMode(qualityMode);
    *renderWidth = (uint32_t)roundf(displayWidth / ratio);
    *renderHeight = (uint32_t)roundf(displayHeight / ratio);
    return true;
}

extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetBaseEventId()
{
    return BaseEventId;
}

// Plugin function to handle a specific rendering event
static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
{
    if (s_Upscaler == nullptr)
        return;

    // User rendering code
    switch (eventID)
    {
    case BaseEventId + FSR3PluginEvent::eDestroyFeature:
    {
        uint32_t featureSlot = (uint32_t)(intptr_t)data;
        s_Upscaler->DestroyFeature(featureSlot);
        break;
    }
    case BaseEventId + FSR3PluginEvent::eExecute:
    {
        auto* execData = (FSR3CommandExecutionData*)data;
        s_Upscaler->Execute(execData);
        break;
    }
    case BaseEventId + FSR3PluginEvent::ePostExecute:
    {
        auto* execData = (FSR3CommandExecutionData*)data;
        s_Upscaler->PostExecute(execData);
        break;
    }
    case BaseEventId + FSR3PluginEvent::eInit:
    {
        auto* initData = (FSR3CommandInitializationData*)data;
        s_Upscaler->InitFeature(initData);
        break;
    }
    }
}

static void UNITY_INTERFACE_API OnSetTextureEvent(int eventID, void* data)
{
    if (s_Upscaler == nullptr)
        return;

    auto* params = (UnityRenderingExtTextureUpdateParamsV2*)data;

    // userData = (featureId & (int) ushort.MaxValue) << 16 | (textureSlot & (int) short.MaxValue) << 1 | (clearTextureTable ? 1 : 0);
    uint32_t featureSlot = (params->userData >> 16) & 0xFFFF;
    uint32_t textureSlot = (params->userData >> 1) & 0x7FFF;
    uint32_t clearTextureTable = params->userData & 0x1;

    // User rendering code
    switch (eventID)
    {
    case kUnityRenderingExtEventUpdateTextureBeginV2:
    {
        if (clearTextureTable)
        {
            s_Upscaler->ClearTextureTable(featureSlot);
        }
        break;
    }
    case kUnityRenderingExtEventUpdateTextureEndV2:
    {
        s_Upscaler->SetTextureSlot(featureSlot, textureSlot, (UnityTextureID)params->textureID, params->width, params->height);
        break;
    }
    }
}