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

namespace FidelityFX
{
    /// <summary>
    /// Port of ffx_cacao_impl.cpp/.h
    /// </summary>
    public class CacaoContext
    {
        private Cacao.Settings _settings;
        private Cacao.BufferSizeInfo _bufferSizeInfo;
        private bool _useDownsampledSsao;

        private readonly CacaoResources _resources = new CacaoResources();

        private CacaoClearLoadCounterPass _clearLoadCounterPass;
        private CacaoPrepareDepthsPass _prepareDepthsPass;
        private CacaoPrepareNormalsPass _prepareNormalsPass;
        private CacaoGenerateBaseSsaoPass _generateBaseSsaoPass;
        private CacaoGenerateImportanceMapPass _generateImportanceMapPass;
        private CacaoGenerateSsaoPass _generateSsaoPass;
        private CacaoEdgeSensitiveBlurPass _edgeSensitiveBlurPass;
        private CacaoBilateralUpscalePass _bilateralUpscalePass;
        private CacaoReinterleavePass _reinterleavePass;

        private ComputeBuffer _constantsBuffer;
        private readonly Cacao.Constants[] _constants = new Cacao.Constants[1];
        private ref Cacao.Constants Constants => ref _constants[0];
        
        public bool Init(CacaoShaders shaders)
        {
            // Create constant buffers
            _constantsBuffer = CreateConstantBuffer<Cacao.Constants>("FFX_CACAO_CONSTANT_BUFFER");

            // Create load counter resource
            if (!_resources.CreateLoadCounter())
            {
                return false;
            }

            // Initialize compute shaders
            CreatePasses(shaders);
            return true;
        }

        private void CreatePasses(CacaoShaders shaders)
        {
            _clearLoadCounterPass = new CacaoClearLoadCounterPass(shaders.clearLoadCounter, _resources);
            _prepareDepthsPass = new CacaoPrepareDepthsPass(shaders.prepareDepths, _resources);
            _prepareNormalsPass = new CacaoPrepareNormalsPass(shaders.prepareNormals, _resources);
            _generateBaseSsaoPass = new CacaoGenerateBaseSsaoPass(shaders.generateSsao, _resources);
            _generateImportanceMapPass = new CacaoGenerateImportanceMapPass(shaders.generateImportanceMap, _resources);
            _generateSsaoPass = new CacaoGenerateSsaoPass(shaders.generateSsao, _resources);
            _edgeSensitiveBlurPass = new CacaoEdgeSensitiveBlurPass(shaders.edgeSensitiveBlur, _resources);
            _bilateralUpscalePass = new CacaoBilateralUpscalePass(shaders.bilateralUpscale, _resources);
            _reinterleavePass = new CacaoReinterleavePass(shaders.reinterleave, _resources);
        }

        public void Destroy()
        {
            DestroyPass(ref _clearLoadCounterPass);
            DestroyPass(ref _prepareDepthsPass);
            DestroyPass(ref _prepareNormalsPass);
            DestroyPass(ref _generateBaseSsaoPass);
            DestroyPass(ref _generateImportanceMapPass);
            DestroyPass(ref _generateSsaoPass);
            DestroyPass(ref _edgeSensitiveBlurPass);
            DestroyPass(ref _bilateralUpscalePass);
            DestroyPass(ref _reinterleavePass);
            
            _resources.DestroyLoadCounter();

            DestroyConstantBuffer(ref _constantsBuffer);
        }

        public bool InitScreenSizeDependentResources(in Cacao.ScreenSizeInfo info)
        {
            _useDownsampledSsao = info.UseDownsampledSsao;
            
            UpdateBufferSizeInfo(info.Width, info.Height, _useDownsampledSsao, out _bufferSizeInfo);

            // Create textures
            return _resources.CreateResources(_bufferSizeInfo);
        }

        public void DestroyScreenSizeDependentResources()
        {
            _resources.DestroyResources();
        }
        
