using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using UnityEngine;
using UnityEngine.Rendering;

namespace FidelityFX
{
    public class Fsr2Context
    {
        private const int MaxQueuedFrames = 16;
        
        private Fsr2.ContextDescription _contextDescription;
        private CommandBuffer _commandBuffer;
        
        private Fsr2Pipeline _depthClipPipeline;
        private Fsr2Pipeline _reconstructPreviousDepthPipeline;
        private Fsr2Pipeline _lockPipeline;
        private Fsr2Pipeline _accumulatePipeline;
        private Fsr2Pipeline _accumulateSharpenPipeline;
        private Fsr2Pipeline _rcasPipeline;
        private Fsr2Pipeline _computeLuminancePyramidPipeline;
        private Fsr2Pipeline _generateReactivePipeline;

        private readonly Fsr2Resources _resources = new Fsr2Resources();

        private ComputeBuffer _fsr2ConstantsBuffer;
        private readonly Fsr2.Fsr2Constants[] _fsr2ConstantsArray = { new Fsr2.Fsr2Constants() };
        private ref Fsr2.Fsr2Constants Constants => ref _fsr2ConstantsArray[0];

        private ComputeBuffer _spdConstantsBuffer;
        private readonly Fsr2.SpdConstants[] _spdConstantsArray = { new Fsr2.SpdConstants() };
        private ref Fsr2.SpdConstants SpdConsts => ref _spdConstantsArray[0];
    
        private ComputeBuffer _rcasConstantsBuffer;
        private readonly Fsr2.RcasConstants[] _rcasConstantsArray = new Fsr2.RcasConstants[1];
        private ref Fsr2.RcasConstants RcasConsts => ref _rcasConstantsArray[0];

        private ComputeBuffer _generateReactiveConstantsBuffer;
        private readonly Fsr2.GenerateReactiveConstants[] _generateReactiveConstantsArray = { new Fsr2.GenerateReactiveConstants() };
        private ref Fsr2.GenerateReactiveConstants GenReactiveConsts => ref _generateReactiveConstantsArray[0];

        private bool _firstExecution;
        private Vector2 _previousJitterOffset;
        private int _resourceFrameIndex;

        public void Create(Fsr2.ContextDescription contextDescription)
        {
            _contextDescription = contextDescription;
            _commandBuffer = new CommandBuffer { name = "FSR2" };
            
            _fsr2ConstantsBuffer = CreateConstantBuffer<Fsr2.Fsr2Constants>();
            _spdConstantsBuffer = CreateConstantBuffer<Fsr2.SpdConstants>();
            _rcasConstantsBuffer = CreateConstantBuffer<Fsr2.RcasConstants>();
            _generateReactiveConstantsBuffer = CreateConstantBuffer<Fsr2.GenerateReactiveConstants>();

            // Set defaults
            _firstExecution = true;
            _resourceFrameIndex = 0;
            
            Constants.displaySize = _contextDescription.DisplaySize;
            
            _resources.Create(_contextDescription);
            CreatePipelines();
        }

