using System;
using Unity.Collections;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;
using UnityEngine.Rendering.Universal;

namespace X.Rendering.Feature
{
    internal class CullResult : IDisposable
    {
        public RenderTexture ResultTex;
        public NativeArray<half> ResultArray;
        public bool ReadDone = false;
        public bool DataReady = false;
        public Action<AsyncGPUReadbackRequest> ReadBackAction;

        private void ReadBack(AsyncGPUReadbackRequest readback)
        {
            if (readback.done && !readback.hasError)
            {
                var data = readback.GetData<half>();
                if (ResultArray.IsCreated)
                {
                    ResultArray.CopyFrom(data);
                }
                ReadDone = true;
            }
            else
            {
                ReadDone = false;
            }

            DataReady = true;
        }

        public CullResult CreateResources(int texSize)
        {
            ResultArray.Dispose();
            ResultArray = new(texSize * texSize, Allocator.Persistent);
            if (ResultTex != null)
            {
                ResultTex.Release();
            }

            ResultTex = new RenderTexture(texSize, texSize, 0, RenderTextureFormat.RHalf);
            ResultTex.filterMode = FilterMode.Point;
            ResultTex.enableRandomWrite = true;
            ResultTex.Create();
            ReadBackAction = ReadBack;
            return this;
        }

        public void Dispose()
        {
            if (ResultTex != null)
            {
                ResultTex.Release();
            }
            ResultTex = null;
            ResultArray.Dispose();
        }
    }

    public class HierarchicalZOcclusionCullFeature : ScriptableRendererFeature
    {
        public enum EDepthPyramidFunc
        {
            CopyDepth,
            Compute,
            SPD
        }

        [Serializable]
        public class Settings
        {
            public EDepthPyramidFunc PyramidFunc = EDepthPyramidFunc.CopyDepth;
            public bool SkipThreeMip;
            // 留历史 depth，做 tow pass cull
            public bool UseThreeFrameReadback = true;
            public ComputeShader ComputeShader;
            public ComputeShader Spd;
            public Material CopyDepth;

            //TODO: 通过使用上一帧的摄像机位置（包括矩阵）和上一帧的深度图做剔除，储存已经被剔除的物体和未被剔除的物体，然后绘制未被剔除的物体到GBuffer（包含深度图），
            //再二次生成HiZ DepthTexture，并对已经被剔除的物体使用一遍新的深度和当前摄像机的位置做一次剔除判断
            public bool UseTowCullPass;
            public bool UseCompute;
            public Material CullMat;
            public ComputeShader CullShader;
        }

        [SerializeField]
        private Settings settings;

        private HizOcclusionCullPass hizPass;
        private static HierarchicalZOcclusionCullFeature instance;

        public override void AddRenderPasses(ScriptableRenderer renderer, ref RenderingData renderingData)
        {
            if ((renderingData.cameraData.cameraType == CameraType.Game))
            {
                renderer.EnqueuePass(hizPass);
            }
        }

        public override void Create()
        {
            hizPass = new(settings);
            instance = this;
        }

        protected override void Dispose(bool disposing)
        {
            base.Dispose(disposing);
            if(hizPass != null)
            {
                hizPass.Dispose();
            }
        }

        internal static CullResult GetCullResult()
        {
            return instance.hizPass.GetCullResult();
        }

        internal class HizOcclusionCullPass : ScriptableRenderPass, IDisposable
        {
            private readonly Settings settings;

            private ProfilingSampler profiler;

            private Vector2Int cachedHardwareTextureSize;
            private bool cachedSkipThreeMip;
            private Vector2Int mip0SizeNOP;
            private Vector2Int depthPyramidTextureSize;
            private int mipLevelCount;
            private Vector2Int[] mipLevelOffsets;
            private Vector2Int[] mipLevelSizes;
            private Vector4[] mipOffsetAndSizes = new Vector4[16];

            private RTHandle[] depthPyramidTexs = new RTHandle[3];
            private static string[] depthPyramidNames = new string[3] { "_DepthPyramidTextureA", "_DepthPyramidTextureB", "_DepthPyramidTextureC" };

            RenderTexture SpdAtomicCounter;


            CullResult[] cullResults;
            HizOcclusionCullPass instance; 