        public void UpdateSettings(in Cacao.Settings settings)
        {
            _settings = settings;
        }
        
        private static readonly ProfilerMarker PrepareMarker = new ProfilerMarker("Prepare Downsampled Depth, Normals and Mips");
        private static readonly ProfilerMarker BasePassMarker = new ProfilerMarker("Generate High Quality Base Pass");
        private static readonly ProfilerMarker BaseSsaoMarker = new ProfilerMarker("Base SSAO");
        private static readonly ProfilerMarker ImportanceMapMarker = new ProfilerMarker("Importance Map");
        private static readonly ProfilerMarker GenerateSsaoMarker = new ProfilerMarker("Generate SSAO");
        private static readonly ProfilerMarker DeinterleavedBlurMarker = new ProfilerMarker("Deinterleaved Blur");
        private static readonly ProfilerMarker BilateralUpsampleMarker = new ProfilerMarker("Bilateral Upsample");
        private static readonly ProfilerMarker ReinterleaveMarker = new ProfilerMarker("Reinterleave");

        public void Draw(CommandBuffer commandBuffer, in Cacao.DispatchInfo dispatchInfo, in Matrix4x4 proj, in Matrix4x4 normalsToView)
        {
            // Update constant buffers
            UpdateConstants(ref Constants, _settings, _bufferSizeInfo, proj, normalsToView);
            commandBuffer.SetBufferData(_constantsBuffer, _constants);
            
            // Clear load counter
            // The same could be accomplished with a ClearRenderTarget on the LoadCounter texture, but Unity does not allow that with async compute
            _clearLoadCounterPass.Execute(commandBuffer, _constantsBuffer);
            
            // Prepare depths, normals and mips
            commandBuffer.BeginSample(PrepareMarker);
            _prepareDepthsPass.Execute(commandBuffer, _constantsBuffer, dispatchInfo, _bufferSizeInfo, _settings.qualityLevel, _useDownsampledSsao);
            _prepareNormalsPass.Execute(commandBuffer, _constantsBuffer, dispatchInfo, _bufferSizeInfo, _settings.generateNormals, _useDownsampledSsao);
            commandBuffer.EndSample(PrepareMarker);
            
            // Base pass for the highest quality setting
            if (_settings.qualityLevel == Cacao.Quality.Highest)
            {
                commandBuffer.BeginSample(BasePassMarker);
                
                commandBuffer.BeginSample(BaseSsaoMarker);
                _generateBaseSsaoPass.Execute(commandBuffer, _constantsBuffer, _bufferSizeInfo);
                commandBuffer.EndSample(BaseSsaoMarker);
                
                commandBuffer.BeginSample(ImportanceMapMarker);
                _generateImportanceMapPass.Execute(commandBuffer, _constantsBuffer, _bufferSizeInfo);
                commandBuffer.EndSample(ImportanceMapMarker);
                
                commandBuffer.EndSample(BasePassMarker);
            }
            
            // Main SSAO generation
            commandBuffer.BeginSample(GenerateSsaoMarker);
            _generateSsaoPass.Execute(commandBuffer, _constantsBuffer, _bufferSizeInfo, _settings.qualityLevel);
            commandBuffer.EndSample(GenerateSsaoMarker);
            
            // De-interleaved blur
            uint blurPassCount = Math.Clamp(_settings.blurPassCount, 0, 8);
            if (blurPassCount > 0)
            {
                commandBuffer.BeginSample(DeinterleavedBlurMarker);
                _edgeSensitiveBlurPass.Execute(commandBuffer, _constantsBuffer, _bufferSizeInfo, _settings.qualityLevel, blurPassCount);
                commandBuffer.EndSample(DeinterleavedBlurMarker);
            }

            if (_useDownsampledSsao)
            {
                // Bilateral upsample
                commandBuffer.BeginSample(BilateralUpsampleMarker);
                _bilateralUpscalePass.Execute(commandBuffer, _constantsBuffer, dispatchInfo, _bufferSizeInfo, _settings.qualityLevel, blurPassCount > 0);
                commandBuffer.EndSample(BilateralUpsampleMarker);
            }
            else
            {
                // Reinterleave
                commandBuffer.BeginSample(ReinterleaveMarker);
                _reinterleavePass.Execute(commandBuffer, _constantsBuffer, dispatchInfo, _bufferSizeInfo, _settings.qualityLevel, blurPassCount > 0);
                commandBuffer.EndSample(ReinterleaveMarker);
            }
        }
        
