#include "common.h"

#include <libsysmodule.h>
#include <agc.h>
#include <agc/core/sync.h>
#include <agc/gnmp/gnmp.h>
#include <psml_mfsr.h>

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

#include "UnityPluginAPI/IUnityGraphicsPS5.h"
#include "UnityPluginAPI/IUnityGraphicsAgcPS5.h"

using namespace sce::Agc;
using namespace sce::Psml;

static IUnityInterfaces* s_UnityInterfaces = nullptr;
static IUnityLog* s_Log = nullptr;
static IUnityGraphics* s_Graphics = nullptr;
static IUnityGraphicsPS5* s_GraphicsPS5 = nullptr;          // Old Gnmp-based graphics API
static IUnityGraphicsAgcPS5* s_GraphicsAgcPS5 = nullptr;    // New Agc-based graphics API (NGGC)
static UnityGfxRenderer s_RendererType = kUnityGfxRendererNull;

MfsrPhysicalMemoryBlock* s_mfsrSharedResourcesMemoryBlocks = nullptr;
MfsrSharedResources s_mfsrSharedResources = nullptr;

// Just use a single context for now to keep things simple
// "Multiple contexts (up to 4) can be generated to perform multiple MFSR processes within an application."
static bool s_mfsrInitialized = false;
static Core::BasicContext s_AgcContext;
static MfsrPhysicalMemoryBlock* s_mfsrMemoryBlocks = nullptr;
static MfsrContext s_mfsrContext = nullptr;
static Core::Texture s_outputColorTexture;

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

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

// Freely defined function to pass a callback to plugin-specific scripts
extern "C" UnityRenderingEventAndData UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API GetRenderEventAndDataFunc()
{
    return OnRenderEventAndData;
}

static void UNITY_INTERFACE_API OnGraphicsDeviceEvent(UnityGfxDeviceEventType eventType)
{
    switch (eventType)
    {
        case kUnityGfxDeviceEventInitialize:
        {
            s_RendererType = s_Graphics->GetRenderer();
            if (s_RendererType != kUnityGfxRendererPS5 && s_RendererType != kUnityGfxRendererPS5NGGC)
            {
                return;
            }

            s_GraphicsPS5 = s_UnityInterfaces->Get<IUnityGraphicsPS5>();
            s_GraphicsAgcPS5 = s_UnityInterfaces->Get<IUnityGraphicsAgcPS5>();
            if (s_GraphicsPS5 == nullptr && s_GraphicsAgcPS5 == nullptr)
            {
                UNITY_LOG_ERROR(s_Log, "Could not obtain PS5 Graphics interface!");
                return;
            }

            // Yeah just don't do the stuff below. It makes Unity act weirdly.
            //UnityD3D12PluginEventConfig config;
            //config.graphicsQueueAccess = kUnityD3D12GraphicsQueueAccess_DontCare;
            //config.flags = kUnityD3D12EventConfigFlag_ModifiesCommandBuffersState;
            //config.ensureActiveRenderTextureIsBound = true; // TODO: not entirely sure if this is necessary
            //s_GraphicsD3D12->ConfigureEvent(1, &config);

            //config.graphicsQueueAccess = kUnityD3D12GraphicsQueueAccess_DontCare;
            //config.flags = 0;
            //config.ensureActiveRenderTextureIsBound = false;
            //s_GraphicsD3D12->ConfigureEvent(2, &config);

            break;
        }
        case kUnityGfxDeviceEventShutdown:
        {
            s_GraphicsPS5 = nullptr;
            s_GraphicsAgcPS5 = nullptr;
            s_RendererType = kUnityGfxRendererNull;
            break;
        }
        case kUnityGfxDeviceEventBeforeReset:
        {
            //TODO: user Direct3D 9 code
            break;
        }
        case kUnityGfxDeviceEventAfterReset:
        {
            //TODO: user Direct3D 9 code
            break;
        }
    };
}

