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

namespace FidelityFX
{
    internal abstract class CacaoPass: IDisposable
    {
        protected readonly CacaoResources Resources;

        protected ComputeShader ComputeShader;
        protected int[] KernelIndices;
        
        protected CacaoPass(CacaoResources resources)
        {
            Resources = resources;
        }

        protected void InitComputeShaders<TKernel>(string passName, ComputeShader shader)
            where TKernel: Enum
        {
            if (shader == null)
            {
                throw new MissingReferenceException($"Shader for CACAO '{passName}' could not be loaded! Please ensure it is included in the project correctly.");
            }

            ComputeShader = shader;
            
            var kernelNames = Enum.GetNames(typeof(TKernel));
            KernelIndices = new int[kernelNames.Length];

            for (int i = 0; i < kernelNames.Length; ++i)
            {
                KernelIndices[i] = shader.FindKernel(kernelNames[i]);
            }
        }

        public void Dispose()
        {
        }
        
        protected static int DispatchSize(uint tileSize, uint totalSize)
        {
            return (int)((totalSize + tileSize - 1) / tileSize);
        }
    }

    internal class CacaoClearLoadCounterPass : CacaoPass
    {
        private enum Kernels
        {
            CS,
        }

        public CacaoClearLoadCounterPass(ComputeShader shader, CacaoResources resources)
            : base(resources)
        {
            InitComputeShaders<Kernels>("clear_load_counter", shader);
        }