        /// <summary>
        /// Update buffer size info for resolution width x height.
        /// </summary>
        /// <param name="width">Screen width.</param>
        /// <param name="height">Screen height.</param>
        /// <param name="useDownsampledSsao">Whether FFX CACAO should use downsampling.</param>
        /// <param name="bsi">Output buffer size info.</param>
        private static void UpdateBufferSizeInfo(uint width, uint height, bool useDownsampledSsao, out Cacao.BufferSizeInfo bsi)
        {
            uint halfWidth = (width + 1) / 2;
            uint halfHeight = (height + 1) / 2;
            uint quarterWidth = (halfWidth + 1) / 2;
            uint quarterHeight = (halfHeight + 1) / 2;
            uint eighthWidth = (quarterWidth + 1) / 2;
            uint eighthHeight = (quarterHeight + 1) / 2;

            uint depthBufferWidth = width;
            uint depthBufferHeight = height;
            uint depthBufferHalfWidth = halfWidth;
            uint depthBufferHalfHeight = halfHeight;
            uint depthBufferQuarterWidth = quarterWidth;
            uint depthBufferQuarterHeight = quarterHeight;

            uint depthBufferXOffset = 0;
            uint depthBufferYOffset = 0;
            uint depthBufferHalfXOffset = 0;
            uint depthBufferHalfYOffset = 0;
            uint depthBufferQuarterXOffset = 0;
            uint depthBufferQuarterYOffset = 0;

            bsi.InputOutputBufferWidth = width;
            bsi.InputOutputBufferHeight = height;
            bsi.DepthBufferXOffset = depthBufferXOffset;
            bsi.DepthBufferYOffset = depthBufferYOffset;
            bsi.DepthBufferWidth = depthBufferWidth;
            bsi.DepthBufferHeight = depthBufferHeight;

            if (useDownsampledSsao)
            {
                bsi.SsaoBufferWidth = quarterWidth;
                bsi.SsaoBufferHeight = quarterHeight;
                bsi.DeinterleavedDepthBufferXOffset = depthBufferQuarterXOffset;
                bsi.DeinterleavedDepthBufferYOffset = depthBufferQuarterYOffset;
                bsi.DeinterleavedDepthBufferWidth = depthBufferQuarterWidth;
                bsi.DeinterleavedDepthBufferHeight = depthBufferQuarterHeight;
                bsi.ImportanceMapWidth = eighthWidth;
                bsi.ImportanceMapHeight = eighthHeight;
                bsi.DownsampledSsaoBufferWidth = halfWidth;
                bsi.DownsampledSsaoBufferHeight = halfHeight;
            }
            else
            {
                bsi.SsaoBufferWidth = halfWidth;
                bsi.SsaoBufferHeight = halfHeight;
                bsi.DeinterleavedDepthBufferXOffset = depthBufferHalfXOffset;
                bsi.DeinterleavedDepthBufferYOffset = depthBufferHalfYOffset;
                bsi.DeinterleavedDepthBufferWidth = depthBufferHalfWidth;
                bsi.DeinterleavedDepthBufferHeight = depthBufferHalfHeight;
                bsi.ImportanceMapWidth = quarterWidth;
                bsi.ImportanceMapHeight = quarterHeight;
                bsi.DownsampledSsaoBufferWidth = 1;
                bsi.DownsampledSsaoBufferHeight = 1;
            }
        }