extern "C" bool UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API InitPssr()
{
    // Check if we're running on PS5 Pro
    if (!sceKernelIsTrinityMode())
    {
        UNITY_LOG_ERROR(s_Log, "Kernel is not running in Trinity mode, PSML is not supported!");
        return false;
    }

    // Load PSML module
    if (sceSysmoduleIsLoaded(SCE_SYSMODULE_PSML) != SCE_SYSMODULE_LOADED)
    {
        int res = sceSysmoduleLoadModule(SCE_SYSMODULE_PSML);
        if (res < SCE_OK)
        {
            std::stringstream msg;
            msg << "Failed to load PSML sysmodule, error code = " << res;
            UNITY_LOG_ERROR(s_Log, msg.str().c_str());
            return false;
        }
    }

    // Initialize MFSR
    int res = initMfsr();
    if (res != SCE_OK && res != SCE_PSML_MFSR_ERROR_ALREADY_INITIALIZED)
    {
        std::stringstream msg;
        msg << "Failed to initialize MFSR, error code = " << res;
        UNITY_LOG_ERROR(s_Log, msg.str().c_str());
        return false;
    }

    // TODO: do some idempotence here, make sure we don't create shared resources twice

    // Set up MFSR shared resources
    {
        // Get required physical memory blocks sizeAlign for initializing MFSR shared resources
        MfsrSharedResourcesInitParameters mfsrSharedResourcesInitParameters;
        mfsrSharedResourcesInitParameters.init();
        mfsrSharedResourcesInitParameters.m_profile = MfsrProfile::kQuality;

        MfsrSharedResourcesInitRequirement mfsrSharedResourcesInitRequirement;
        getMfsrSharedResourcesInitRequirement(&mfsrSharedResourcesInitRequirement, &mfsrSharedResourcesInitParameters);

        // Allocate physical memory blocks for MFSR shared resource and setup shared resources init param
        s_mfsrSharedResourcesMemoryBlocks = new MfsrPhysicalMemoryBlock[mfsrSharedResourcesInitRequirement.m_physicalMemoryBlockCount];

        size_t blockSize = mfsrSharedResourcesInitRequirement.m_physicalMemoryBlockSize;
        size_t blockAlign = mfsrSharedResourcesInitRequirement.m_physicalMemoryBlockAlignment;

        size_t m_mfsrSharedResourcesMemorySize = mfsrSharedResourcesInitRequirement.m_physicalMemoryBlockCount * blockSize;
        off_t m_mfsrSharedResourcesMemoryAddress;
        sceKernelAllocateMainDirectMemory(m_mfsrSharedResourcesMemorySize, blockAlign, SCE_KERNEL_MTYPE_C_SHARED, &m_mfsrSharedResourcesMemoryAddress);

        for (uint32_t i = 0; i < mfsrSharedResourcesInitRequirement.m_physicalMemoryBlockCount; i++)
        {
            s_mfsrSharedResourcesMemoryBlocks[i].m_size = blockSize;
            s_mfsrSharedResourcesMemoryBlocks[i].m_offset = m_mfsrSharedResourcesMemoryAddress + blockSize * i;
        }

        mfsrSharedResourcesInitParameters.m_physicalMemoryBlocks = s_mfsrSharedResourcesMemoryBlocks;
        mfsrSharedResourcesInitParameters.m_physicalMemoryBlockCount = mfsrSharedResourcesInitRequirement.m_physicalMemoryBlockCount;

        // Create MFSR shared resources with the memory
        createMfsrSharedResources(&s_mfsrSharedResources, &mfsrSharedResourcesInitParameters);
    }

    s_mfsrInitialized = true;

    // Set up command buffer
    if (s_RendererType == kUnityGfxRendererPS5NGGC && s_GraphicsAgcPS5 != nullptr)
    {
        const size_t agcContextSize = 128 * 1024;
        s_AgcContext.m_dcb.init(s_GraphicsAgcPS5->AllocateGPUMemory(agcContextSize, Alignment::kCommandBuffer), agcContextSize);
        s_AgcContext.m_bdr.init(&s_AgcContext.m_dcb, &s_AgcContext.m_dcb);
        s_AgcContext.m_sb.init(256, &s_AgcContext.m_dcb, &s_AgcContext.m_dcb);
    }

    UNITY_LOG(s_Log, "PSSR plugin initialized");
    return true;
}

extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API ReleasePssr()
{
    s_mfsrInitialized = false;

    if (s_GraphicsAgcPS5 != nullptr && s_AgcContext.m_dcb.m_bottom != nullptr)
    {
        s_GraphicsAgcPS5->ReleaseGPUMemory(s_AgcContext.m_dcb.m_bottom);
        s_AgcContext.m_dcb.clear();
    }

    if (s_mfsrSharedResources != nullptr)
    {
        releaseMfsrSharedResources(s_mfsrSharedResources);
        s_mfsrSharedResources = nullptr;
    }

    if (s_mfsrSharedResourcesMemoryBlocks != nullptr)
    {
        delete[] s_mfsrSharedResourcesMemoryBlocks;
        s_mfsrSharedResourcesMemoryBlocks = nullptr;
    }
}

typedef struct pssr_init_params_s
{
    uint32_t displayWidth;
    uint32_t displayHeight;
    uint32_t maxRenderWidth;
    uint32_t maxRenderHeight;
} pssr_init_params_t;

