/*
 * Copyright (c) 2026 Mykhailo Mamedov. All rights reserved.
 *
 * RESEARCH PREVIEW / REFERENCE ONLY:
 * This source code is provided solely for the purpose of reviewing
 * the author's research methods and implementation.
 *
 * NO LICENSE GRANTED:
 * This code is NOT for distribution, modification, or use in any
 * project (commercial or otherwise). Unauthorized copying or
 * use of this code is strictly prohibited.
 *
 * For inquiries regarding use or licensing, contact: ua.modin@gmail.com
 * 
 * Simple UE Shader dispatch for testing.
 * 
 */

#include "testPlugMi.h"
#include "Interfaces/IPluginManager.h"
#include "ShaderParameterStruct.h"
#include "GlobalShader.h"
#include "HAL/IConsoleManager.h"
#include "RenderGraphBuilder.h"
#include "RenderGraphUtils.h"
#include "SceneViewExtension.h"

#define LOCTEXT_NAMESPACE "FtestPlugMiModule"

TRefCountPtr<IPooledRenderTarget> MyInternalTexturePtr;

class FMyTestCS : public FGlobalShader
{
public:
    DECLARE_GLOBAL_SHADER(FMyTestCS);
    SHADER_USE_PARAMETER_STRUCT(FMyTestCS, FGlobalShader);

    BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
        SHADER_PARAMETER_RDG_TEXTURE_UAV(RWTexture2D<float4>, OutputTexture)
    END_SHADER_PARAMETER_STRUCT()

    static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters) {
        return IsFeatureLevelSupported(Parameters.Platform, ERHIFeatureLevel::SM5);
    }

    /// <summary>
    /// Shader dispatch proc.
    /// </summary>
    static void Dispatch(FRDGBuilder& GraphBuilder, FRDGTextureRef OutputTexture) {
        TShaderMapRef<FMyTestCS> ComputeShader(GetGlobalShaderMap(GMaxRHIFeatureLevel));
        if (ComputeShader.IsValid()) {
            FMyTestCS::FParameters* PassParameters = GraphBuilder.AllocParameters<FMyTestCS::FParameters>();
            PassParameters->OutputTexture = GraphBuilder.CreateUAV(OutputTexture);
            FComputeShaderUtils::AddPass(
                GraphBuilder,
                RDG_EVENT_NAME("testPass"),
                ComputeShader,
                PassParameters,
                FIntVector(32, 32, 1)
            );
        }
    }
};

IMPLEMENT_GLOBAL_SHADER(FMyTestCS, "/Plugin/testPlugMi/Private/textureWrite.usf", "MainCS", SF_Compute);

class FMyTestViewExtension : public FSceneViewExtensionBase
{
public:
    FMyTestViewExtension(const FAutoRegister& AutoRegister) : FSceneViewExtensionBase(AutoRegister) {}

    /// <summary>
    /// Create temporary texture and dispatch shader, boilerplates.
    /// </summary>
    virtual void PreRenderView_RenderThread(FRDGBuilder& GraphBuilder, FSceneView& InView) override {

        FRDGTextureDesc Desc = FRDGTextureDesc::Create2D(
            FIntPoint(1024, 1024),
            PF_FloatRGBA,
            FClearValueBinding::Black,
            TexCreate_ShaderResource | TexCreate_UAV
        );

        FRDGTextureRef textureTest = GraphBuilder.CreateTexture(Desc, TEXT("textureTestShaderOutput"));
        FMyTestCS::Dispatch(GraphBuilder, textureTest);
        GraphBuilder.QueueTextureExtraction(textureTest, &MyInternalTexturePtr);
    }

    virtual void SetupViewFamily(FSceneViewFamily& InViewFamily) override {}
    virtual void SetupView(FSceneViewFamily& InViewFamily, FSceneView& InView) override {}
    virtual void BeginRenderViewFamily(FSceneViewFamily& InViewFamily) override {}
};


TSharedPtr<FMyTestViewExtension> TestViewExtension;


void FtestPlugMiModule::StartupModule() {

    // Register shader earlier (PostConfigInit).
    FString ShaderDir = FPaths::Combine(IPluginManager::Get().FindPlugin(TEXT("testPlugMi"))->GetBaseDir(), TEXT("Shaders"));
    AddShaderSourceDirectoryMapping(TEXT("/Plugin/testPlugMi"), ShaderDir);

    // Adds a new view extension on post engine init.
    FCoreDelegates::OnPostEngineInit.AddLambda([]() {
        TestViewExtension = TSharedPtr<FMyTestViewExtension>(FSceneViewExtensions::NewExtension<FMyTestViewExtension>());
        UE_LOG(LogTemp, Warning, TEXT("Extension Registered Safely via Delegate!"));
    });
}

void FtestPlugMiModule::ShutdownModule() {
    if (TestViewExtension.IsValid()) {
        TestViewExtension.Reset();
    }
}

#undef LOCTEXT_NAMESPACE
	
IMPLEMENT_MODULE(FtestPlugMiModule, testPlugMi)