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

namespace FidelityFX.FrameGen
{
    public class FrameInterpolationContext
    {
        private FrameInterpolation.ContextDescription _contextDescription;

        private FrameInterpolationPass _reconstructAndDilatePass;
        private FrameInterpolationPass _setupPass;
        private FrameInterpolationPass _reconstructPreviousDepthPass;
        private FrameInterpolationPass _gameMotionVectorFieldPass;
        private FrameInterpolationPass _opticalFlowVectorFieldPass;
        private FrameInterpolationPass _disocclusionMaskPass;
        private FrameInterpolationPass _interpolationPass;
        private FrameInterpolationPass _inpaintingPyramidPass;
        private FrameInterpolationPass _inpaintingPass;
        private FrameInterpolationPass _gameVectorFieldInpaintingPyramidPass;
        private FrameInterpolationPass _debugViewPass;
        
        private readonly FrameInterpolationResources _resources = new FrameInterpolationResources();
        
        private ComputeBuffer _frameInterpolationConstantsBuffer;
        private readonly FrameInterpolation.Constants[] _frameInterpolationConstantsArray = { new FrameInterpolation.Constants() };
        private ref FrameInterpolation.Constants Constants => ref _frameInterpolationConstantsArray[0];
        
        private ComputeBuffer _spdConstantsBuffer;
        private readonly FrameInterpolation.InpaintingPyramidConstants[] _spdConstantsArray = { new FrameInterpolation.InpaintingPyramidConstants() };
        private ref FrameInterpolation.InpaintingPyramidConstants SpdConstants => ref _spdConstantsArray[0];
        
        private readonly CustomSampler _sampler = CustomSampler.Create("Frame Interpolation");
        private readonly CustomSampler _prepareSampler = CustomSampler.Create("Frame Interpolation - Prepare");

        private bool _firstExecution;
        private bool _asyncSupported;
        private ulong _previousFrameID;
        private ulong _dispatchCount;
        
        public void Create(in FrameInterpolation.ContextDescription contextDescription)
        {
            _contextDescription = contextDescription;

            _frameInterpolationConstantsBuffer = CreateConstantBuffer<FrameInterpolation.Constants>();
            _spdConstantsBuffer = CreateConstantBuffer<FrameInterpolation.InpaintingPyramidConstants>();

            _firstExecution = true;
            _asyncSupported = (_contextDescription.flags & FrameInterpolation.InitializationFlags.EnableAsyncSupport) == FrameInterpolation.InitializationFlags.EnableAsyncSupport;

            Constants.maxRenderSize = _contextDescription.maxRenderSize;
            Constants.displaySize = _contextDescription.displaySize;
            Constants.displaySizeRcp.x = 1.0f / _contextDescription.displaySize.x;
            Constants.displaySizeRcp.y = 1.0f / _contextDescription.displaySize.y;
            Constants.interpolationRectBase = Vector2Int.zero;
            Constants.interpolationRectSize = _contextDescription.displaySize;
            
            _resources.Create(_contextDescription);
            CreatePasses();
        }

        private void CreatePasses()
        {
            _reconstructAndDilatePass = new FrameInterpolationReconstructAndDilatePass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _setupPass = new FrameInterpolationSetupPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _reconstructPreviousDepthPass = new FrameInterpolationReconstructPreviousDepthPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _gameMotionVectorFieldPass = new FrameInterpolationGameMotionVectorFieldPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _opticalFlowVectorFieldPass = new FrameInterpolationOpticalFlowVectorFieldPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _disocclusionMaskPass = new FrameInterpolationDisocclusionMaskPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _interpolationPass = new FrameInterpolationInterpolationPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _inpaintingPyramidPass = new FrameInterpolationInpaintingPyramidPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _inpaintingPass = new FrameInterpolationInpaintingPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _gameVectorFieldInpaintingPyramidPass = new FrameInterpolationGameVectorFieldInpaintingPyramidPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
            _debugViewPass = new FrameInterpolationDebugViewPass(_contextDescription, _resources, _frameInterpolationConstantsBuffer);
        }

        public void Destroy()
        {
            DestroyPass(ref _debugViewPass);
            DestroyPass(ref _gameVectorFieldInpaintingPyramidPass);
            DestroyPass(ref _inpaintingPass);
            DestroyPass(ref _inpaintingPyramidPass);
            DestroyPass(ref _interpolationPass);
            DestroyPass(ref _disocclusionMaskPass);
            DestroyPass(ref _opticalFlowVectorFieldPass);
            DestroyPass(ref _gameMotionVectorFieldPass);
            DestroyPass(ref _reconstructPreviousDepthPass);
            DestroyPass(ref _setupPass);
            DestroyPass(ref _reconstructAndDilatePass);
            
            _resources.Destroy();
            
            DestroyConstantBuffer(ref _spdConstantsBuffer);
            DestroyConstantBuffer(ref _frameInterpolationConstantsBuffer);
        }

        public void Prepare(CommandBuffer commandBuffer, FrameInterpolation.PrepareDescription prepareDescription)
        {
            commandBuffer.BeginSample(_prepareSampler);
            
            commandBuffer.EndSample(_prepareSampler);
        }
        
        public void Dispatch(CommandBuffer commandBuffer, FrameInterpolation.DispatchDescription dispatchDescription)
        {
            commandBuffer.BeginSample(_sampler);
            
            commandBuffer.EndSample(_sampler);
        }
        
        private static ComputeBuffer CreateConstantBuffer<TConstants>() where TConstants: struct
        {
            return new ComputeBuffer(1, Marshal.SizeOf<TConstants>(), ComputeBufferType.Constant);
        }
        
        private static void DestroyConstantBuffer(ref ComputeBuffer bufferRef)
        {
            if (bufferRef == null)
                return;
            
            bufferRef.Release();
            bufferRef = null;
        }

        private static void DestroyPass(ref FrameInterpolationPass pass)
        {
            if (pass == null)
                return;
            
            pass.Dispose();
            pass = null;
        }
    }
}