        public void Execute(CommandBuffer commandBuffer, ComputeBuffer constants)
        {
            int kernelIndex = KernelIndices[(int)Kernels.CS];
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CacaoShaderIDs.CbSsaoConstantsBuffer, constants, 0, Marshal.SizeOf<Cacao.Constants>());
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavLoadCounterBuffer, Resources.LoadCounter);
            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, 1, 1, 1);
        }
    }

    internal class CacaoPrepareDepthsPass : CacaoPass
    {
        private enum Kernels
        {
            CS_Native,
            CS_NativeHalf,
            CS_NativeAndMips,
            CS_Downsampled,
            CS_DownsampledHalf,
            CS_DownsampledAndMips,
        }
        
        public CacaoPrepareDepthsPass(ComputeShader shader, CacaoResources resources)
            : base(resources)
        {
            InitComputeShaders<Kernels>("prepare_depths", shader);
        }
        
        public void Execute(CommandBuffer commandBuffer, ComputeBuffer constants, in Cacao.DispatchInfo dispatchInfo, in Cacao.BufferSizeInfo bsi, Cacao.Quality quality, bool useDownsampledSsao)
        {
            Kernels kernel;
            int dispatchWidth, dispatchHeight;

            switch (quality)
            {
                case Cacao.Quality.Lowest:
                    dispatchWidth = DispatchSize(CacaoConsts.PrepareDepthsHalfWidth, bsi.DeinterleavedDepthBufferWidth);
                    dispatchHeight = DispatchSize(CacaoConsts.PrepareDepthsHalfHeight, bsi.DeinterleavedDepthBufferHeight);
                    kernel = useDownsampledSsao ? Kernels.CS_DownsampledHalf : Kernels.CS_NativeHalf;
                    break;
                case Cacao.Quality.Low:
                    dispatchWidth = DispatchSize(CacaoConsts.PrepareDepthsWidth, bsi.DeinterleavedDepthBufferWidth);
                    dispatchHeight = DispatchSize(CacaoConsts.PrepareDepthsHeight, bsi.DeinterleavedDepthBufferHeight);
                    kernel = useDownsampledSsao ? Kernels.CS_Downsampled : Kernels.CS_Native;
                    break;
                default:
                    dispatchWidth = DispatchSize(CacaoConsts.PrepareDepthsAndMipsWidth, bsi.DeinterleavedDepthBufferWidth);
                    dispatchHeight = DispatchSize(CacaoConsts.PrepareDepthsAndMipsHeight, bsi.DeinterleavedDepthBufferHeight);
                    kernel = useDownsampledSsao ? Kernels.CS_DownsampledAndMips : Kernels.CS_NativeAndMips;
                    break;
            }

            int kernelIndex = KernelIndices[(int)kernel];
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CacaoShaderIDs.CbSsaoConstantsBuffer, constants, 0, Marshal.SizeOf<Cacao.Constants>());
            
            var depthView = dispatchInfo.DepthView;
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvDepthIn, depthView.RenderTarget, depthView.MipLevel, depthView.SubElement);

            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavDeinterleavedDepth, Resources.DeinterleavedDepths);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavDownsampledDepthMip0, Resources.DeinterleavedDepths, 0);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavDownsampledDepthMip1, Resources.DeinterleavedDepths, 1);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavDownsampledDepthMip2, Resources.DeinterleavedDepths, 2);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavDownsampledDepthMip3, Resources.DeinterleavedDepths, 3);

            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, dispatchWidth, dispatchHeight, 1);
        }
    }

    internal class CacaoPrepareNormalsPass : CacaoPass
    {
        private enum Kernels
        {
            CS_Native,
            CS_NativeFromInputNormals,
            CS_Downsampled,
            CS_DownsampledFromInputNormals,
        }

        public CacaoPrepareNormalsPass(ComputeShader shader, CacaoResources resources)
            : base(resources)
        {
            InitComputeShaders<Kernels>("prepare_normals", shader);
        }
        
        public void Execute(CommandBuffer commandBuffer, ComputeBuffer constants, in Cacao.DispatchInfo dispatchInfo, in Cacao.BufferSizeInfo bsi, bool generateNormals, bool useDownsampledSsao)
        {
            Kernels kernel;
            int dispatchWidth, dispatchHeight;

            if (generateNormals)
            {
                dispatchWidth = DispatchSize(CacaoConsts.PrepareNormalsWidth, bsi.SsaoBufferWidth);
                dispatchHeight = DispatchSize(CacaoConsts.PrepareNormalsHeight, bsi.SsaoBufferHeight);
                kernel = useDownsampledSsao ? Kernels.CS_Downsampled : Kernels.CS_Native;
            }
            else
            {
                dispatchWidth = DispatchSize(CacaoConsts.PrepareNormalsFromInputNormalsWidth, bsi.SsaoBufferWidth);
                dispatchHeight = DispatchSize(CacaoConsts.PrepareNormalsFromInputNormalsHeight, bsi.SsaoBufferHeight);
                kernel = useDownsampledSsao ? Kernels.CS_DownsampledFromInputNormals : Kernels.CS_NativeFromInputNormals;
            }
            
            int kernelIndex = KernelIndices[(int)kernel];
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CacaoShaderIDs.CbSsaoConstantsBuffer, constants, 0, Marshal.SizeOf<Cacao.Constants>());

            if (generateNormals)
            {
                var depthView = dispatchInfo.DepthView;
                commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvDepthIn, depthView.RenderTarget, depthView.MipLevel, depthView.SubElement);
            }
            else if (dispatchInfo.NormalsView.HasValue)
            {
                var normalsView = dispatchInfo.NormalsView.Value;
                commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvNormalIn, normalsView.RenderTarget, normalsView.MipLevel, normalsView.SubElement);
            }
            else
            {
                commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvNormalIn, (Texture)null);
            }

            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavDeinterleavedNormals, Resources.DeinterleavedNormals);
            
            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, dispatchWidth, dispatchHeight, 1);
        }
    }

    internal class CacaoGenerateBaseSsaoPass : CacaoPass
    {
        private enum Kernels
        {
            CS_Q3Base
        }

        public CacaoGenerateBaseSsaoPass(ComputeShader shader, CacaoResources resources)
            : base(resources)
        {
            InitComputeShaders<Kernels>("generate_base_ssao", shader);
        }
        
        public void Execute(CommandBuffer commandBuffer, ComputeBuffer constants, in Cacao.BufferSizeInfo bsi)
        {
            int kernelIndex = KernelIndices[(int)Kernels.CS_Q3Base];
            int dispatchWidth = DispatchSize(CacaoConsts.GenerateWidth, bsi.SsaoBufferWidth);
            int dispatchHeight = DispatchSize(CacaoConsts.GenerateHeight, bsi.SsaoBufferHeight);

            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CacaoShaderIDs.CbSsaoConstantsBuffer, constants, 0, Marshal.SizeOf<Cacao.Constants>());
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvDeinterleavedDepth, Resources.DeinterleavedDepths);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvDeinterleavedNormals, Resources.DeinterleavedNormals);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavSsaoBufferPing, Resources.SsaoBufferPong);   // FFX_CACAO_UAV_SSAO_REMAP_TO_PONG
            
            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, dispatchWidth, dispatchHeight, 4);
        }
    }

    internal class CacaoGenerateImportanceMapPass : CacaoPass
    {
        private enum Kernels
        {
            CS_Generate,
            CS_PostprocessA,
            CS_PostprocessB,
        }

        public CacaoGenerateImportanceMapPass(ComputeShader shader, CacaoResources resources)
            : base(resources)
        {
            InitComputeShaders<Kernels>("generate_importance_map", shader);
        }
        
        public void Execute(CommandBuffer commandBuffer, ComputeBuffer constants, in Cacao.BufferSizeInfo bsi)
        {
            int dispatchWidth = DispatchSize(CacaoConsts.ImportanceMapWidth, bsi.ImportanceMapWidth);
            int dispatchHeight = DispatchSize(CacaoConsts.ImportanceMapHeight, bsi.ImportanceMapHeight);
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CacaoShaderIDs.CbSsaoConstantsBuffer, constants, 0, Marshal.SizeOf<Cacao.Constants>());

            int kernelIndex = KernelIndices[(int)Kernels.CS_Generate];
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvSsaoBufferPong, Resources.SsaoBufferPong);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavImportanceMap, Resources.ImportanceMap);
            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, dispatchWidth, dispatchHeight, 1);

            kernelIndex = KernelIndices[(int)Kernels.CS_PostprocessA];
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvImportanceMap, Resources.ImportanceMap);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavImportanceMapPong, Resources.ImportanceMapPong);
            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, dispatchWidth, dispatchHeight, 1);
            
            kernelIndex = KernelIndices[(int)Kernels.CS_PostprocessB];
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvImportanceMapPong, Resources.ImportanceMapPong);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavImportanceMap, Resources.ImportanceMap);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavLoadCounterBuffer, Resources.LoadCounter);
            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, dispatchWidth, dispatchHeight, 1);
        }
    }

    internal class CacaoGenerateSsaoPass : CacaoPass
    {
        private enum Kernels
        {
            CS_Q0,
            CS_Q1,
            CS_Q2,
            CS_Q3,
        }

        public CacaoGenerateSsaoPass(ComputeShader shader, CacaoResources resources)
            : base(resources)
        {
            InitComputeShaders<Kernels>("generate_ssao", shader);
        }
        
        public void Execute(CommandBuffer commandBuffer, ComputeBuffer constants, in Cacao.BufferSizeInfo bsi, Cacao.Quality quality)
        {
            int kernelIndex = KernelIndices[(int)Kernels.CS_Q0 + Math.Max(0, (int)quality - 1)];
            
            int dispatchWidth, dispatchHeight, dispatchDepth;
            switch (quality)
            {
                case Cacao.Quality.Lowest:
                case Cacao.Quality.Low:
                case Cacao.Quality.Medium:
                    dispatchWidth = DispatchSize(CacaoConsts.GenerateSparseWidth, bsi.SsaoBufferWidth);
                    dispatchWidth = (dispatchWidth + 4) / 5;
                    dispatchHeight = DispatchSize(CacaoConsts.GenerateSparseHeight, bsi.SsaoBufferHeight);
                    dispatchDepth = 5;
                    break;
                case Cacao.Quality.High:
                case Cacao.Quality.Highest:
                default:
                    dispatchWidth = DispatchSize(CacaoConsts.GenerateWidth, bsi.SsaoBufferWidth);
                    dispatchHeight = DispatchSize(CacaoConsts.GenerateHeight, bsi.SsaoBufferHeight);
                    dispatchDepth = 1;
                    break;
            }

            dispatchDepth *= (quality == Cacao.Quality.Lowest) ? 2 : 4; // 2 layers for lowest, 4 for all others
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CacaoShaderIDs.CbSsaoConstantsBuffer, constants, 0, Marshal.SizeOf<Cacao.Constants>());
            
            if (quality == Cacao.Quality.Highest)
            {
                // Use adaptive descriptor set
                commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvLoadCounter, Resources.LoadCounter);
                commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvImportanceMap, Resources.ImportanceMap);
                commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvSsaoBufferPong, Resources.SsaoBufferPong);
            }
                
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvDeinterleavedDepth, Resources.DeinterleavedDepths);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvDeinterleavedNormals, Resources.DeinterleavedNormals);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavSsaoBufferPing, Resources.SsaoBufferPing);
            
            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, dispatchWidth, dispatchHeight, dispatchDepth);
        }
    }

    internal class CacaoEdgeSensitiveBlurPass : CacaoPass
    {
        private enum Kernels
        {
            CS_1Pass,
            CS_2Pass,
            CS_3Pass,
            CS_4Pass,
            CS_5Pass,
            CS_6Pass,
            CS_7Pass,
            CS_8Pass,
        }

        public CacaoEdgeSensitiveBlurPass(ComputeShader shader, CacaoResources resources)
            : base(resources)
        {
            InitComputeShaders<Kernels>("edge_sensitive_blur", shader);
        }
        
        public void Execute(CommandBuffer commandBuffer, ComputeBuffer constants, in Cacao.BufferSizeInfo bsi, Cacao.Quality quality, uint blurPassCount)
        {
            uint w = 4 * CacaoConsts.BlurWidth - 2 * blurPassCount;
            uint h = 3 * CacaoConsts.BlurHeight - 2 * blurPassCount;
            int dispatchWidth = DispatchSize(w, bsi.SsaoBufferWidth);
            int dispatchHeight = DispatchSize(h, bsi.SsaoBufferHeight);
            int dispatchDepth = (quality == Cacao.Quality.Lowest) ? 2 : 4;

            int kernelIndex = KernelIndices[(int)Kernels.CS_1Pass + blurPassCount - 1];
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CacaoShaderIDs.CbSsaoConstantsBuffer, constants, 0, Marshal.SizeOf<Cacao.Constants>());
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvSsaoBufferPing, Resources.SsaoBufferPing);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavSsaoBufferPong, Resources.SsaoBufferPong);
            
            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, dispatchWidth, dispatchHeight, dispatchDepth);
        }
    }

    internal class CacaoBilateralUpscalePass : CacaoPass
    {
        private enum Kernels
        {
            CS_5x5Smart,
            CS_5x5NonSmart,
            CS_5x5Half,
        }

        public CacaoBilateralUpscalePass(ComputeShader shader, CacaoResources resources)
            : base(resources)
        {
            InitComputeShaders<Kernels>("bilateral_upscale", shader);
        }
        
        public void Execute(CommandBuffer commandBuffer, ComputeBuffer constants, in Cacao.DispatchInfo dispatchInfo, in Cacao.BufferSizeInfo bsi, Cacao.Quality quality, bool usedBlur)
        {
            Kernels kernel;
            switch (quality)
            {
                case Cacao.Quality.Lowest:
                    kernel = Kernels.CS_5x5Half;
                    break;
                case Cacao.Quality.Low:
                case Cacao.Quality.Medium:
                    kernel = Kernels.CS_5x5NonSmart;
                    break;
                case Cacao.Quality.High:
                case Cacao.Quality.Highest:
                default:
                    kernel = Kernels.CS_5x5Smart;
                    break;
            }
            
            int kernelIndex = KernelIndices[(int)kernel];
            int dispatchWidth = DispatchSize(2 * CacaoConsts.BilateralUpscaleWidth, bsi.InputOutputBufferWidth);
            int dispatchHeight = DispatchSize(2 * CacaoConsts.BilateralUpscaleHeight, bsi.InputOutputBufferHeight);

            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CacaoShaderIDs.CbSsaoConstantsBuffer, constants, 0, Marshal.SizeOf<Cacao.Constants>());
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvSsaoBufferPing, usedBlur ? Resources.SsaoBufferPong : Resources.SsaoBufferPing);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvDeinterleavedDepth, Resources.DeinterleavedDepths);
            
            var depthView = dispatchInfo.DepthView;
            var outputView = dispatchInfo.OutputView;
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvDepthIn, depthView.RenderTarget, depthView.MipLevel, depthView.SubElement);
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavOutput, outputView.RenderTarget, outputView.MipLevel, outputView.SubElement);
            
            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, dispatchWidth, dispatchHeight, 1);
        }
    }

    internal class CacaoReinterleavePass : CacaoPass
    {
        private enum Kernels
        {
            CS_Smart,
            CS_NonSmart,
            CS_NonSmartHalf,
        }

        public CacaoReinterleavePass(ComputeShader shader, CacaoResources resources)
            : base(resources)
        {
            InitComputeShaders<Kernels>("reinterleave_apply", shader);
        }
        
        public void Execute(CommandBuffer commandBuffer, ComputeBuffer constants, in Cacao.DispatchInfo dispatchInfo, in Cacao.BufferSizeInfo bsi, Cacao.Quality quality, bool usedBlur)
        {
            Kernels kernel;
            switch (quality)
            {
                case Cacao.Quality.Lowest:
                    kernel = Kernels.CS_NonSmartHalf;
                    break;
                case Cacao.Quality.Low:
                    kernel = Kernels.CS_NonSmart;
                    break;
                default:
                    kernel = Kernels.CS_Smart;
                    break;
            }

            int kernelIndex = KernelIndices[(int)kernel];
            int dispatchWidth = DispatchSize(CacaoConsts.ApplyWidth, bsi.InputOutputBufferWidth);
            int dispatchHeight = DispatchSize(CacaoConsts.ApplyHeight, bsi.InputOutputBufferHeight);

            commandBuffer.SetComputeConstantBufferParam(ComputeShader, CacaoShaderIDs.CbSsaoConstantsBuffer, constants, 0, Marshal.SizeOf<Cacao.Constants>());
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.SrvSsaoBufferPing, usedBlur ? Resources.SsaoBufferPong : Resources.SsaoBufferPing);

            var outputView = dispatchInfo.OutputView;
            commandBuffer.SetComputeTextureParam(ComputeShader, kernelIndex, CacaoShaderIDs.UavOutput, outputView.RenderTarget, outputView.MipLevel, outputView.SubElement);
            
            commandBuffer.DispatchCompute(ComputeShader, kernelIndex, dispatchWidth, dispatchHeight, 1);
        }
    }
}