            public HizOcclusionCullPass(Settings settings)
            {
                profiler = new ProfilingSampler(nameof(HizOcclusionCullPass));

                renderPassEvent = RenderPassEvent.BeforeRenderingPostProcessing;
                this.settings = settings;
                mipLevelOffsets = new Vector2Int[16];
                mipLevelSizes = new Vector2Int[16];
                SpdAtomicCounter = new RenderTexture(1, 1, 0, GraphicsFormat.R32_UInt) { name = "FSR2_SpdAtomicCounter", enableRandomWrite = true };
                SpdAtomicCounter.Create();

                int texSize = 64;

                cullResults = new CullResult[3]
                {
                    new CullResult().CreateResources(texSize),
                    new CullResult().CreateResources(texSize),
                    new CullResult().CreateResources(texSize),
                };
            }

            public static int HizIndex { get; private set; }

            private int GetHizIndex()
            {
                if(settings.UseThreeFrameReadback)
                {
                    HizIndex = Time.frameCount % 3;
                }
                else
                {
                    HizIndex = 0;
                }

                return HizIndex;
            }
            
            public CullResult GetCullResult()
            {
                var hizIndex = GetHizIndex();
                return cullResults[GetLastCullIndex(hizIndex)];
            }

            public void ComputePackedMipChainInfo(Vector2Int viewportSize)
            {
                if (cachedHardwareTextureSize == viewportSize && cachedSkipThreeMip == settings.SkipThreeMip)
                {
                    return;
                }

                cachedHardwareTextureSize = viewportSize;
                cachedSkipThreeMip = settings.SkipThreeMip;

                mip0SizeNOP = viewportSize;
                int resizeX = viewportSize.x;
                int resizeY = viewportSize.y;

                resizeX = Mathf.IsPowerOfTwo(viewportSize.x) ? viewportSize.x : Mathf.NextPowerOfTwo(viewportSize.x);
                resizeY = Mathf.IsPowerOfTwo(viewportSize.y) ? viewportSize.y : Mathf.NextPowerOfTwo(viewportSize.y);
                //if (resizeX > viewportSize.x)
                //    resizeX /= 2;
                //if (resizeY > viewportSize.y)
                //    resizeY /= 2;

                Vector2Int hardwareTextureSize = new Vector2Int(resizeX, resizeY);
                mipLevelOffsets[0] = Vector2Int.zero;
                mipLevelSizes[0] = Vector2Int.zero;
                Vector2Int mipSize = hardwareTextureSize;
                Vector2Int texSize = Vector2Int.zero;
                mipLevelCount = 1 + (int)math.floor(math.log2(math.max(resizeX, resizeY)));
   
                for (int mipLevel = 1; mipLevel < mipLevelCount; mipLevel++)
                {
                    mipSize.x = System.Math.Max(1, (mipSize.x + 1) >> 1);
                    mipSize.y = System.Math.Max(1, (mipSize.y + 1) >> 1);

                    mipLevelSizes[mipLevel] = mipSize;

                    Vector2Int prevMipBegin = mipLevelOffsets[mipLevel - 1];
                    Vector2Int prevMipEnd = prevMipBegin + mipLevelSizes[mipLevel - 1];

                    Vector2Int mipBegin = new Vector2Int();

                    if ((mipLevel & 1) != 0)
                    {
                        mipBegin.x = prevMipBegin.x;
                        mipBegin.y = prevMipEnd.y;
                    }
                    else
                    {
                        mipBegin.x = prevMipEnd.x;
                        mipBegin.y = prevMipBegin.y;
                    }
                    if (settings.SkipThreeMip && mipLevel == 4) //跳过前面3级的Mip,从第四级开始算
                    {
                        mipBegin = Vector2Int.zero;
                        texSize = Vector2Int.zero;
                    }

                    mipLevelOffsets[mipLevel] = mipBegin;

                    texSize.x = System.Math.Max(texSize.x, mipBegin.x + mipSize.x);
                    texSize.y = System.Math.Max(texSize.y, mipBegin.y + mipSize.y);
                }

                mipLevelSizes[0] = hardwareTextureSize;
                //RT实际大小
                depthPyramidTextureSize = new Vector2Int((int)Mathf.Ceil((float)texSize.x), (int)Mathf.Ceil((float)texSize.y));

                for (int i = 0;i < mipLevelSizes.Length; i++)
                {
                    mipOffsetAndSizes[i] = new Vector4(mipLevelOffsets[i].x, mipLevelOffsets[i].y, mipLevelSizes[i].x,
                        mipLevelSizes[i].y);
                }
            }

            private static int[] depthMipId = null;

