using System.Runtime.InteropServices;
using FidelityFX;
using UnityEngine.Experimental.Rendering;

namespace UnityEngine.Rendering.PostProcessing
{
    internal class SGSR2Upscaler: Upscaler
    {
        public static bool IsSupported => SystemInfo.supportsComputeShaders;
        
        private RenderTexture _colorLuma;
        private RenderTexture _motionDepthAlpha;
        private RenderTexture _motionDepthClipAlpha;
        private readonly RenderTexture[] _lumaHistory = new RenderTexture[2];
        private readonly RenderTexture[] _upscaleHistory = new RenderTexture[2];

        private readonly ConstantsBuffer<SGSR2.Params> _paramsBuffer = new();

        private uint _frameCount = 0;
        
        public override void CreateContext(PostProcessRenderContext context, Upscaling config)
        {
            CreateRenderTexture(ref _colorLuma, "ColorLuma", config.MaxRenderSize, GraphicsFormat.R32_UInt, true);
            CreateRenderTexture(ref _motionDepthAlpha, "MotionDepthAlpha", config.MaxRenderSize, GraphicsFormat.R16G16B16A16_SFloat, true);
            CreateRenderTexture(ref _motionDepthClipAlpha, "MotionDepthClipAlpha", config.MaxRenderSize, GraphicsFormat.R16G16B16A16_SFloat, true);
            CreateRenderTextureArray(_lumaHistory, "LumaHistory", config.MaxRenderSize, GraphicsFormat.R32_UInt, true);
            CreateRenderTextureArray(_upscaleHistory, "History", config.UpscaleSize, context.sourceFormat, true);
            
            _paramsBuffer.Init();
        }

        public override void DestroyContext()
        {
            base.DestroyContext();
            
            _paramsBuffer.Destroy();
            
            DestroyRenderTextureArray(_upscaleHistory);
            DestroyRenderTextureArray(_lumaHistory);
            DestroyRenderTexture(ref _motionDepthClipAlpha);
            DestroyRenderTexture(ref _motionDepthAlpha);
            DestroyRenderTexture(ref _colorLuma);
        }

        public override void Render(PostProcessRenderContext context, Upscaling config)
        {
            var cmd = context.command;
            cmd.BeginSample("SGSR2");
            
            ref var parms = ref _paramsBuffer.Value;
            parms.renderSize = config.GetScaledRenderSize(context.camera);
            parms.displaySize = config.UpscaleSize;
            parms.renderSizeRcp = new Vector2(1.0f / parms.renderSize.x, 1.0f / parms.renderSize.y);
            parms.displaySizeRcp = new Vector2(1.0f / parms.displaySize.x, 1.0f / parms.displaySize.y);
            parms.jitterOffset = config.JitterOffset;
            parms.clipToPrevClip = Matrix4x4.identity; // TODO: clipToPrevClip => (previous_view_proj * inv_vp)
            parms.preExposure = config.preExposure;
            parms.cameraFovAngleHor = Mathf.Tan(context.camera.fieldOfView * Mathf.Deg2Rad * 0.5f) * (float)parms.renderSize.x / parms.renderSize.y;
            parms.cameraNear = context.camera.nearClipPlane;
            parms.minLerpContribution = 0f;
            parms.bSameCamera = 0u;
            parms.reset = config.Reset ? 1u : 0u;
            _paramsBuffer.UpdateBufferData(cmd);

            if (_frameCount == 0 || config.Reset)
            {
                cmd.SetRenderTarget(_lumaHistory[0]);
                cmd.ClearRenderTarget(false, true, Color.clear);
                cmd.SetRenderTarget(_lumaHistory[1]);
                cmd.ClearRenderTarget(false, true, Color.clear);
                cmd.SetRenderTarget(_upscaleHistory[0]);
                cmd.ClearRenderTarget(false, true, Color.clear);
                cmd.SetRenderTarget(_upscaleHistory[1]);
                cmd.ClearRenderTarget(false, true, Color.clear);
            }

            Convert(cmd, context, config);
            Activate(cmd, context);
            Upscale(cmd, context);
            
            cmd.EndSample("SGSR2");
            _frameCount++;
        }

        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, "Params", _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, "Params", _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, "Params", _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);
        }
    }
}