        private void CreatePipelines()
        {
            _computeLuminancePyramidPipeline = new Fsr2ComputeLuminancePyramidPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer, _spdConstantsBuffer);
            _reconstructPreviousDepthPipeline = new Fsr2ReconstructPreviousDepthPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer);
            _depthClipPipeline = new Fsr2DepthClipPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer);
            _lockPipeline = new Fsr2LockPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer);
            _accumulatePipeline = new Fsr2AccumulatePipeline(_contextDescription, _resources, _fsr2ConstantsBuffer);
            _accumulateSharpenPipeline = new Fsr2AccumulateSharpenPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer);
            _rcasPipeline = new Fsr2RcasPipeline(_contextDescription, _resources, _fsr2ConstantsBuffer, _rcasConstantsBuffer);
            _generateReactivePipeline = new Fsr2GenerateReactivePipeline(_contextDescription, _resources, _generateReactiveConstantsBuffer);
        }
        
        public void Destroy()
        {
            DestroyPipeline(ref _generateReactivePipeline);
            DestroyPipeline(ref _computeLuminancePyramidPipeline);
            DestroyPipeline(ref _rcasPipeline);
            DestroyPipeline(ref _accumulateSharpenPipeline);
            DestroyPipeline(ref _accumulatePipeline);
            DestroyPipeline(ref _lockPipeline);
            DestroyPipeline(ref _reconstructPreviousDepthPipeline);
            DestroyPipeline(ref _depthClipPipeline);
            
            _resources.Destroy();
            
            DestroyConstantBuffer(ref _generateReactiveConstantsBuffer);
            DestroyConstantBuffer(ref _rcasConstantsBuffer);
            DestroyConstantBuffer(ref _spdConstantsBuffer);
            DestroyConstantBuffer(ref _fsr2ConstantsBuffer);

            _commandBuffer.Dispose();
            _commandBuffer = null;
        }

        public void Dispatch(Fsr2.DispatchDescription dispatchParams)
        {
            _commandBuffer.Clear();
            Dispatch(dispatchParams, _commandBuffer);
            Graphics.ExecuteCommandBuffer(_commandBuffer);
        }
        
        public void Dispatch(Fsr2.DispatchDescription dispatchParams, CommandBuffer commandBuffer)
        {
            // TODO: validation & debug checking
            
            if (_firstExecution)
            {
                commandBuffer.SetRenderTarget(_resources.LockStatus[0]);
                commandBuffer.ClearRenderTarget(false, true, Color.clear);
                commandBuffer.SetRenderTarget(_resources.LockStatus[1]);
                commandBuffer.ClearRenderTarget(false, true, Color.clear);
            }
            
            int frameIndex = _resourceFrameIndex % 2;
            bool resetAccumulation = dispatchParams.Reset || _firstExecution;
            _firstExecution = false;

            // TODO: auto-exposure flag
            if (dispatchParams.Exposure == null) dispatchParams.Exposure = _resources.DefaultExposure;
            if (dispatchParams.Reactive == null) dispatchParams.Reactive = _resources.DefaultReactive;
            Fsr2Pipeline.RegisterResources(commandBuffer, _contextDescription, dispatchParams);
            
            SetupConstants(dispatchParams, resetAccumulation);
            
            // Reactive mask bias
            const int threadGroupWorkRegionDim = 8;
            int dispatchSrcX = (Constants.renderSize.x + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            int dispatchSrcY = (Constants.renderSize.y + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            int dispatchDstX = (_contextDescription.DisplaySize.x + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            int dispatchDstY = (_contextDescription.DisplaySize.y + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;

            // Clear reconstructed depth for max depth store.
            if (resetAccumulation)
            {
                commandBuffer.SetRenderTarget(_resources.LockStatus[frameIndex ^ 1]);
                commandBuffer.ClearRenderTarget(false, true, Color.clear);
                
                commandBuffer.SetRenderTarget(_resources.InternalUpscaled[frameIndex ^ 1]);
                commandBuffer.ClearRenderTarget(false, true, Color.clear);
                
                commandBuffer.SetRenderTarget(_resources.SceneLuminance);
                commandBuffer.ClearRenderTarget(false, true, Color.clear);
                
                // Auto exposure always used to track luma changes in locking logic
                commandBuffer.SetRenderTarget(_resources.AutoExposure);
                commandBuffer.ClearRenderTarget(false, true, new Color(-1f, 1e8f, 0f, 0f));
            }
            
            // Auto exposure
            SetupSpdConstants(dispatchParams, out var dispatchThreadGroupCount);
            
            // Initialize constant buffers data
            _fsr2ConstantsBuffer.SetData(_fsr2ConstantsArray);
            _spdConstantsBuffer.SetData(_spdConstantsArray);
            
            // Compute luminance pyramid
            _computeLuminancePyramidPipeline.ScheduleDispatch(commandBuffer, dispatchParams, frameIndex, dispatchThreadGroupCount.x, dispatchThreadGroupCount.y);

            // Reconstruct previous depth
            _reconstructPreviousDepthPipeline.ScheduleDispatch(commandBuffer, dispatchParams, frameIndex, dispatchSrcX, dispatchSrcY);

            // Depth clip
            _depthClipPipeline.ScheduleDispatch(commandBuffer, dispatchParams, frameIndex, dispatchSrcX, dispatchSrcY);

            // Create locks
            _lockPipeline.ScheduleDispatch(commandBuffer, dispatchParams, frameIndex, dispatchSrcX, dispatchSrcY);

            bool sharpenEnabled = dispatchParams.EnableSharpening;
            
            // Accumulate
            var accumulatePipeline = sharpenEnabled ? _accumulateSharpenPipeline : _accumulatePipeline;
            accumulatePipeline.ScheduleDispatch(commandBuffer, dispatchParams, frameIndex, dispatchDstX, dispatchDstY);

            if (sharpenEnabled)
            {
                // Compute the constants
                SetupRcasConstants(dispatchParams);
                _rcasConstantsBuffer.SetData(_rcasConstantsArray);
                
                // Dispatch RCAS
                const int threadGroupWorkRegionDimRcas = 16;
                int threadGroupsX = (Screen.width + threadGroupWorkRegionDimRcas - 1) / threadGroupWorkRegionDimRcas;
                int threadGroupsY = (Screen.height + threadGroupWorkRegionDimRcas - 1) / threadGroupWorkRegionDimRcas;
                _rcasPipeline.ScheduleDispatch(commandBuffer, dispatchParams, frameIndex, threadGroupsX, threadGroupsY);
            }

            _resourceFrameIndex = (_resourceFrameIndex + 1) % MaxQueuedFrames;

            Fsr2Pipeline.UnregisterResources(commandBuffer);
        }

        public void GenerateReactiveMask(Fsr2.GenerateReactiveDescription dispatchParams)
        {
            _commandBuffer.Clear();
            GenerateReactiveMask(dispatchParams, _commandBuffer);
            Graphics.ExecuteCommandBuffer(_commandBuffer);
        }

        public void GenerateReactiveMask(Fsr2.GenerateReactiveDescription dispatchParams, CommandBuffer commandBuffer)
        {
            const int threadGroupWorkRegionDim = 8;
            int dispatchSrcX = (dispatchParams.RenderSize.x + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;
            int dispatchSrcY = (dispatchParams.RenderSize.y + (threadGroupWorkRegionDim - 1)) / threadGroupWorkRegionDim;

            GenReactiveConsts.scale = dispatchParams.Scale;
            GenReactiveConsts.threshold = dispatchParams.CutoffThreshold;
            GenReactiveConsts.binaryValue = dispatchParams.BinaryValue;
            GenReactiveConsts.flags = dispatchParams.Flags;
            _generateReactiveConstantsBuffer.SetData(_generateReactiveConstantsArray);
            
            ((Fsr2GenerateReactivePipeline)_generateReactivePipeline).ScheduleDispatch(_commandBuffer, dispatchParams, dispatchSrcX, dispatchSrcY);
        }

        private void SetupConstants(Fsr2.DispatchDescription dispatchParams, bool resetAccumulation)
        {
            ref Fsr2.Fsr2Constants constants = ref Constants;
            
            constants.jitterOffset = dispatchParams.JitterOffset;
            constants.renderSize = dispatchParams.RenderSize;
            constants.maxRenderSize = _contextDescription.MaxRenderSize;
            constants.inputColorResourceDimensions = dispatchParams.RenderSize; // We have no way to query the actual width & height of the input resources, so trust that it matches the render size
            // TODO: if we want to support dynamic resolution, we may have to differentiate between renderSize and inputColorResourceDimensions again
            // TODO: come to think of it, couldn't we use Unity's ScalableBufferManager for dealing with render scale altogether? So we don't have to deal with this nasty double camera business?

            // Compute the horizontal FOV for the shader from the vertical one
            float aspectRatio = (float)dispatchParams.RenderSize.x / dispatchParams.RenderSize.y;
            float cameraAngleHorizontal = Mathf.Atan(Mathf.Tan(dispatchParams.CameraFovAngleVertical / 2.0f) * aspectRatio) * 2.0f;
            constants.tanHalfFOV = Mathf.Tan(cameraAngleHorizontal * 0.5f);
            constants.viewSpaceToMetersFactor = (dispatchParams.ViewSpaceToMetersFactor > 0.0f) ? dispatchParams.ViewSpaceToMetersFactor : 1.0f;

            // Compute params to enable device depth to view space depth computation in shader
            constants.deviceToViewDepth = SetupDeviceDepthToViewSpaceDepthParams(dispatchParams);
            
            // To be updated if resource is larger than the actual image size
            constants.downscaleFactor = new Vector2((float)constants.renderSize.x / _contextDescription.DisplaySize.x, (float)constants.renderSize.y / _contextDescription.DisplaySize.y);
            constants.previousFramePreExposure = constants.preExposure;
            constants.preExposure = (dispatchParams.PreExposure != 0) ? dispatchParams.PreExposure : 1.0f;
            
            // Motion vector data
            Vector2Int motionVectorsTargetSize = (_contextDescription.Flags & Fsr2.InitializationFlags.EnableDisplayResolutionMotionVectors) != 0 ? constants.displaySize : constants.renderSize;
            constants.motionVectorScale = dispatchParams.MotionVectorScale / motionVectorsTargetSize;
            
            // Compute jitter cancellation
            if ((_contextDescription.Flags & Fsr2.InitializationFlags.EnableMotionVectorsJitterCancellation) != 0)
            {
                constants.motionVectorJitterCancellation = (_previousJitterOffset - constants.jitterOffset) / motionVectorsTargetSize;
                _previousJitterOffset = constants.jitterOffset;
            }

            int jitterPhaseCount = Fsr2.GetJitterPhaseCount(dispatchParams.RenderSize.x, _contextDescription.DisplaySize.x);
            if (resetAccumulation || constants.jitterPhaseCount == 0)
            {
                constants.jitterPhaseCount = jitterPhaseCount;
            }
            else
            {
                int jitterPhaseCountDelta = (int)(jitterPhaseCount - constants.jitterPhaseCount);
                if (jitterPhaseCountDelta > 0)
                    constants.jitterPhaseCount++;
                else if (jitterPhaseCountDelta < 0)
                    constants.jitterPhaseCount--;
            }
            
            // Convert delta time to seconds and clamp to [0, 1]
            constants.deltaTime = Mathf.Clamp01(dispatchParams.FrameTimeDelta);

            if (resetAccumulation)
                constants.frameIndex = 0;
            else
                constants.frameIndex++;

            // Shading change usage of the SPD mip levels
            constants.lumaMipLevelToUse = Fsr2Pipeline.ShadingChangeMipLevel;

            float mipDiv = 2 << constants.lumaMipLevelToUse;
            constants.lumaMipDimensions.x = (int)(constants.maxRenderSize.x / mipDiv);
            constants.lumaMipDimensions.y = (int)(constants.maxRenderSize.y / mipDiv);
        }
        
        private Vector4 SetupDeviceDepthToViewSpaceDepthParams(Fsr2.DispatchDescription dispatchParams)
        {
            bool inverted = (_contextDescription.Flags & Fsr2.InitializationFlags.EnableDepthInverted) != 0;
            bool infinite = (_contextDescription.Flags & Fsr2.InitializationFlags.EnableDepthInfinite) != 0;

            // make sure it has no impact if near and far plane values are swapped in dispatch params
            // the flags "inverted" and "infinite" will decide what transform to use
            float min = Mathf.Min(dispatchParams.CameraNear, dispatchParams.CameraFar);
            float max = Mathf.Max(dispatchParams.CameraNear, dispatchParams.CameraFar);

            if (inverted)
            {
                (min, max) = (max, min);
            }

            float q = max / (min - max);
            float d = -1.0f;

            Vector4 matrixElemC = new Vector4(q, -1.0f - Mathf.Epsilon, q, 0.0f + Mathf.Epsilon);
            Vector4 matrixElemE = new Vector4(q * min, -min - Mathf.Epsilon, q * min, max);
            
            // Revert x and y coords
            float aspect = (float)dispatchParams.RenderSize.x / dispatchParams.RenderSize.y;
            float cotHalfFovY = Mathf.Cos(0.5f * dispatchParams.CameraFovAngleVertical) / Mathf.Sin(0.5f * dispatchParams.CameraFovAngleVertical);

            int matrixIndex = (inverted ? 2 : 0) + (infinite ? 1 : 0);
            return new Vector4(
                d * matrixElemC[matrixIndex],
                matrixElemE[matrixIndex],
                aspect / cotHalfFovY,
                1.0f / cotHalfFovY);
        }

        private void SetupRcasConstants(Fsr2.DispatchDescription dispatchParams)
        {
            int sharpnessIndex = Mathf.RoundToInt(Mathf.Clamp01(dispatchParams.Sharpness) * (RcasConfigs.Count - 1));
            RcasConsts = RcasConfigs[sharpnessIndex];
        }

        private void SetupSpdConstants(Fsr2.DispatchDescription dispatchParams, out Vector2Int dispatchThreadGroupCount)
        {
            RectInt rectInfo = new RectInt(0, 0, dispatchParams.RenderSize.x, dispatchParams.RenderSize.y);
            SpdSetup(rectInfo, out dispatchThreadGroupCount, out var workGroupOffset, out var numWorkGroupsAndMips);
            
            // Downsample
            ref Fsr2.SpdConstants spdConstants = ref SpdConsts;
            spdConstants.numWorkGroups = (uint)numWorkGroupsAndMips.x;
            spdConstants.mips = (uint)numWorkGroupsAndMips.y;
            spdConstants.workGroupOffsetX = (uint)workGroupOffset.x;
            spdConstants.workGroupOffsetY = (uint)workGroupOffset.y;
            spdConstants.renderSizeX = (uint)dispatchParams.RenderSize.x;
            spdConstants.renderSizeY = (uint)dispatchParams.RenderSize.y;
        }

        private static void SpdSetup(RectInt rectInfo, 
            out Vector2Int dispatchThreadGroupCount, out Vector2Int workGroupOffset, out Vector2Int numWorkGroupsAndMips, int mips = -1)
        {
            workGroupOffset = new Vector2Int(rectInfo.x / 64, rectInfo.y / 64);

            int endIndexX = (rectInfo.x + rectInfo.width - 1) / 64;
            int endIndexY = (rectInfo.y + rectInfo.height - 1) / 64;

            dispatchThreadGroupCount = new Vector2Int(endIndexX + 1 - workGroupOffset.x, endIndexY + 1 - workGroupOffset.y);

            numWorkGroupsAndMips = new Vector2Int(dispatchThreadGroupCount.x * dispatchThreadGroupCount.y, mips);
            if (mips < 0)
            {
                float resolution = Math.Max(rectInfo.width, rectInfo.height);
                numWorkGroupsAndMips.y = Math.Min(Mathf.FloorToInt(Mathf.Log(resolution, 2.0f)), 12);
            }
        }

        /// <summary>
        /// The FSR2 C++ codebase uses floats bitwise converted to ints to pass sharpness parameters to the RCAS shader.
        /// This is not possible in C# without enabling unsafe code compilation, so to avoid that we instead use a table of precomputed values.
        /// </summary>
        private static readonly List<Fsr2.RcasConstants> RcasConfigs = new List<Fsr2.RcasConstants>()
        {
            new Fsr2.RcasConstants(1048576000u, 872428544u),
            new Fsr2.RcasConstants(1049178080u, 877212745u),
            new Fsr2.RcasConstants(1049823372u, 882390168u),
            new Fsr2.RcasConstants(1050514979u, 887895276u),
            new Fsr2.RcasConstants(1051256227u, 893859143u),
            new Fsr2.RcasConstants(1052050675u, 900216232u),
            new Fsr2.RcasConstants(1052902144u, 907032080u),
            new Fsr2.RcasConstants(1053814727u, 914306687u),
            new Fsr2.RcasConstants(1054792807u, 922105590u),
            new Fsr2.RcasConstants(1055841087u, 930494326u),
            new Fsr2.RcasConstants(1056964608u, 939538432u),
            new Fsr2.RcasConstants(1057566688u, 944322633u),
            new Fsr2.RcasConstants(1058211980u, 949500056u),
            new Fsr2.RcasConstants(1058903587u, 955005164u),
            new Fsr2.RcasConstants(1059644835u, 960969031u),
            new Fsr2.RcasConstants(1060439283u, 967326120u),
            new Fsr2.RcasConstants(1061290752u, 974141968u),
            new Fsr2.RcasConstants(1062203335u, 981416575u),
            new Fsr2.RcasConstants(1063181415u, 989215478u),
            new Fsr2.RcasConstants(1064229695u, 997604214u),
            new Fsr2.RcasConstants(1065353216u, 1006648320),
        };
        
        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 DestroyPipeline(ref Fsr2Pipeline pipeline)
        {
            if (pipeline == null)
                return;
            
            pipeline.Dispose();
            pipeline = null;
        }
    }
}