typedef struct pssr_dispatch_params_s
{
    sce::Agc::Core::Texture* color;
    sce::Agc::Core::Texture* depth;
    sce::Agc::Core::Texture* motionVectors;
    sce::Agc::Core::Texture* exposure;
    sce::Agc::Core::Texture* reactiveMask;
    sce::Agc::Core::Texture* outputColor;

    uint32_t renderWidth;
    uint32_t renderHeight;

    Vector2f jitter;
    Vector2f motionVectorScale;

    Matrix4f camProjection;
    Vector3f camForward;
    Vector3f camUp;
    Vector3f camRight;
    Vector3d camPosition;
    float pad;

    float camNear;
    float camFar;
    float preExposure;
    uint32_t resetHistory;

} pssr_dispatch_params_t;

extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API CreatePssrContext(const pssr_init_params_t* params, Core::Texture** outputColorTexture)
{
    MfsrContextInitParameters initParams = {};
    initParams.init();
    initParams.m_outputColorWidth = params->displayWidth;
    initParams.m_outputColorHeight = params->displayHeight;
    initParams.m_sharedResources = s_mfsrSharedResources;

    // Keep a copy of the previous depth internally
    MfsrKeepCopyTextureSpec depthCopySpec = {};
    depthCopySpec.m_width = params->maxRenderWidth;
    depthCopySpec.m_height = params->maxRenderHeight;
    depthCopySpec.m_bytesPerPixel = 4;  // R32_SFloat
    initParams.m_maxSizeDepthForKeepCopy = &depthCopySpec;

    // Keep a copy of the previous motion vectors internally
    MfsrKeepCopyTextureSpec mvCopySpec = {};
    mvCopySpec.m_width = params->maxRenderWidth;
    mvCopySpec.m_height = params->maxRenderHeight;
    mvCopySpec.m_bytesPerPixel = 4;  // R16G16_SFloat
    initParams.m_maxSizeMotionVectorsForKeepCopy = &mvCopySpec;

    // Song and dance to set up memory used by this MFSR context
    {
        MfsrContextInitRequirement mfsrInitRequirement = {};
        getMfsrContextInitRequirement(&mfsrInitRequirement, &initParams);

        s_mfsrMemoryBlocks = new MfsrPhysicalMemoryBlock[mfsrInitRequirement.m_physicalMemoryBlockCount];

        size_t size = mfsrInitRequirement.m_physicalMemoryBlockSize;
        size_t align = mfsrInitRequirement.m_physicalMemoryBlockAlignment;

        size_t m_mfsrMemorySize = mfsrInitRequirement.m_physicalMemoryBlockCount * size;
        off_t m_mfsrMemoryPhysicalAddress;
        sceKernelAllocateMainDirectMemory(m_mfsrMemorySize, align, SCE_KERNEL_MTYPE_C_SHARED, &m_mfsrMemoryPhysicalAddress);

        for (uint32_t i = 0; i < mfsrInitRequirement.m_physicalMemoryBlockCount; i++)
        {
            s_mfsrMemoryBlocks[i].m_size = size;
            s_mfsrMemoryBlocks[i].m_offset = m_mfsrMemoryPhysicalAddress + mfsrInitRequirement.m_physicalMemoryBlockSize * i;
        }

        initParams.m_sharedResources = s_mfsrSharedResources;
        initParams.m_physicalMemoryBlocks = s_mfsrMemoryBlocks;
        initParams.m_physicalMemoryBlockCount = mfsrInitRequirement.m_physicalMemoryBlockCount;
    }

    // Set up output color texture with the correct specifications (k11_11_10Float with tile mode kStandard256B)
    {
        Core::TextureSpec texSpec;
        texSpec.init();
        texSpec.setAllowNullptr(true);
        texSpec.m_width = params->displayWidth;
        texSpec.m_height = params->displayHeight;
        texSpec.m_format = { Core::TypedFormat::k11_11_10Float, Core::Swizzle::kRGB1_R3S34 };
        texSpec.setTileMode(Core::Texture::TileMode::kStandard256B);

        SizeAlign sizeAlign = Core::getSize(&texSpec);
        if (s_GraphicsPS5 != nullptr)
        {
            texSpec.m_dataAddress = s_GraphicsPS5->AllocateGPUMemory(sizeAlign.m_size, MfsrAlignment::kOutputTexture);
        }
        else if (s_GraphicsAgcPS5 != nullptr)
        {
            texSpec.m_dataAddress = s_GraphicsAgcPS5->AllocateGPUMemory(sizeAlign.m_size, MfsrAlignment::kOutputTexture);
        }

        int res = Core::initialize(&s_outputColorTexture, &texSpec);
        if (res != SCE_OK)
        {
            std::stringstream msg;
            msg << "Failed to create MFSR output color texture, error code = " << res;
            UNITY_LOG_ERROR(s_Log, msg.str().c_str());
            return res;
        }

        if (outputColorTexture != nullptr)
        {
            *outputColorTexture = &s_outputColorTexture;
        }
    }

    // Finally, create the actual MFSR context
    int res = createMfsrContext(&s_mfsrContext, &initParams);
    if (res != SCE_OK)
    {
        std::stringstream msg;
        msg << "Failed to create MFSR context, error code = " << res;
        UNITY_LOG_ERROR(s_Log, msg.str().c_str());
        return res;
    }

    UNITY_LOG(s_Log, "Created PSSR context!");
    return SCE_OK;
}

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

    if (s_RendererType == kUnityGfxRendererPS5 && s_GraphicsPS5 == nullptr)
        return;

    if (s_RendererType == kUnityGfxRendererPS5NGGC && s_GraphicsAgcPS5 == nullptr)
        return;

    // User rendering code
    switch (eventID)
    {
        case 0: // Initialize PSSR
        {
            auto* params = (pssr_init_params_t*)data;
            CreatePssrContext(params, nullptr);
            break;
        }
        case 1: // Execute PSSR
        {
            if (s_mfsrContext == nullptr)
                break;

            // TODO: how do get command buffer from PS5 graphics interface? :thinking:
            // I guess we're hoping that s_GraphicsAgcPS5->SubmitGraphics is smart enough to add the submit pointer to the currently recording command list (which seems to be the case, phew!)
            auto* params = (pssr_dispatch_params_t*)data;

            sce::Agc::DrawCommandBuffer* cmd = nullptr;
            if (s_RendererType == kUnityGfxRendererPS5)
            {
                auto* gfxc = (sce::Agc::Gnmp::LightweightGfxContext*)s_GraphicsPS5->GetGfxContext();
                cmd = &gfxc->m_agcCtx.m_dcb;
            }
            else if (s_RendererType == kUnityGfxRendererPS5NGGC)
            {
                Core::BasicContext& ctx = s_AgcContext.reset();
                cmd = &ctx.m_dcb;
            }

            cmd->pushMarker("Dispatch MFSR");

            DispatchMfsrParameters dispatchParams = {};
            dispatchParams.init();
            dispatchParams.m_outputColor = params->outputColor;
            dispatchParams.m_color = params->color;
            dispatchParams.m_motionVectors = params->motionVectors;
            dispatchParams.m_prevDepth = nullptr;   // MFSR context keeps a copy of the previous depth internally
            dispatchParams.m_depth = params->depth;
            dispatchParams.m_exposure = params->exposure;
            dispatchParams.m_reactiveMask = params->reactiveMask;
            dispatchParams.m_prevMotionVectors = nullptr;   // MFSR context keeps a copy of the previous motion vectors internally
            dispatchParams.m_preExposure = params->preExposure;
            dispatchParams.m_projectionNoJitter = params->camProjection;
            dispatchParams.m_camForward = params->camForward;
            dispatchParams.m_camUp = params->camUp;
            dispatchParams.m_camRight = params->camRight;
            dispatchParams.m_camPosD = params->camPosition;
            dispatchParams.m_nearZ = params->camNear;
            dispatchParams.m_farZ = params->camFar;
            dispatchParams.m_renderWidth = params->renderWidth;
            dispatchParams.m_renderHeight = params->renderHeight;
            dispatchParams.m_jitter = params->jitter;
            dispatchParams.m_motionVectorScale = params->motionVectorScale;
            dispatchParams.m_reset = params->resetHistory != 0;
            dispatchParams.m_flags = MfsrOptionFlags::kReverseDepth;    // TODO: make flags available as plugin input

            dispatchMfsr(s_mfsrContext, cmd, &dispatchParams);

            cmd->popMarker();

            if (s_RendererType == kUnityGfxRendererPS5NGGC)
            {
                s_GraphicsAgcPS5->SubmitGraphics(cmd->getSubmitPointer(), cmd->getSubmitSize());
            }
            break;
        }
        case 2: // Destroy PSSR context
        {
            // TODO: release color input texture

            // TODO: this still crashes often due to race conditions between dispatch above and release here
            // The main cause will be that the MFSR context is released before the already dispatched commands are executed on the GPU
            // Just adding mutexes everywhere is not going to fix that
            // A cheap but slightly dirty solution would be to run this code only at EndOfFrame from Unity
            // 
            // Another issue is that we reuse the same static s_mfsrContext pointer when recreating the MFSR context, which can lead to situations where dispatch uses the wrong context object
            // Multi-buffering those pointers will help with that
            if (s_mfsrContext != nullptr)
            {
                releaseMfsrContext(s_mfsrContext);
                s_mfsrContext = nullptr;
            }

            if (s_mfsrMemoryBlocks != nullptr)
            {
                delete[] s_mfsrMemoryBlocks;
                s_mfsrMemoryBlocks = nullptr;
            }

            UNITY_LOG(s_Log, "Destroyed PSSR context");
            break;
        }
    }
}