            private void DoCopyDepth(CommandBuffer cmd, Texture depthTex)
            {
                if (depthMipId == null)
                {
                    depthMipId = new int[32];
                    for (int i = 0; i < depthMipId.Length; i++)
                    {
                        depthMipId[i] = Shader.PropertyToID("_DepthMip" + i);
                    }
                }

                var sampleName = "DepthDownSample";
                cmd.BeginSample(sampleName);
                cmd.SetViewProjectionMatrices(Matrix4x4.identity, Matrix4x4.identity);
                cmd.SetGlobalTexture(HizShaderIds.DepthInputId, depthTex);
                for (int i = 0; i < mipLevelCount; i++)
                {
                    var index = i;
                    if (settings.SkipThreeMip && i < 3)
                    {
                        continue;
                    }
                    else if (settings.SkipThreeMip && i == 3)
                    {
                        index = 0;
                    }
                        
                    var mipSize = mipLevelSizes[i];
                    var inputMipSize = index == 0 ? mip0SizeNOP : mipLevelSizes[i - 1];
                    var texId = depthMipId[i];
                    cmd.SetGlobalVector(HizShaderIds.InputScaleAndMaxIndexId, new Vector4(inputMipSize.x / (float)mipSize.x, inputMipSize.y / (float)mipSize.y, inputMipSize.x - 1, inputMipSize.y - 1));
                    cmd.GetTemporaryRT(texId, mipSize.x, mipSize.y, 0, FilterMode.Point, RenderTextureFormat.RFloat);
                    cmd.SetRenderTarget(texId, RenderBufferLoadAction.DontCare, RenderBufferStoreAction.Store);
                    cmd.DrawMesh(RenderingUtils.fullscreenMesh, Matrix4x4.identity, settings.CopyDepth, 0, 0);
                    cmd.SetGlobalTexture(HizShaderIds.DepthInputId, texId);
                }

                cmd.EndSample(sampleName);

                sampleName = "DepthCombine";
                cmd.BeginSample(sampleName);
                var hizIndex = GetHizIndex();
                RTHandle hizRt = depthPyramidTexs[hizIndex];
           
                RenderingUtils.ReAllocateIfNeeded(ref hizRt, new RenderTextureDescriptor()
                {
                    width = depthPyramidTextureSize.x,
                    height = depthPyramidTextureSize.y,
                    dimension = TextureDimension.Tex2D,
                    colorFormat = RenderTextureFormat.RFloat,
                    msaaSamples = 1,
                }, filterMode : FilterMode.Point, name: depthPyramidNames[hizIndex]);
                depthPyramidTexs[hizIndex] = hizRt;
                    
                cmd.SetRenderTarget(hizRt, RenderBufferLoadAction.DontCare, RenderBufferStoreAction.Store);
                for (int i = 1; i < mipLevelCount; i++)
                {
                    if (settings.SkipThreeMip && i < 3 + 1)
                    {
                        continue;
                    }
                    
                    var texId = depthMipId[i];
                    cmd.SetGlobalTexture(HizShaderIds.DepthInputId, texId);
                    var mipSize = mipLevelSizes[i];
                    var mipOffset = mipLevelOffsets[i];
                    cmd.SetViewport(new Rect(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y));
                    cmd.DrawMesh(RenderingUtils.fullscreenMesh, Matrix4x4.identity, settings.CopyDepth, 0, 1);
                }
                
                cmd.SetGlobalTexture(HizShaderIds.DepthPyramidTexId, hizRt);
                cmd.EndSample(sampleName);

                for (int i = 0; i < depthMipId.Length; i++)
                {
                    var texId = depthMipId[i];
                    cmd.ReleaseTemporaryRT(texId);
                }
            }

