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

namespace FidelityFX
{
    public class Fsr2Context
    {
        private Fsr2.ContextDescription _contextDescription;

        private ComputeShader _prepareInputColorShader;
        private ComputeShader _depthClipShader;
        private ComputeShader _reconstructPreviousDepthShader;
        private ComputeShader _lockShader;
        private ComputeShader _accumulateShader;
        private ComputeShader _generateReactiveShader;
        private ComputeShader _rcasShader;
        private ComputeShader _computeLuminancePyramidShader;
        private ComputeShader _tcrAutogenShader;
        
        private ComputeBuffer _fsr2ConstantsBuffer;
        private readonly Fsr2Constants[] _fsr2ConstantsArray = { new Fsr2Constants() };

        private ComputeBuffer _spdConstantsBuffer;
        private readonly SpdConstants[] _spdConstantsArray = { new SpdConstants() };
    
        private ComputeBuffer _rcasConstantsBuffer;
        private readonly RcasConstants[] _rcasConstantsArray = new RcasConstants[1];

        private ComputeBuffer _generateReactiveConstantsBuffer;
        private readonly GenerateReactiveConstants[] _generateReactiveConstantsArray = { new GenerateReactiveConstants() };

        public void Create(Fsr2.ContextDescription contextDescription)
        {
            _contextDescription = contextDescription;
            
            _fsr2ConstantsBuffer = CreateConstantBuffer<Fsr2Constants>();
            _spdConstantsBuffer = CreateConstantBuffer<SpdConstants>();
            _rcasConstantsBuffer = CreateConstantBuffer<RcasConstants>();

            // Set defaults
            _fsr2ConstantsArray[0].displaySize = _contextDescription.DisplaySize;
            
            // Generate the data for the LUT
            const uint lanczos2LutWidth = 128;
            short[] lanczos2Weights = new short[lanczos2LutWidth];
            for (uint currentLanczosWidthIndex = 0; currentLanczosWidthIndex < lanczos2LutWidth; ++currentLanczosWidthIndex)
            {
                float x = 2.0f * currentLanczosWidthIndex / (lanczos2LutWidth - 1);
                float y = Fsr2.Lanczos2(x);
                lanczos2Weights[currentLanczosWidthIndex] = (short)Mathf.Round(y * 32767.0f);
            }
            
            InitShaders();
            
            // TODO: create resources, i.e. render textures used for intermediate results.
            // Note that "aliasable" resources should be equivalent to GetTemporary render textures
            // UAVs *may* be an issue with the PS4 not handling simultaneous reading and writing to an RT properly
            // Unity does have Graphics.SetRandomWriteTarget for enabling UAV on ComputeBuffers or RTs
            // Unity doesn't do 1D textures so just default to Texture2D
        }

        private void InitShaders()
        {
            LoadComputeShader("FSR2/ffx_fsr2_compute_luminance_pyramid_pass", ref _computeLuminancePyramidShader);
            LoadComputeShader("FSR2/ffx_fsr2_rcas_pass", ref _rcasShader);
            LoadComputeShader("FSR2/ffx_fsr2_prepare_input_color_pass", ref _prepareInputColorShader);
            LoadComputeShader("FSR2/ffx_fsr2_depth_clip_pass", ref _depthClipShader);
            LoadComputeShader("FSR2/ffx_fsr2_reconstruct_previous_depth_pass", ref _reconstructPreviousDepthShader);
            LoadComputeShader("FSR2/ffx_fsr2_lock_pass", ref _lockShader);
            LoadComputeShader("FSR2/ffx_fsr2_accumulate_pass", ref _accumulateShader);
            LoadComputeShader("FSR2/ffx_fsr2_autogen_reactive_pass", ref _generateReactiveShader);
            LoadComputeShader("FSR2/ffx_fsr2_tcr_autogen_pass", ref _tcrAutogenShader);
        }

        public void Dispatch(Fsr2.DispatchDescription dispatchDescription)
        {
            _fsr2ConstantsArray[0].preExposure = dispatchDescription.PreExposure;
            _fsr2ConstantsBuffer.SetData(_fsr2ConstantsArray);

            if (dispatchDescription.EnableSharpening)
            {
                int sharpnessIndex = Mathf.RoundToInt(Mathf.Clamp01(dispatchDescription.Sharpness) * (RcasConfigs.Count - 1));
                _rcasConstantsArray[0] = RcasConfigs[sharpnessIndex];
                _rcasConstantsBuffer.SetData(_rcasConstantsArray);
            
                // Run the RCAS sharpening filter on the upscaled image
                int rcasKernel = _rcasShader.FindKernel("CS");
                _rcasShader.SetTexture(rcasKernel, "r_input_exposure", dispatchDescription.Exposure);
                _rcasShader.SetTexture(rcasKernel, "r_rcas_input", dispatchDescription.Input);
                _rcasShader.SetTexture(rcasKernel, "rw_upscaled_output", dispatchDescription.Output);
                _rcasShader.SetConstantBuffer("cbFSR2", _fsr2ConstantsBuffer, 0, Marshal.SizeOf<Fsr2Constants>());
                _rcasShader.SetConstantBuffer("cbRCAS", _rcasConstantsBuffer, 0, Marshal.SizeOf<RcasConstants>());

                const int threadGroupWorkRegionDimRcas = 16;
                int threadGroupsX = (Screen.width + threadGroupWorkRegionDimRcas - 1) / threadGroupWorkRegionDimRcas;
                int threadGroupsY = (Screen.height + threadGroupWorkRegionDimRcas - 1) / threadGroupWorkRegionDimRcas;
            
                _rcasShader.Dispatch(rcasKernel, threadGroupsX, threadGroupsY, 1);
            }
            else
            {
                Graphics.Blit(dispatchDescription.Input, dispatchDescription.Output);
            }
        }

