using System.Runtime.InteropServices;

namespace UnityEngine.Rendering.PostProcessing
{
    internal class SGSR2Upscaler_3PassCS: SGSR2Upscaler
    {
        protected override string VariantName => "SGSR2 3-Pass CS";
        
        protected override void DoRender(CommandBuffer cmd, PostProcessRenderContext context, Upscaling config)
        {
            Convert(cmd, context, config);
            Activate(cmd, context);
            Upscale(cmd, context);
        }

        private void Convert(CommandBuffer cmd, PostProcessRenderContext context, Upscaling config)
        {
            var shader = context.resources.computeShaders.sgsr2Upscaler.threePassCompute.convert;
            int kernelIndex = shader.FindKernel("CS");
            
            cmd.SetComputeConstantBufferParam(shader, "cbSGSR2", _paramsBuffer, 0, Marshal.SizeOf<SGSR2.Params>());
            cmd.SetComputeTextureParam(shader, kernelIndex, "InputOpaqueColor", config.ColorOpaqueOnly);
            cmd.SetComputeTextureParam(shader, kernelIndex, "InputColor", context.source);
            cmd.SetComputeTextureParam(shader, kernelIndex, "InputDepth", BuiltinRenderTextureType.CameraTarget, 0, RenderTextureSubElement.Depth);
            cmd.SetComputeTextureParam(shader, kernelIndex, "InputVelocity", BuiltinRenderTextureType.MotionVectors);
            cmd.SetComputeTextureParam(shader, kernelIndex, "MotionDepthAlphaBuffer", _motionDepthAlpha);
            cmd.SetComputeTextureParam(shader, kernelIndex, "YCoCgColor", _colorLuma);
            
            const int threadGroupWorkRegionDim = 8;
            int dispatchSrcX = (_paramsBuffer.Value.renderSize.x + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            int dispatchSrcY = (_paramsBuffer.Value.renderSize.y + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            cmd.DispatchCompute(shader, kernelIndex, dispatchSrcX, dispatchSrcY, 1);
        }

        private void Activate(CommandBuffer cmd, PostProcessRenderContext context)
        {
            var shader = context.resources.computeShaders.sgsr2Upscaler.threePassCompute.activate;
            int kernelIndex = shader.FindKernel("CS");
            uint frameIndex = _frameCount % 2;
            
            cmd.SetComputeConstantBufferParam(shader, "cbSGSR2", _paramsBuffer, 0, Marshal.SizeOf<SGSR2.Params>());
            cmd.SetComputeTextureParam(shader, kernelIndex, "PrevLumaHistory", _lumaHistory[frameIndex ^ 1]);
            cmd.SetComputeTextureParam(shader, kernelIndex, "MotionDepthAlphaBuffer", _motionDepthAlpha);
            cmd.SetComputeTextureParam(shader, kernelIndex, "YCoCgColor", _colorLuma);
            cmd.SetComputeTextureParam(shader, kernelIndex, "MotionDepthClipAlphaBuffer", _motionDepthClipAlpha);
            cmd.SetComputeTextureParam(shader, kernelIndex, "LumaHistory", _lumaHistory[frameIndex]);
            
            const int threadGroupWorkRegionDim = 8;
            int dispatchSrcX = (_paramsBuffer.Value.renderSize.x + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            int dispatchSrcY = (_paramsBuffer.Value.renderSize.y + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            cmd.DispatchCompute(shader, kernelIndex, dispatchSrcX, dispatchSrcY, 1);
        }

        private void Upscale(CommandBuffer cmd, PostProcessRenderContext context)
        {
            var shader = context.resources.computeShaders.sgsr2Upscaler.threePassCompute.upscale;
            int kernelIndex = shader.FindKernel("CS");
            uint frameIndex = _frameCount % 2;
            
            cmd.SetComputeConstantBufferParam(shader, "cbSGSR2", _paramsBuffer, 0, Marshal.SizeOf<SGSR2.Params>());
            cmd.SetComputeTextureParam(shader, kernelIndex, "PrevHistoryOutput", _upscaleHistory[frameIndex ^ 1]);
            cmd.SetComputeTextureParam(shader, kernelIndex, "MotionDepthClipAlphaBuffer", _motionDepthClipAlpha);
            cmd.SetComputeTextureParam(shader, kernelIndex, "YCoCgColor", _colorLuma);
            cmd.SetComputeTextureParam(shader, kernelIndex, "SceneColorOutput", context.destination);
            cmd.SetComputeTextureParam(shader, kernelIndex, "HistoryOutput", _upscaleHistory[frameIndex]);
            
            const int threadGroupWorkRegionDim = 8;
            int dispatchDstX = (_paramsBuffer.Value.displaySize.x + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            int dispatchDstY = (_paramsBuffer.Value.displaySize.y + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            cmd.DispatchCompute(shader, kernelIndex, dispatchDstX, dispatchDstY, 1);
        }
    }
}