            private void DoComputeDepth(CommandBuffer cmd, Texture depthTex)
            {
                var hizIndex = GetHizIndex();
                RTHandle hizBuffer = depthPyramidTexs[hizIndex];
                cmd.BeginSample("Depth-Downsample");

                RenderingUtils.ReAllocateIfNeeded(ref hizBuffer, new RenderTextureDescriptor()
                {
                    width = depthPyramidTextureSize.x,
                    height = depthPyramidTextureSize.y,
                    dimension = TextureDimension.Tex2D,
                    colorFormat = RenderTextureFormat.RFloat,
                    msaaSamples = 1,
                    enableRandomWrite = true,
                    sRGB = false,
                }, filterMode: FilterMode.Point, name: depthPyramidNames[hizIndex]);
                depthPyramidTexs[hizIndex] = hizBuffer;

                int outputMipCnt = mipLevelCount - 1;
                int dispatchCnt = Mathf.CeilToInt(outputMipCnt / 4f);
                int mipCnt = outputMipCnt;
                Matrix4x4 matrix4X4 = new Matrix4x4();
                for (int i = 0; i <= dispatchCnt; i++) 
                {
                    int  startMip = i * 4;
                    int mip1 = startMip + 1;
                    int mip2 = startMip + 2;
                    int mip3 = startMip + 3;
                    int mip4 = startMip + 4;
                    var outputMipSize = mipLevelSizes[mip1];
                    if(outputMipSize.x == 0 || outputMipSize.y == 0)
                    {
                        break;
                    }
                    int kernelId = 0;
                    if (i == 0)
                    {
                        kernelId = 1;
                        if (settings.SkipThreeMip)
                        {
                            kernelId = 2;
                        }
                    }

                    cmd.SetComputeTextureParam(settings.ComputeShader, kernelId, HizShaderIds.DepthInputId, i == 0 ? depthTex : hizBuffer);
                    cmd.SetComputeTextureParam(settings.ComputeShader, kernelId, "_DepthMipChain", hizBuffer);

                    int inputMipIndex = startMip;
                    var inputMipOffset = mipLevelOffsets[inputMipIndex];
                    var inputMipSize = inputMipIndex == 0 ? mip0SizeNOP : mipLevelSizes[inputMipIndex];

                    cmd.SetComputeVectorParam(settings.ComputeShader, "_InputMipOffsetAndSize", new Vector4(inputMipOffset.x, inputMipOffset.y, inputMipSize.x, inputMipSize.y));
                    cmd.SetComputeIntParam(settings.ComputeShader, "_MipCount", Mathf.Min(mipCnt, 4));

                    ref var mipOffset = ref mipLevelOffsets[mip1];
                    ref var mipSize = ref mipLevelSizes[mip1];
                    matrix4X4.SetRow(0, new Vector4(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y));
                    mipOffset = ref mipLevelOffsets[mip2];
                    mipSize = ref mipLevelSizes[mip2];
                    matrix4X4.SetRow(1, new Vector4(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y));
                    mipOffset = ref mipLevelOffsets[mip3];
                    mipSize = ref mipLevelSizes[mip3];
                    matrix4X4.SetRow(2, new Vector4(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y));
                    mipOffset = ref mipLevelOffsets[mip4];
                    mipSize = ref mipLevelSizes[mip4];
                    matrix4X4.SetRow(3, new Vector4(mipOffset.x, mipOffset.y, mipSize.x, mipSize.y));

                    cmd.SetComputeMatrixParam(settings.ComputeShader, "_MipOffsetAndSizeArray", matrix4X4);

                    // XXX: 可以测试是否有正面作用
                    // cmd.SetExecutionFlags(CommandBufferExecutionFlags.AsyncCompute);
                    cmd.DispatchCompute(settings.ComputeShader, kernelId, Mathf.CeilToInt(outputMipSize.x / 8f), Mathf.CeilToInt(outputMipSize.y / 8f), 1);
                    mipCnt = mipCnt - 4;
                }
                cmd.EndSample("Depth-Downsample");

                cmd.SetGlobalTexture(HizShaderIds.DepthPyramidTexId, hizBuffer);
            }

            private void DoSpdDepth(CommandBuffer cmd, Texture depthTex)
            {
                var hizIndex = GetHizIndex();
                RTHandle hizBuffer = depthPyramidTexs[hizIndex];

                RenderingUtils.ReAllocateIfNeeded(ref hizBuffer, new RenderTextureDescriptor()
                {
                    width = depthPyramidTextureSize.x,
                    height = depthPyramidTextureSize.y,
                    dimension = TextureDimension.Tex2D,
                    colorFormat = RenderTextureFormat.RFloat,
                    msaaSamples = 1,
                    enableRandomWrite = true,
                }, filterMode: FilterMode.Point, name: depthPyramidNames[hizIndex]);
                depthPyramidTexs[hizIndex] = hizBuffer;

                var dispatchX = Mathf.CeilToInt(mipOffsetAndSizes[0].z / 64f);
                var dispatchY = Mathf.CeilToInt(mipOffsetAndSizes[0].w / 64f);
                cmd.SetComputeIntParam(settings.Spd, "mips", mipLevelCount);
                cmd.SetComputeIntParam(settings.Spd, "numWorkGroups", dispatchX * dispatchY);
                cmd.SetComputeVectorParam(settings.Spd, "inputTextureSize", new Vector4(mip0SizeNOP.x, mip0SizeNOP.y, 0, 0));
                

                cmd.SetComputeVectorArrayParam(settings.Spd, "_MipOffsetAndSizeArray", mipOffsetAndSizes);

                cmd.SetComputeTextureParam(settings.Spd, 0, "_InputDepth", depthTex);
                cmd.SetComputeTextureParam(settings.Spd, 0, "_OutDepth", hizBuffer);
                cmd.SetComputeTextureParam(settings.Spd, 0, "rw_internal_global_atomic", SpdAtomicCounter);
                
                cmd.DispatchCompute(settings.Spd, 0, dispatchX, dispatchY, 1);
                cmd.SetGlobalTexture(HizShaderIds.DepthPyramidTexId, hizBuffer);
            }