        public void Destroy()
        {
            DestroyConstantBuffer(ref _rcasConstantsBuffer);
            DestroyConstantBuffer(ref _spdConstantsBuffer);
            DestroyConstantBuffer(ref _fsr2ConstantsBuffer);

            DestroyComputeShader(ref _tcrAutogenShader);
            DestroyComputeShader(ref _generateReactiveShader);
            DestroyComputeShader(ref _accumulateShader);
            DestroyComputeShader(ref _lockShader);
            DestroyComputeShader(ref _reconstructPreviousDepthShader);
            DestroyComputeShader(ref _depthClipShader);
            DestroyComputeShader(ref _prepareInputColorShader);
            DestroyComputeShader(ref _rcasShader);
            DestroyComputeShader(ref _computeLuminancePyramidShader);
        }

        [Serializable, StructLayout(LayoutKind.Sequential)]
        private struct Fsr2Constants
        {
            public Vector2Int renderSize;
            public Vector2Int maxRenderSize;
            public Vector2Int displaySize;
            public Vector2Int inputColorResourceDimensions;
            public Vector2Int lumaMipDimensions;
            public int lumaMipLevelToUse;
            public int frameIndex;
            
            public Vector4 deviceToViewDepth;
            public Vector2 jitterOffset;
            public Vector2 motionVectorScale;
            public Vector2 downscaleFactor;
            public Vector2 motionVectorJitterCancellation;
            public float preExposure;
            public float previousFramePreExposure;
            public float tanHalfFOV;
            public float jitterPhaseCount;
            public float deltaTime;
            public float dynamicResChangeFactor;
            public float viewSpaceToMetersFactor;
        }
        
        [Serializable, StructLayout(LayoutKind.Sequential)]
        private struct SpdConstants
        {
            public uint mips;
            public uint numWorkGroups;
            public uint workGroupOffsetX, workGroupOffsetY;
            public uint renderSizeX, renderSizeY;
        }

        [Serializable, StructLayout(LayoutKind.Sequential)]
        private struct GenerateReactiveConstants
        {
            public float scale;
            public float threshold;
            public float binaryValue;
            public uint flags;
        }
        
        [Serializable, StructLayout(LayoutKind.Sequential)]
        private struct RcasConstants
        {
            public RcasConstants(uint sharpness, uint halfSharp)
            {
                this.sharpness = sharpness;
                this.halfSharp = halfSharp;
                dummy0 = dummy1 = 0;
            }
        
            public readonly uint sharpness;
            public readonly uint halfSharp;
            public readonly uint dummy0;
            public readonly uint dummy1;
        }

        private static readonly List<RcasConstants> RcasConfigs = new()
        {
            new(1048576000u, 872428544u),
            new(1049178080u, 877212745u),
            new(1049823372u, 882390168u),
            new(1050514979u, 887895276u),
            new(1051256227u, 893859143u),
            new(1052050675u, 900216232u),
            new(1052902144u, 907032080u),
            new(1053814727u, 914306687u),
            new(1054792807u, 922105590u),
            new(1055841087u, 930494326u),
            new(1056964608u, 939538432u),
            new(1057566688u, 944322633u),
            new(1058211980u, 949500056u),
            new(1058903587u, 955005164u),
            new(1059644835u, 960969031u),
            new(1060439283u, 967326120u),
            new(1061290752u, 974141968u),
            new(1062203335u, 981416575u),
            new(1063181415u, 989215478u),
            new(1064229695u, 997604214u),
            new(1065353216u, 1006648320),
        };
        
        private static ComputeBuffer CreateConstantBuffer<TConstants>() where TConstants: struct
        {
            return new ComputeBuffer(1, Marshal.SizeOf<TConstants>(), ComputeBufferType.Constant);
        }
        
        private void LoadComputeShader(string name, ref ComputeShader shaderRef)
        {
            if (shaderRef == null)
                shaderRef = _contextDescription.Callbacks.LoadComputeShader(name);
        }
        
        private static void DestroyConstantBuffer(ref ComputeBuffer bufferRef)
        {
            if (bufferRef == null)
                return;
            
            bufferRef.Release();
            bufferRef = null;
        }

        private void DestroyComputeShader(ref ComputeShader shaderRef)
        {
            if (shaderRef == null)
                return;

            _contextDescription.Callbacks.UnloadComputeShader(shaderRef);
            shaderRef = null;
        }
    }
}