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

namespace FidelityFX.FrameGen
{
    internal abstract class FrameInterpolationPass: IDisposable
    {
        protected readonly FrameInterpolation.ContextDescription ContextDescription;
        protected readonly FrameInterpolationResources Resources;
        protected readonly ComputeBuffer Constants;
        
        protected ComputeShader ComputeShader;
        protected int KernelIndex;
        
        protected CustomSampler Sampler;

        protected bool AsyncSupported => (ContextDescription.flags & FrameInterpolation.InitializationFlags.EnableAsyncSupport) == FrameInterpolation.InitializationFlags.EnableAsyncSupport;
        
        protected FrameInterpolationPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants)
        {
            ContextDescription = contextDescription;
            Resources = resources;
            Constants = constants;
        }

        public virtual void Dispose()
        {
        }
        
        public void ScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ = 1)
        {
            commandBuffer.BeginSample(Sampler);
            DoScheduleDispatch(commandBuffer, dispatchParams, frameIndex, dispatchX, dispatchY, dispatchZ);
            commandBuffer.EndSample(Sampler);
        }

        protected abstract void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ);

        protected void InitComputeShader(string passName, ComputeShader shader)
        {
            InitComputeShader(passName, shader, ContextDescription.flags);
        }
        
        private void InitComputeShader(string passName, ComputeShader shader, FrameInterpolation.InitializationFlags flags)
        {
            if (shader == null)
            {
                throw new MissingReferenceException($"Shader for Frame Interpolation pass '{passName}' could not be loaded! Please ensure it is included in the project correctly.");
            }

            ComputeShader = shader;
            KernelIndex = ComputeShader.FindKernel("CS");
            Sampler = CustomSampler.Create(passName);
            
            if ((flags & FrameInterpolation.InitializationFlags.EnableDisplayResolutionMotionVectors) == 0) ComputeShader.EnableKeyword("FFX_FRAMEINTERPOLATION_OPTION_LOW_RES_MOTION_VECTORS");
            if ((flags & FrameInterpolation.InitializationFlags.EnableJitterMotionVectors) != 0) ComputeShader.EnableKeyword("FFX_FRAMEINTERPOLATION_OPTION_JITTERED_MOTION_VECTORS");
            if ((flags & FrameInterpolation.InitializationFlags.EnableDepthInverted) != 0) ComputeShader.EnableKeyword("FFX_FRAMEINTERPOLATION_OPTION_INVERTED_DEPTH");
        }
    }

    internal class FrameInterpolationReconstructAndDilatePass : FrameInterpolationPass
    {
        public FrameInterpolationReconstructAndDilatePass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Reconstruct and Dilate", contextDescription.shaders.reconstructAndDilate);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
        }

        public void ScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.PrepareDescription prepareParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            int doubleBufferId = AsyncSupported ? frameIndex : 0;
            
            commandBuffer.BeginSample(Sampler);

            ref var mvs = ref prepareParams.motionVectors;
            ref var depth = ref prepareParams.depth;
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvInputMotionVectors, mvs.RenderTarget, mvs.MipLevel, mvs.SubElement);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvInputDepth, depth.RenderTarget, depth.MipLevel, depth.SubElement);

            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavReconstructedDepthPreviousFrame, Resources.ReconstructedDepth[doubleBufferId]);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavDilatedMotionVectors, Resources.DilatedMotionVectors[doubleBufferId]);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavDilatedDepth, Resources.DilatedDepth[doubleBufferId]);
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
            
            commandBuffer.EndSample(Sampler);
        }
    }
    
    internal class FrameInterpolationSetupPass : FrameInterpolationPass
    {
        public FrameInterpolationSetupPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Setup", contextDescription.shaders.setup);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            ref var scd = ref dispatchParams.opticalFlowSceneChangeDetection;
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvOpticalFlowSceneChangeDetection, scd.RenderTarget, scd.MipLevel, scd.SubElement);

            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavGameMotionVectorFieldX, Resources.GameMotionVectorFieldX);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavGameMotionVectorFieldY, Resources.GameMotionVectorFieldY);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavOpticalFlowMotionVectorFieldX, Resources.OpticalFlowMotionVectorFieldX);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavOpticalFlowMotionVectorFieldY, Resources.OpticalFlowMotionVectorFieldY);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavDisocclusionMask, Resources.DisocclusionMask);
            commandBuffer.SetComputeBufferParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavCounters, Resources.Counters);
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
        }
    }
    
    internal class FrameInterpolationReconstructPreviousDepthPass : FrameInterpolationPass
    {
        public FrameInterpolationReconstructPreviousDepthPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Reconstruct Previous Depth", contextDescription.shaders.reconstructPreviousDepth);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            int doubleBufferId = AsyncSupported ? frameIndex : 0;
            ref var backBuf = ref dispatchParams.currentBackBuffer;
            ref var bbNoHud = ref dispatchParams.currentBackBuffer_HUDLess;
            
            // TODO: verify that we need the buffers from *this* frame (probably yes)
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvDilatedMotionVectors, Resources.DilatedMotionVectors[doubleBufferId]);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvDilatedDepth, Resources.DilatedDepth[doubleBufferId]);
            if (bbNoHud.IsValid)
                commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvCurrentInterpolationSource, bbNoHud.RenderTarget, bbNoHud.MipLevel, bbNoHud.SubElement);
            else
                commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvCurrentInterpolationSource, backBuf.RenderTarget, backBuf.MipLevel, backBuf.SubElement);

            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavReconstructedDepthInterpolatedFrame, Resources.ReconstructedDepthInterpolatedFrame);
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
        }
    }
    
    internal class FrameInterpolationGameMotionVectorFieldPass : FrameInterpolationPass
    {
        public FrameInterpolationGameMotionVectorFieldPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Game Motion Vector Field", contextDescription.shaders.gameMotionVectorField);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            int doubleBufferId = AsyncSupported ? frameIndex : 0;
            ref var backBuf = ref dispatchParams.currentBackBuffer;
            ref var bbNoHud = ref dispatchParams.currentBackBuffer_HUDLess;

            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvDilatedMotionVectors, Resources.DilatedMotionVectors[doubleBufferId]);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvDilatedDepth, Resources.DilatedDepth[doubleBufferId]);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvPreviousInterpolationSource, Resources.PreviousInterpolationSource);
            if (bbNoHud.IsValid)
                commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvCurrentInterpolationSource, bbNoHud.RenderTarget, bbNoHud.MipLevel, bbNoHud.SubElement);
            else
                commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvCurrentInterpolationSource, backBuf.RenderTarget, backBuf.MipLevel, backBuf.SubElement);

            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavGameMotionVectorFieldX, Resources.GameMotionVectorFieldX);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavGameMotionVectorFieldY, Resources.GameMotionVectorFieldY);
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
        }
    }
    
    internal class FrameInterpolationOpticalFlowVectorFieldPass : FrameInterpolationPass
    {
        public FrameInterpolationOpticalFlowVectorFieldPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Optical Flow Vector Field", contextDescription.shaders.opticalFlowVectorField);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            int doubleBufferId = AsyncSupported ? frameIndex : 0;
            ref var ofVector = ref dispatchParams.opticalFlowVector;
            ref var backBuf = ref dispatchParams.currentBackBuffer;
            ref var bbNoHud = ref dispatchParams.currentBackBuffer_HUDLess;

            if (dispatchParams.opticalFlowScale.x > 0f)
                commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvOpticalFlowVector, ofVector.RenderTarget, ofVector.MipLevel, ofVector.SubElement);
            else
                commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvOpticalFlowVector, BuiltinRenderTextureType.None);

            // TODO this might error... if so, bind an empty placeholder resource
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvOpticalFlowConfidence, BuiltinRenderTextureType.None);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvDilatedDepth, Resources.DilatedDepth[doubleBufferId]);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvPreviousInterpolationSource, Resources.PreviousInterpolationSource);
            if (bbNoHud.IsValid)
                commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvCurrentInterpolationSource, bbNoHud.RenderTarget, bbNoHud.MipLevel, bbNoHud.SubElement);
            else
                commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvCurrentInterpolationSource, backBuf.RenderTarget, backBuf.MipLevel, backBuf.SubElement);

            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavOpticalFlowMotionVectorFieldX, Resources.OpticalFlowMotionVectorFieldX);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavOpticalFlowMotionVectorFieldY, Resources.OpticalFlowMotionVectorFieldY);
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
        }
    }
    
    internal class FrameInterpolationDisocclusionMaskPass : FrameInterpolationPass
    {
        public FrameInterpolationDisocclusionMaskPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Disocclusion Mask", contextDescription.shaders.disocclusionMask);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            int doubleBufferId = AsyncSupported ? frameIndex : 0;
            
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvGameMotionVectorFieldX, Resources.GameMotionVectorFieldX);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvGameMotionVectorFieldY, Resources.GameMotionVectorFieldY);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvReconstructedDepthPreviousFrame, Resources.ReconstructedDepth[doubleBufferId]);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvDilatedDepth, Resources.DilatedDepth[doubleBufferId]);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvReconstructedDepthInterpolatedFrame, Resources.ReconstructedDepthInterpolatedFrame);
            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.SrvInpaintingPyramid, Resources.InpaintingPyramid);

            commandBuffer.SetComputeTextureParam(ComputeShader, KernelIndex, FrameInterpolationShaderIDs.UavDisocclusionMask, Resources.DisocclusionMask);
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
        }
    }
    
    internal class FrameInterpolationInterpolationPass : FrameInterpolationPass
    {
        public FrameInterpolationInterpolationPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Interpolation", contextDescription.shaders.interpolation);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            throw new NotImplementedException();
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
        }
    }
    
    internal class FrameInterpolationInpaintingPyramidPass : FrameInterpolationPass
    {
        public FrameInterpolationInpaintingPyramidPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Inpainting Pyramid", contextDescription.shaders.inpaintingPyramid);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            throw new NotImplementedException();
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
        }
    }
    
    internal class FrameInterpolationInpaintingPass : FrameInterpolationPass
    {
        public FrameInterpolationInpaintingPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Inpainting", contextDescription.shaders.inpainting);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            throw new NotImplementedException();
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
        }
    }
    
    internal class FrameInterpolationGameVectorFieldInpaintingPyramidPass : FrameInterpolationPass
    {
        public FrameInterpolationGameVectorFieldInpaintingPyramidPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Game Vector Field Inpainting Pyramid", contextDescription.shaders.gameVectorFieldInpaintingPyramid);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            throw new NotImplementedException();
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
        }
    }
    
    internal class FrameInterpolationDebugViewPass : FrameInterpolationPass
    {
        public FrameInterpolationDebugViewPass(FrameInterpolation.ContextDescription contextDescription, FrameInterpolationResources resources, ComputeBuffer constants) 
            : base(contextDescription, resources, constants)
        {
            InitComputeShader("Debug View", contextDescription.shaders.debugView);
        }

        protected override void DoScheduleDispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchParams, int frameIndex, int dispatchX, int dispatchY, int dispatchZ)
        {
            throw new NotImplementedException();
            
            commandBuffer.SetComputeConstantBufferParam(ComputeShader, FrameInterpolationShaderIDs.CbFrameInterpolation, Constants, 0, Marshal.SizeOf<FrameInterpolation.Constants>());
            
            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, dispatchZ);
        }
    }
}