            private void DoCulling(ScriptableRenderContext context, ref RenderingData renderingData)
            {
                var cmd = renderingData.commandBuffer;;
                cmd.BeginSample("Hiz-Culling");
                var hizIndex = GetHizIndex();
                var cullResult = cullResults[hizIndex];
                var world2Project = renderingData.cameraData.GetGPUProjectionMatrix() *
                    renderingData.cameraData.GetViewMatrix();
                if (settings.UseCompute)
                {

                }
                else
                {
                    cmd.SetGlobalMatrix(HizShaderIds.GPUCullingVPId, world2Project);
                    cmd.SetViewProjectionMatrices(Matrix4x4.identity, Matrix4x4.identity);

                    cmd.SetRenderTarget(cullResult.ResultTex, RenderBufferLoadAction.DontCare, RenderBufferStoreAction.Store);
                    cmd.DrawMesh(RenderingUtils.fullscreenMesh, Matrix4x4.identity, settings.CullMat, 0, 0);
                    cmd.RequestAsyncReadback(cullResult.ResultTex, 0, cullResult.ReadBackAction);
                    //cmd.RequestAsyncReadbackIntoNativeArray(ref cullResult.ResultArray, cullResult.ResultTex, 0, cullResult.ReadBackAction);
                    cullResult.DataReady = false;
                }
                cmd.EndSample("Hiz-Culling");
            }

            public override void Execute(ScriptableRenderContext context, ref RenderingData renderingData)
            {
                var cmd = renderingData.commandBuffer;
                Texture depthTex = Shader.GetGlobalTexture(HizShaderIds.CameraDepthTextureId);
                if(depthTex == null)
                {
                    return;
                }
                using var scp = new ProfilingScope(cmd, profiler);
                ComputePackedMipChainInfo(new Vector2Int(depthTex.width, depthTex.height));

                switch (settings.PyramidFunc)
                {
                    case EDepthPyramidFunc.CopyDepth:
                        DoCopyDepth(cmd, depthTex);
                        break;
                    case EDepthPyramidFunc.Compute:
                        DoComputeDepth(cmd, depthTex);
                        break;
                    case EDepthPyramidFunc.SPD:
                        DoSpdDepth(cmd, depthTex);
                        break;
                    default:
                        break;
                }

                cmd.SetGlobalVector(HizShaderIds.Mip0SizeId, new Vector4(mip0SizeNOP.x, mip0SizeNOP.y, 0, 0));
                if (settings.SkipThreeMip)
                {
                    cmd.SetGlobalVector(HizShaderIds.MipmapLevelMinMaxIndexId, new Vector4(4, mipLevelCount - 1, 0, 0));
                }
                else
                {
                    cmd.SetGlobalVector(HizShaderIds.MipmapLevelMinMaxIndexId, new Vector4(1, mipLevelCount - 1, 0, 0));
                }
                cmd.SetGlobalVectorArray(HizShaderIds.MipOffsetAndSizeArrayId, mipOffsetAndSizes);
                DoCulling(context, ref renderingData);
            }



            private int GetLastCullIndex(int hizIndex)
            {
                if (!settings.UseThreeFrameReadback)
                {
                    return 0;
                }

                switch (hizIndex)
                {
                    case 0:
                        return 1;
                    case 1:
                        return 2;
                    case 2:
                        return 0;
                    default:
                        throw new ArgumentException("参数错误 hizIndex:" + hizIndex);
                }
            }


            public void Dispose()
            {
                if (SpdAtomicCounter != null)
                {
                    SpdAtomicCounter.Release();
                    SpdAtomicCounter = null;
                }

                for (int i = 0; i < depthPyramidTexs.Length; i++)
                {
                    var rt = depthPyramidTexs[i];
                    rt?.Release();
                }

                for (int i = 0; i < cullResults.Length; i++)
                {
                    cullResults[i]?.Dispose();
                }
            }
        }
    }
}