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

namespace FidelityFX
{
    public class Fsr2Context
    {
        private Fsr2.ContextDescription _contextDescription;
        
        private ComputeShader _rcasComputeShader;
        
        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];

        public void Create(Fsr2.ContextDescription contextDescription)
        {
            _contextDescription = contextDescription;
            
            if (_rcasComputeShader == null)
                _rcasComputeShader = _contextDescription.Callbacks.LoadComputeShader("FSR2/ffx_fsr2_rcas_pass");
            
            _fsr2ConstantsBuffer = new ComputeBuffer(1, Marshal.SizeOf<Fsr2Constants>(), ComputeBufferType.Constant);
            _spdConstantsBuffer = new ComputeBuffer(1, Marshal.SizeOf<SpdConstants>(), ComputeBufferType.Constant);
            _rcasConstantsBuffer = new ComputeBuffer(1, Marshal.SizeOf<RcasConstants>(), ComputeBufferType.Constant);
        }
        
        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 = _rcasComputeShader.FindKernel("CS");
                _rcasComputeShader.SetTexture(rcasKernel, "r_exposure", dispatchDescription.Exposure);
                _rcasComputeShader.SetTexture(rcasKernel, "r_rcas_input", dispatchDescription.Input);
                _rcasComputeShader.SetTexture(rcasKernel, "rw_upscaled_output", dispatchDescription.Output);
                _rcasComputeShader.SetConstantBuffer("cbFSR2", _fsr2ConstantsBuffer, 0, Marshal.SizeOf<Fsr2Constants>());
                _rcasComputeShader.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;
            
                _rcasComputeShader.Dispatch(rcasKernel, threadGroupsX, threadGroupsY, 1);
            }
            else
            {
                Graphics.Blit(dispatchDescription.Input, dispatchDescription.Output);
            }
        }

        public void Destroy()
        {
            if (_rcasConstantsBuffer != null)
            {
                _rcasConstantsBuffer.Release();
                _rcasConstantsBuffer = null;
            }

            if (_spdConstantsBuffer != null)
            {
                _spdConstantsBuffer.Release();
                _spdConstantsBuffer = null;
            }

            if (_fsr2ConstantsBuffer != null)
            {
                _fsr2ConstantsBuffer.Release();
                _fsr2ConstantsBuffer = null;
            }
        }
        
        [Serializable, StructLayout(LayoutKind.Sequential)]
        private struct Fsr2Constants
        {
            public Vector2Int renderSize;
            public Vector2Int displaySize;
            public uint lumaMipDimensionsX, lumaMipDimensionsY;
            public uint lumaMipLevelToUse;
            public uint frameIndex;
            public Vector2 displaySizeRcp;
            public Vector2 jitterOffset;
            public Vector4 deviceToViewDepth;
            public Vector2 depthClipUVScale;
            public Vector2 postLockStatusUVScale;
            public Vector2 reactiveMaskDimRcp;
            public Vector2 motionVectorScale;
            public Vector2 downscaleFactor;
            public float preExposure;
            public float tanHalfFOV;
            public Vector2 motionVectorJitterCancellation;
            public float jitterPhaseCount;
            public float lockInitialLifetime;
            public float lockTickDelta;
            public float deltaTime;
            public float dynamicResChangeFactor;
            public float lumaMipRcp;
        }
        
        [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 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),
        };
    }
}