#include <libsysmodule.h>
#include <agc.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;

static bool s_mfsrInitialized = false;
static size_t s_mfsrSharedResourcesMemorySize = 0;
static off_t s_mfsrSharedResourcesMemoryAddress = 0;
static MfsrPhysicalMemoryBlock* s_mfsrSharedResourcesMemoryBlocks = nullptr;
static MfsrSharedResources s_mfsrSharedResources = nullptr;

struct PssrContext
{
    Core::BasicContext agcContext;

    size_t mfsrMemorySize = 0;
    off_t mfsrMemoryPhysicalAddress = 0;
    MfsrPhysicalMemoryBlock* mfsrMemoryBlocks = nullptr;

    MfsrContext mfsrContext = nullptr;
    Core::Texture outputColorTexture;
};

// "Multiple contexts (up to 4) can be generated to perform multiple MFSR processes within an application."
static const int MaxNumContexts = 4;
static PssrContext s_contexts[MaxNumContexts];

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

            break;
        }
        case kUnityGfxDeviceEventShutdown:
        {
            s_GraphicsPS5 = nullptr;
            s_GraphicsAgcPS5 = nullptr;
            s_RendererType = kUnityGfxRendererNull;
            break;
        }
        case kUnityGfxDeviceEventBeforeReset:
        {
            break;
        }
        case kUnityGfxDeviceEventAfterReset:
        {
            break;
        }
    };
}

extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API PSSR_Init()
{
    // 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 -1;
    }

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

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

    if (s_mfsrInitialized)
    {
        return SCE_OK;
    }

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

        s_mfsrSharedResourcesMemorySize = mfsrSharedResourcesInitRequirement.m_physicalMemoryBlockCount * blockSize;
        sceKernelAllocateMainDirectMemory(s_mfsrSharedResourcesMemorySize, blockAlign, SCE_KERNEL_MTYPE_C_SHARED, &s_mfsrSharedResourcesMemoryAddress);

        for (uint32_t i = 0; i < mfsrSharedResourcesInitRequirement.m_physicalMemoryBlockCount; i++)
        {
            s_mfsrSharedResourcesMemoryBlocks[i].m_size = blockSize;
            s_mfsrSharedResourcesMemoryBlocks[i].m_offset = s_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;

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

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

    // Destroy the MFSR shared resources
    if (s_mfsrSharedResources != nullptr)
    {
        releaseMfsrSharedResources(s_mfsrSharedResources);
        s_mfsrSharedResources = nullptr;
    }

    // Free the memory used for MFSR shared resource data
    if (s_mfsrSharedResourcesMemoryAddress != 0)
    {
        sceKernelReleaseDirectMemory(s_mfsrSharedResourcesMemoryAddress, s_mfsrSharedResourcesMemorySize);
        s_mfsrSharedResourcesMemoryAddress = 0;
        s_mfsrSharedResourcesMemorySize = 0;
    }

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

typedef struct pssr_init_params_s
{
    uint32_t contextIndex;

    uint32_t displayWidth;
    uint32_t displayHeight;
    uint32_t maxRenderWidth;
    uint32_t maxRenderHeight;

    uint32_t autoKeepCopies;
} pssr_init_params_t;

typedef struct pssr_dispatch_params_s
{
    uint32_t contextIndex;

    sce::Agc::Core::Texture* color;
    sce::Agc::Core::Texture* depth;
    sce::Agc::Core::Texture* prevDepth;
    sce::Agc::Core::Texture* motionVectors;
    sce::Agc::Core::Texture* prevMotionVectors;
    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 camProjectionNoJitter;
    Vector3f camForward;
    Vector3f camUp;
    Vector3f camRight;
    Vector3d camPosition;
    float camNear;
    float camFar;
    float preExposure;
    uint32_t resetHistory;
    MfsrOptionFlags flags;

} pssr_dispatch_params_t;

extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API PSSR_CreateContext(const pssr_init_params_t* params, Core::Texture** outputColorTexture)
{
    if (params->contextIndex < 0 || params->contextIndex >= MaxNumContexts)
    {
        std::stringstream msg;
        msg << "Invalid PSSR context index: " << params->contextIndex;
        UNITY_LOG_ERROR(s_Log, msg.str().c_str());
        return -1;
    }

    PssrContext& context = s_contexts[params->contextIndex];

    MfsrContextInitParameters initParams = {};
    initParams.init();
    initParams.m_outputColorWidth = params->displayWidth;
    initParams.m_outputColorHeight = params->displayHeight;
    initParams.m_sharedResources = s_mfsrSharedResources;

    MfsrKeepCopyTextureSpec depthCopySpec = {}, mvCopySpec = {};
    if (params->autoKeepCopies)
    {
        // Keep a copy of the previous depth internally
        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
        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);

        context.mfsrMemoryBlocks = new MfsrPhysicalMemoryBlock[mfsrInitRequirement.m_physicalMemoryBlockCount];

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

        context.mfsrMemorySize = mfsrInitRequirement.m_physicalMemoryBlockCount * size;
        sceKernelAllocateMainDirectMemory(context.mfsrMemorySize, align, SCE_KERNEL_MTYPE_C_SHARED, &context.mfsrMemoryPhysicalAddress);

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

        initParams.m_sharedResources = s_mfsrSharedResources;
        initParams.m_physicalMemoryBlocks = context.mfsrMemoryBlocks;
        initParams.m_physicalMemoryBlockCount = mfsrInitRequirement.m_physicalMemoryBlockCount;
    }

    // Set up output color texture with the correct specifications (k11_11_10Float with tile mode kStandard256B)
    if (outputColorTexture != nullptr)
    {
        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(&context.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;
        }

        *outputColorTexture = &context.outputColorTexture;
    }

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

    // Finally, create the actual MFSR context
    int res = createMfsrContext(&context.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;
}

extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API PSSR_DestroyContext(uint32_t contextIndex)
{
    if (contextIndex < 0 || contextIndex >= MaxNumContexts)
    {
        std::stringstream msg;
        msg << "Invalid PSSR context index: " << contextIndex;
        UNITY_LOG_ERROR(s_Log, msg.str().c_str());
        return;
    }

    PssrContext& context = s_contexts[contextIndex];

    // Destroy the MFSR context
    if (context.mfsrContext != nullptr)
    {
        releaseMfsrContext(context.mfsrContext);
        context.mfsrContext = nullptr;
    }

    // Destroy the NGGC command buffer
    if (s_GraphicsAgcPS5 != nullptr && context.agcContext.m_dcb.m_bottom != nullptr)
    {
        s_GraphicsAgcPS5->ReleaseGPUMemory(context.agcContext.m_dcb.m_bottom);
        context.agcContext.m_dcb.clear();
    }

    // Destroy the output color texture
    if (context.outputColorTexture.getDataAddress())
    {
        if (s_GraphicsPS5 != nullptr)
        {
            s_GraphicsPS5->ReleaseGPUMemory(context.outputColorTexture.getDataAddress());
        }
        else if (s_GraphicsAgcPS5 != nullptr)
        {
            s_GraphicsAgcPS5->ReleaseGPUMemory(context.outputColorTexture.getDataAddress());
        }

        context.outputColorTexture.init();
    }

    // Free the memory used for MFSR context data
    if (context.mfsrMemoryPhysicalAddress != 0)
    {
        sceKernelReleaseDirectMemory(context.mfsrMemoryPhysicalAddress, context.mfsrMemorySize);
        context.mfsrMemoryPhysicalAddress = 0;
        context.mfsrMemorySize = 0;
    }

    if (context.mfsrMemoryBlocks != nullptr)
    {
        delete[] context.mfsrMemoryBlocks;
        context.mfsrMemoryBlocks = nullptr;
    }

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

extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API PSSR_Dispatch(const pssr_dispatch_params_t* params)
{
    if (params->contextIndex < 0 || params->contextIndex >= MaxNumContexts)
    {
        std::stringstream msg;
        msg << "Invalid PSSR context index: " << params->contextIndex;
        UNITY_LOG_ERROR(s_Log, msg.str().c_str());
        return;
    }

    PssrContext& context = s_contexts[params->contextIndex];
    if (context.mfsrContext == nullptr)
        return;

    // How we obtain a command buffer to dispatch to depends on which graphics API we're using in Unity
    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 = context.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 = params->prevDepth;
    dispatchParams.m_depth = params->depth;
    dispatchParams.m_exposure = params->exposure;
    dispatchParams.m_reactiveMask = params->reactiveMask;
    dispatchParams.m_prevMotionVectors = params->prevMotionVectors;
    dispatchParams.m_preExposure = params->preExposure;
    dispatchParams.m_projectionNoJitter = params->camProjectionNoJitter;
    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_motionVectorGamma = 1.0f;
    dispatchParams.m_reset = params->resetHistory != 0;
    dispatchParams.m_flags = params->flags;

    dispatchMfsr(context.mfsrContext, cmd, &dispatchParams);

    cmd->popMarker();

    if (s_RendererType == kUnityGfxRendererPS5NGGC)
    {
        s_GraphicsAgcPS5->SubmitGraphics(cmd->getSubmitPointer(), cmd->getSubmitSize());
    }
}

// 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: // Create PSSR context
        {
            auto* params = (pssr_init_params_t*)data;
            PSSR_CreateContext(params, nullptr);
            break;
        }
        case 1: // Execute PSSR
        {
            auto* params = (pssr_dispatch_params_t*)data;
            PSSR_Dispatch(params);
            break;
        }
        case 2: // Destroy PSSR context
        {
            uint32_t* contextIndex = (uint32_t*)data;
            PSSR_DestroyContext(*contextIndex);
            break;
        }
    }
}