using System;
using UnityEngine;
using UnityEngine.Rendering;
using UnityEngine.Rendering.Universal;

namespace X.Rendering.Feature
{
    public class AutoExposure : ScriptableRendererFeature
    {
        [Serializable]
        public class Settings
        {
            public RenderPassEvent renderPassEvent = RenderPassEvent.BeforeRenderingPostProcessing;
            public Material BlitMat;
            public ComputeShader Compute;
        }

        [SerializeField]
        Settings settings;
        AutoExposurePass autoExposurePass;

        public override void AddRenderPasses(ScriptableRenderer renderer, ref RenderingData renderingData)
        {
            renderer.EnqueuePass(autoExposurePass);
        }

        public override void Create()
        {
            autoExposurePass = new(settings);
        }

        class AutoExposurePass : ScriptableRenderPass
        {
            private Settings settings;
            private int computeCounter;
            private ProfilingSampler profiler;
            private ComputeBuffer outputBuffer;

            private static readonly int Dims = Shader.PropertyToID("_Dims");
            private static readonly int Source = Shader.PropertyToID("_Source");
            private static readonly int TempOut = Shader.PropertyToID("_OutBuffer");
            private static readonly int ScreenExposureProp = Shader.PropertyToID("_Exposure");
            private static readonly int ExposureRangeProp = Shader.PropertyToID("_Range");
            private static readonly int WhitePointProp = Shader.PropertyToID("_WhitePoint");
            private static readonly int MaxBrightness = Shader.PropertyToID("_MaxBrightness");
            private static readonly int BlitTextureId = Shader.PropertyToID("_BlitTexture");
            

            private const int ThreadGroupSize = 32;

            private Vector2Int lastSize = Vector2Int.zero;

            private static float targetExposure  = 0.1f;
            private static float currentExposure = 0.1f;
            private static AutoExposureVolumeProfile autoExposureVolume = null;
            private static float lastTime = 0.0f;
            private static bool autoExposureing = false;

            public AutoExposurePass(Settings settings)
            {
                this.settings = settings;
                renderPassEvent = settings.renderPassEvent;
                profiler = new(nameof(AutoExposurePass));
            }

            public override void Execute(ScriptableRenderContext context, ref RenderingData renderingData)
            {
                var stack = VolumeManager.instance.stack;
                autoExposureVolume = stack.GetComponent<AutoExposureVolumeProfile>();
                if (!autoExposureVolume || !autoExposureVolume.IsActive())
                {
                    return;
                }

                var cmd = renderingData.commandBuffer;
                using var scp = new ProfilingScope(cmd, profiler);
                if (computeCounter % Mathf.Max(autoExposureVolume.framesPerCompute.value, 1.0f) == 0)
                {
                    computeCounter = 0;
                    var textureDesc = renderingData.cameraData.cameraTargetDescriptor;
                    
                    Vector2Int size = new(textureDesc.width, textureDesc.height);
                    if(outputBuffer ==  null || lastSize != size)
                    {
                        lastSize = size;
                        outputBuffer?.Dispose();
                        outputBuffer = new ComputeBuffer(Mathf.CeilToInt((float)size.y / ThreadGroupSize * 2), 4, ComputeBufferType.Structured)
                        {
                            name = "Auto Exposure Output Buffer"
                        };
                    }
                    cmd.SetComputeIntParams(settings.Compute, Dims, size.x, size.y);
                    
                    cmd.SetComputeTextureParam(settings.Compute, 0, Source, renderingData.cameraData.renderer.cameraColorTargetHandle);
                    cmd.SetComputeBufferParam(settings.Compute, 0, TempOut, outputBuffer);
                    cmd.DispatchCompute(settings.Compute, 0, size.y, 1, 1);
                    AsyncGPUReadback.Request(outputBuffer, OnCompleteReadBack);
                }
                ++computeCounter;
                if (!autoExposureing)
                {
                    //return;
                }
                UpdateExposure();
                settings.BlitMat.SetFloat(ScreenExposureProp, currentExposure);
                settings.BlitMat.SetVector(ExposureRangeProp, autoExposureVolume.exposureRange.value);
                settings.BlitMat.SetFloat(WhitePointProp, autoExposureVolume.whitePoint.value);
                settings.BlitMat.SetFloat(MaxBrightness, autoExposureVolume.brightnessLimit.value);
                var renderer = renderingData.cameraData.renderer;
                settings.BlitMat.SetTexture(BlitTextureId, renderer.cameraColorTargetHandle);
                var destination = renderer.GetCameraColorFrontBuffer(cmd);
                cmd.SetRenderTarget(destination, loadAction: RenderBufferLoadAction.DontCare, storeAction: RenderBufferStoreAction.Store);
                cmd.DrawProcedural(Matrix4x4.identity, settings.BlitMat, 0, MeshTopology.Triangles, 3);
                renderer.SwapColorBuffer(cmd);
            }

            private static void OnCompleteReadBack(AsyncGPUReadbackRequest request)
            {
                if (request.hasError)
                {
                    return;
                }

                // Get Compute result
                float[] groupValues = request.GetData<float>().ToArray();

                // Add up all the workgroup's results
                double pixelLumTotal = 0;
                uint pixels = 0;
                for (int i = 0; i < groupValues.Length / 2; i++)
                {
                    pixelLumTotal += groupValues[i * 2];
                    pixels += (uint)Mathf.CeilToInt(groupValues[i * 2 + 1]);
                }

                // Average the results and set as the target exposure
                targetExposure = (float)(pixelLumTotal / pixels);
                autoExposureing = true;
            }

            public static float ExpDecay(float a, float b, float decay, float deltaTime) => b + (a - b) * Mathf.Exp(-decay * deltaTime);

            private static void UpdateExposure()
            {
                // Delta time, using the built-in one seems to flicker a lot, weird
                float deltaTime = Time.time - lastTime;
                lastTime = Time.time;

                // Exposure difference, skip if zero
                float diff = targetExposure - currentExposure;
                if (Mathf.Approximately(diff,0))
                {
                    autoExposureing = false;
                    return;
                }

                float decay = diff > 0 ? autoExposureVolume.increaseSpeed.value : autoExposureVolume.decreaseSpeed.value;
                currentExposure = ExpDecay(currentExposure, targetExposure, decay, deltaTime);

                if (float.IsNaN(currentExposure))
                {
                    currentExposure = 0.1f;
                }
            }
        }
    }
}