        /// <summary>
        /// Update the contents of the FFX CACAO constant buffer (an FFX_CACAO_Constants struct). Note, this function does not update per pass constants.
        /// </summary>
        /// <param name="consts">Cacao.Constants constant buffer.</param>
        /// <param name="settings">Cacao.Settings settings.</param>
        /// <param name="bufferSizeInfo">Cacao.BufferSizeInfo buffer size info.</param>
        /// <param name="proj">Projection matrix for the frame.</param>
        /// <param name="normalsToView">Normals world space to view space matrix for the frame.</param>
        private static void UpdateConstants(ref Cacao.Constants consts, in Cacao.Settings settings, in Cacao.BufferSizeInfo bufferSizeInfo, in Matrix4x4 proj, in Matrix4x4 normalsToView)
        {
            consts.BilateralSigmaSquared = settings.bilateralSigmaSquared;
            consts.BilateralSimilarityDistanceSigma = settings.bilateralSimilarityDistanceSigma;

            if (settings.generateNormals)
            {
                consts.NormalsWorldToViewspaceMatrix = Matrix4x4.identity;
            }
            else
            {
                consts.NormalsWorldToViewspaceMatrix = normalsToView;
            }

            // used to get average load per pixel; 9.0 is there to compensate for only doing every 9th InterlockedAdd in PSPostprocessImportanceMapB for performance reasons
            consts.LoadCounterAvgDiv = 9.0f / (bufferSizeInfo.ImportanceMapWidth * bufferSizeInfo.ImportanceMapHeight * 255);

            float depthLinearizeMul = -proj.m32;    // float depthLinearizeMul = ( clipFar * clipNear ) / ( clipFar - clipNear );
            float depthLinearizeAdd = proj.m22;     // float depthLinearizeAdd = clipFar / ( clipFar - clipNear );
            // correct the handedness issue. need to make sure this below is correct, but I think it is.
            if (depthLinearizeMul * depthLinearizeAdd < 0)
                depthLinearizeAdd = -depthLinearizeAdd;

            consts.DepthUnpackConsts.x = depthLinearizeMul;
            consts.DepthUnpackConsts.y = depthLinearizeAdd;

            float tanHalfFovY = 1.0f / proj.m11;    // = tanf( drawContext.Camera.GetYFOV( ) * 0.5f );
            float tanHalfFovX = 1.0f / proj.m00;    // = tanHalfFovY * drawContext.Camera.GetAspect( );
            consts.CameraTanHalfFOV.x = tanHalfFovX;
            consts.CameraTanHalfFOV.y = tanHalfFovY;

            consts.NDCToViewMul.x = consts.CameraTanHalfFOV.x * 2.0f;
            consts.NDCToViewMul.y = consts.CameraTanHalfFOV.y * -2.0f;
            consts.NDCToViewAdd.x = consts.CameraTanHalfFOV.x * -1.0f;
            consts.NDCToViewAdd.y = consts.CameraTanHalfFOV.y * 1.0f;

            float ratio = (float)bufferSizeInfo.InputOutputBufferWidth / bufferSizeInfo.DepthBufferWidth;
            float border = (1.0f - ratio) / 2.0f;
            for (int i = 0; i < 2; ++i)
            {
                consts.DepthBufferUVToViewMul[i] = consts.NDCToViewMul[i] / ratio;
                consts.DepthBufferUVToViewAdd[i] = consts.NDCToViewAdd[i] - consts.NDCToViewMul[i] * border / ratio;
            }

            consts.EffectRadius = Mathf.Clamp(settings.radius, 0.0f, 100000.0f);
            consts.EffectShadowStrength = Mathf.Clamp(settings.shadowMultiplier * 4.3f, 0.0f, 10.0f);
            consts.EffectShadowPow = Mathf.Clamp(settings.shadowPower, 0.0f, 10.0f);
            consts.EffectShadowClamp = Mathf.Clamp(settings.shadowClamp, 0.0f, 1.0f);
            consts.EffectFadeOutMul = -1.0f / (settings.fadeOutTo - settings.fadeOutFrom);
            consts.EffectFadeOutAdd = settings.fadeOutFrom / (settings.fadeOutTo - settings.fadeOutFrom) + 1.0f;
            consts.EffectHorizonAngleThreshold = Mathf.Clamp(settings.horizonAngleThreshold, 0.0f, 1.0f);

            // 1.2 seems to be around the best trade off - 1.0 means on-screen radius will stop/slow growing when the camera is at 1.0 distance, so, depending on FOV, basically filling up most of the screen
            // This setting is viewspace-dependent and not screen size dependent intentionally, so that when you change FOV the effect stays (relatively) similar.
            float effectSamplingRadiusNearLimit = settings.radius * 1.2f;

            // if the depth precision is switched to 32bit float, this can be set to something closer to 1 (0.9999 is fine)
            consts.DepthPrecisionOffsetMod = 0.9992f;
            
            // Special settings for lowest quality level - just nerf the effect a tiny bit
            if (settings.qualityLevel <= Cacao.Quality.Low)
            {
                //consts.EffectShadowStrength *= 0.9f;
                effectSamplingRadiusNearLimit *= 1.50f;

                if (settings.qualityLevel == Cacao.Quality.Lowest)
                    consts.EffectRadius *= 0.8f;
            }

            effectSamplingRadiusNearLimit /= tanHalfFovY;   // to keep the effect same regardless of FOV

            consts.EffectSamplingRadiusNearLimitRec = 1.0f / effectSamplingRadiusNearLimit;
            consts.AdaptiveSampleCountLimit = settings.adaptiveQualityLimit;
            consts.NegRecEffectRadius = -1.0f / consts.EffectRadius;
            consts.InvSharpness = Mathf.Clamp01(1.0f - settings.sharpness);
            consts.DetailAOStrength = settings.detailShadowStrength;
            
            // set buffer size constants.
            consts.SSAOBufferDimensions = new Vector2(bufferSizeInfo.SsaoBufferWidth, bufferSizeInfo.SsaoBufferHeight);
            consts.SSAOBufferInverseDimensions = new Vector2(1.0f / bufferSizeInfo.SsaoBufferWidth, 1.0f / bufferSizeInfo.SsaoBufferHeight);

            consts.DepthBufferDimensions = new Vector2(bufferSizeInfo.DepthBufferWidth, bufferSizeInfo.DepthBufferHeight);
            consts.DepthBufferInverseDimensions = new Vector2(1.0f / bufferSizeInfo.DepthBufferWidth, 1.0f / bufferSizeInfo.DepthBufferHeight);

            consts.DepthBufferOffset = new Vector2Int((int)bufferSizeInfo.DepthBufferXOffset, (int)bufferSizeInfo.DepthBufferYOffset);

            consts.InputOutputBufferDimensions = new Vector2(bufferSizeInfo.InputOutputBufferWidth, bufferSizeInfo.InputOutputBufferHeight);
            consts.InputOutputBufferInverseDimensions = new Vector2(1.0f / bufferSizeInfo.InputOutputBufferWidth, 1.0f / bufferSizeInfo.InputOutputBufferHeight);

            consts.ImportanceMapDimensions = new Vector2(bufferSizeInfo.ImportanceMapWidth, bufferSizeInfo.ImportanceMapHeight);
            consts.ImportanceMapInverseDimensions = new Vector2(1.0f / bufferSizeInfo.ImportanceMapWidth, 1.0f / bufferSizeInfo.ImportanceMapHeight);

            consts.DeinterleavedDepthBufferDimensions = new Vector2(bufferSizeInfo.DeinterleavedDepthBufferWidth, bufferSizeInfo.DeinterleavedDepthBufferHeight);
            consts.DeinterleavedDepthBufferInverseDimensions = new Vector2(1.0f / bufferSizeInfo.DeinterleavedDepthBufferWidth, 1.0f / bufferSizeInfo.DeinterleavedDepthBufferHeight);

            consts.DeinterleavedDepthBufferOffset = new Vector2(bufferSizeInfo.DeinterleavedDepthBufferXOffset, bufferSizeInfo.DeinterleavedDepthBufferYOffset);
            consts.DeinterleavedDepthBufferNormalisedOffset = consts.DeinterleavedDepthBufferOffset / consts.DeinterleavedDepthBufferDimensions;

            if (!settings.generateNormals)
            {
                consts.NormalsUnpackMul = 2.0f;     // inputs->NormalsUnpackMul;
                consts.NormalsUnpackAdd = -1.0f;    // inputs->NormalsUnpackAdd;
            }
            else
            {
                consts.NormalsUnpackMul = 2.0f;
                consts.NormalsUnpackAdd = -1.0f;
            }
            
            consts.BlurNumPasses = settings.qualityLevel == Cacao.Quality.Lowest ? 2 : 4;
            
            float additionalAngleOffset = settings.temporalSupersamplingAngleOffset;  // if using temporal supersampling approach (like "Progressive Rendering Using Multi-frame Sampling" from GPU Pro 7, etc.)
            float additionalRadiusScale = settings.temporalSupersamplingRadiusOffset;  // if using temporal supersampling approach (like "Progressive Rendering Using Multi-frame Sampling" from GPU Pro 7, etc.)

            for (int passId = 0; passId < 4; ++passId)
            {
                ref Vector4 perPassFullResUVOffset = ref consts.GetPerPassFullResUVOffset(passId);
                perPassFullResUVOffset.x = (float)(passId % 2) / bufferSizeInfo.SsaoBufferWidth;
                perPassFullResUVOffset.y = (float)(passId / 2) / bufferSizeInfo.SsaoBufferHeight;
                perPassFullResUVOffset.z = 0;
                perPassFullResUVOffset.w = 0;

                ref Cacao.PatternRotScaleMatrices patternRotScaleMatrices = ref consts.GetPatternRotScaleMatrices(passId);
                const int subPassCount = 5;
                for (int subPass = 0; subPass < subPassCount; subPass++)
                {
                    int a = passId;
                    int b = SpMap[subPass];

                    float angle0 = (a + (float)b / subPassCount) * Mathf.PI * 0.5f;
                    float ca = Mathf.Cos(angle0);
                    float sa = Mathf.Sin(angle0);

                    float scale = 1.0f + (a - 1.5f + (b - (subPassCount - 1.0f) * 0.5f) / subPassCount) * 0.07f;

                    ref Vector4 patternRotScaleMatrix = ref patternRotScaleMatrices.GetPatternRotScaleMatrix(subPass);
                    patternRotScaleMatrix.x = scale * ca;
                    patternRotScaleMatrix.y = scale * -sa;
                    patternRotScaleMatrix.z = -scale * sa;
                    patternRotScaleMatrix.w = -scale * ca;
                }
            }
        }

        private static readonly int[] SpMap = { 0, 1, 4, 3, 2 };
        
        private static ComputeBuffer CreateConstantBuffer<TConstants>(string name) where TConstants: struct
        {
            return new ComputeBuffer(1, Marshal.SizeOf<TConstants>(), ComputeBufferType.Constant) { name = name };
        }
        
        private static void DestroyConstantBuffer(ref ComputeBuffer bufferRef)
        {
            if (bufferRef == null)
                return;
            
            bufferRef.Release();
            bufferRef = null;
        }

        private static void DestroyPass<TPass>(ref TPass pass)
            where TPass: CacaoPass
        {
            if (pass == null)
                return;
            
            pass.Dispose();
            pass = null;
        }
    }
}