using System;
using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;

namespace ArmASR
{
    /// <summary>
    /// Helper class for bundling and managing persistent resources required by the ASR process.
    /// This includes lookup tables, default fallback resources and double-buffered resources that get swapped between frames.
    /// </summary>
    internal class AsrResources
    {
        public Texture2D DefaultExposure;
        public Texture2D DefaultReactive;
        public Texture2D LanczosLut;
        public Texture2D MaximumBiasLut;
        public RenderTexture SpdAtomicCounter;
        public RenderTexture SceneLuminance;
        public readonly RenderTexture[] AutoExposure = new RenderTexture[2];
        public readonly RenderTexture[] DilatedMotionVectors = new RenderTexture[2];
        public readonly RenderTexture[] LockStatus = new RenderTexture[2];
        public readonly RenderTexture[] InternalUpscaled = new RenderTexture[2];
        public readonly RenderTexture[] InternalReactive = new RenderTexture[2];
        public readonly RenderTexture[] LumaHistory = new RenderTexture[2];

        public void Create(in Asr.ContextDescription contextDescription)
        {
            // Generate the data for the LUT
            const int lanczos2LutWidth = 128;
            float[] lanczos2Weights = new float[lanczos2LutWidth];
            for (int currentLanczosWidthIndex = 0; currentLanczosWidthIndex < lanczos2LutWidth; ++currentLanczosWidthIndex)
            {
                float x = 2.0f * currentLanczosWidthIndex / (lanczos2LutWidth - 1);
                float y = Asr.Lanczos2(x);
                lanczos2Weights[currentLanczosWidthIndex] = y;
            }

            float[] maximumBias = new float[MaximumBiasTextureWidth * MaximumBiasTextureHeight];
            for (int i = 0; i < maximumBias.Length; ++i)
            {
                maximumBias[i] = MaximumBias[i] / 2.0f;
            }

            GetFormatRequirements(contextDescription, out bool isBalancedOrPerformance, out _, out _, out GraphicsFormat r16Format, out GraphicsFormat rg16Format);

            // Resource FSR2_LanczosLutData: FFX_RESOURCE_USAGE_READ_ONLY, FFX_SURFACE_FORMAT_R16_SNORM, FFX_RESOURCE_FLAGS_NONE
            // R16_SNorm textures are not supported by Unity on most platforms, strangely enough. So instead we use R32_SFloat and upload pre-normalized float data.
            LanczosLut = new Texture2D(lanczos2LutWidth, 1, GraphicsFormat.R32_SFloat, TextureCreationFlags.None) { name = "ASR_LanczosLutData" };
            LanczosLut.SetPixelData(lanczos2Weights, 0);
            LanczosLut.Apply();
            
            // Resource FSR2_MaximumUpsampleBias: FFX_RESOURCE_USAGE_READ_ONLY, FFX_SURFACE_FORMAT_R16_SNORM, FFX_RESOURCE_FLAGS_NONE
            MaximumBiasLut = new Texture2D(MaximumBiasTextureWidth, MaximumBiasTextureHeight, GraphicsFormat.R32_SFloat, TextureCreationFlags.None) { name = "ASR_MaximumUpsampleBias" };
            MaximumBiasLut.SetPixelData(maximumBias, 0);
            MaximumBiasLut.Apply();
            
            // Resource FSR2_DefaultExposure: FFX_RESOURCE_USAGE_READ_ONLY, FFX_SURFACE_FORMAT_R32G32_FLOAT, FFX_RESOURCE_FLAGS_NONE
            DefaultExposure = new Texture2D(1, 1, GraphicsFormat.R32G32_SFloat, TextureCreationFlags.None) { name = "ASR_DefaultExposure" };
            DefaultExposure.SetPixel(0, 0, Color.clear);
            DefaultExposure.Apply();

            // Resource FSR2_DefaultReactivityMask: FFX_RESOURCE_USAGE_READ_ONLY, FFX_SURFACE_FORMAT_R8_UNORM, FFX_RESOURCE_FLAGS_NONE
            DefaultReactive = new Texture2D(1, 1, GraphicsFormat.R8_UNorm, TextureCreationFlags.None) { name = "ASR_DefaultReactivityMask" };
            DefaultReactive.SetPixel(0, 0, Color.clear);
            DefaultReactive.Apply();
            
            // Resource FSR2_SpdAtomicCounter: FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R32_UINT, FFX_RESOURCE_FLAGS_ALIASABLE
            // Despite what the original FSR2 codebase says, this resource really isn't aliasable. Resetting this counter to 0 every frame breaks auto-exposure on MacOS Metal.
            SpdAtomicCounter = new RenderTexture(1, 1, 0, GraphicsFormat.R32_UInt) { name = "ASR_SpdAtomicCounter", enableRandomWrite = true };
            SpdAtomicCounter.Create();

            // Resource FSR2_ExposureMips: FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R16_FLOAT, FFX_RESOURCE_FLAGS_ALIASABLE
            // This is a rather special case: it's an aliasable resource, but because we require a mipmap chain and bind specific mip levels per shader, we can't easily use temporary RTs for this.
            int w = contextDescription.MaxRenderSize.x / 2, h = contextDescription.MaxRenderSize.y / 2;
            int mipCount = 1 + Mathf.FloorToInt(Mathf.Log(Math.Max(w, h), 2.0f));
            SceneLuminance = new RenderTexture(w, h, 0, r16Format, mipCount) { name = "ASR_ExposureMips", enableRandomWrite = true, useMipMap = true, autoGenerateMips = false };
            SceneLuminance.Create();
            
            // Resource FSR2_AutoExposure: FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R32G32_FLOAT, FFX_RESOURCE_FLAGS_NONE
            CreateDoubleBufferedResource(AutoExposure, "ASR_AutoExposure", Vector2Int.one, rg16Format, enableRandomWrite: true);
            
            // Resources FSR2_InternalDilatedVelocity1/2: FFX_RESOURCE_USAGE_RENDERTARGET | FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R16G16_FLOAT, FFX_RESOURCE_FLAGS_NONE
            CreateDoubleBufferedResource(DilatedMotionVectors, "ASR_InternalDilatedVelocity", contextDescription.MaxRenderSize, GraphicsFormat.R16G16_SFloat);
            
            // Resources FSR2_LockStatus1/2: FFX_RESOURCE_USAGE_RENDERTARGET | FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R16G16_FLOAT, FFX_RESOURCE_FLAGS_NONE
            CreateDoubleBufferedResource(LockStatus, "ASR_LockStatus", contextDescription.DisplaySize, GraphicsFormat.R16G16_SFloat);
            
            // Resources FSR2_InternalUpscaled1/2: FFX_RESOURCE_USAGE_RENDERTARGET | FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R16G16B16A16_FLOAT, FFX_RESOURCE_FLAGS_NONE
            CreateDoubleBufferedResource(InternalUpscaled, "ASR_InternalUpscaled", contextDescription.DisplaySize, !isBalancedOrPerformance ? GraphicsFormat.R16G16B16A16_SFloat : GraphicsFormat.B10G11R11_UFloatPack32);

            // Additional textures used by either balanced or performance presets
            if (isBalancedOrPerformance)
            {
                // Resources FSR2_InternalReactive1/2: FFXM_RESOURCE_USAGE_RENDERTARGET, FFXM_SURFACE_FORMAT_R8_SNORM, FFXM_RESOURCE_FLAGS_NONE
                CreateDoubleBufferedResource(InternalReactive, "ASR_InternalReactive", contextDescription.DisplaySize, GraphicsFormat.R8_SNorm);
            }
            else    // Quality preset specific
            {
                // Resources FSR2_LumaHistory1/2: FFX_RESOURCE_USAGE_RENDERTARGET | FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R8G8B8A8_UNORM, FFX_RESOURCE_FLAGS_NONE
                CreateDoubleBufferedResource(LumaHistory, "ASR_LumaHistory", contextDescription.DisplaySize, GraphicsFormat.R8G8B8A8_UNorm);
            }
        }
        
        // Set up shared aliasable resources, i.e. temporary render textures
        // These do not need to persist between frames, but they do need to be available between passes
        public static void CreateAliasableResources(CommandBuffer commandBuffer, in Asr.ContextDescription contextDescription, in Asr.DispatchDescription dispatchParams)
        {
            Vector2Int displaySize = contextDescription.DisplaySize;
            Vector2Int maxRenderSize = contextDescription.MaxRenderSize;
            
            GetFormatRequirements(contextDescription, out _, out bool preparedInputColorNeedsFp16, out GraphicsFormat r8Format, out _, out _);

            // FSR2_ReconstructedPrevNearestDepth: FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R32_UINT, FFX_RESOURCE_FLAGS_ALIASABLE
            commandBuffer.GetTemporaryRT(AsrShaderIDs.UavReconstructedPrevNearestDepth, maxRenderSize.x, maxRenderSize.y, 0, default, GraphicsFormat.R32_UInt, 1, true);

            // FSR2_DilatedDepth: FFX_RESOURCE_USAGE_RENDERTARGET | FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R32_FLOAT, FFX_RESOURCE_FLAGS_ALIASABLE
            commandBuffer.GetTemporaryRT(AsrShaderIDs.UavDilatedDepth, maxRenderSize.x, maxRenderSize.y, 0, default, GraphicsFormat.R32_SFloat, 1);

            // FSR2_LockInputLuma: FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R16_FLOAT, FFX_RESOURCE_FLAGS_ALIASABLE
            commandBuffer.GetTemporaryRT(AsrShaderIDs.UavLockInputLuma, maxRenderSize.x, maxRenderSize.y, 0, default, GraphicsFormat.R16_SFloat, 1);
            
            // FSR2_DilatedReactiveMasks: FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R8G8_UNORM, FFX_RESOURCE_FLAGS_ALIASABLE
            commandBuffer.GetTemporaryRT(AsrShaderIDs.UavDilatedReactiveMasks, maxRenderSize.x, maxRenderSize.y, 0, default, GraphicsFormat.R8G8_UNorm, 1);
            
            // FSR2_PreparedInputColor: FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R16G16B16A16_FLOAT, FFX_RESOURCE_FLAGS_ALIASABLE
            commandBuffer.GetTemporaryRT(AsrShaderIDs.UavPreparedInputColor, maxRenderSize.x, maxRenderSize.y, 0, default, preparedInputColorNeedsFp16 ? GraphicsFormat.R16G16B16A16_SFloat : GraphicsFormat.R8G8B8A8_UNorm, 1);
            
            // FSR2_NewLocks: FFX_RESOURCE_USAGE_UAV, FFX_SURFACE_FORMAT_R8_UNORM, FFX_RESOURCE_FLAGS_ALIASABLE
            commandBuffer.GetTemporaryRT(AsrShaderIDs.UavNewLocks, displaySize.x, displaySize.y, 0, default, r8Format, 1, true);
        }

        private static void GetFormatRequirements(in Asr.ContextDescription contextDescription, 
            out bool isBalancedOrPerformance, out bool preparedInputColorNeedsFP16,
            out GraphicsFormat r8Format, out GraphicsFormat r16Format, out GraphicsFormat rg16Format)
        {
            bool applyPerfModeOptimizations = contextDescription.Variant == Asr.Variant.Performance;
            bool applyBalancedModeOptimizations = contextDescription.Variant == Asr.Variant.Balanced;
            isBalancedOrPerformance = applyBalancedModeOptimizations || applyPerfModeOptimizations;
            preparedInputColorNeedsFP16 = !applyPerfModeOptimizations;

            // OpenGLES 3.2 specific: We need to work around some GLES limitations for some resources.
            bool isOpenGLES = SystemInfo.graphicsDeviceType == GraphicsDeviceType.OpenGLES3;
            r8Format = isOpenGLES ? GraphicsFormat.R32_SFloat : GraphicsFormat.R8_UNorm;
            r16Format = isOpenGLES ? GraphicsFormat.R32_SFloat : GraphicsFormat.R16_SFloat;
            rg16Format = isOpenGLES ? GraphicsFormat.R16G16B16A16_SFloat : GraphicsFormat.R16G16_SFloat;
        }

        public static void DestroyAliasableResources(CommandBuffer commandBuffer)
        {
            // Release all the aliasable resources used this frame
            commandBuffer.ReleaseTemporaryRT(AsrShaderIDs.UavReconstructedPrevNearestDepth);
            commandBuffer.ReleaseTemporaryRT(AsrShaderIDs.UavDilatedDepth);
            commandBuffer.ReleaseTemporaryRT(AsrShaderIDs.UavLockInputLuma);
            commandBuffer.ReleaseTemporaryRT(AsrShaderIDs.UavDilatedReactiveMasks);
            commandBuffer.ReleaseTemporaryRT(AsrShaderIDs.UavPreparedInputColor);
            commandBuffer.ReleaseTemporaryRT(AsrShaderIDs.UavNewLocks);
        }

        private static void CreateDoubleBufferedResource(RenderTexture[] resource, string name, Vector2Int size, GraphicsFormat format, bool enableRandomWrite = false)
        {
            for (int i = 0; i < 2; ++i)
            {
                resource[i] = new RenderTexture(size.x, size.y, 0, format) { name = name + (i + 1), enableRandomWrite = enableRandomWrite };
                resource[i].Create();
            }
        }

        public void Destroy()
        {
            DestroyResource(LumaHistory);
            DestroyResource(InternalReactive);
            DestroyResource(InternalUpscaled);
            DestroyResource(LockStatus);
            DestroyResource(DilatedMotionVectors);
            DestroyResource(AutoExposure);
            DestroyResource(ref SceneLuminance);
            DestroyResource(ref DefaultReactive);
            DestroyResource(ref DefaultExposure);
            DestroyResource(ref MaximumBiasLut);
            DestroyResource(ref LanczosLut);
        }
        
        private static void DestroyResource(ref Texture2D resource)
        {
            Asr.DestroyObject(resource);
            resource = null;
        }

        private static void DestroyResource(ref RenderTexture resource)
        {
            if (resource == null)
                return;
            
            resource.Release();
            resource = null;
        }

        private static void DestroyResource(RenderTexture[] resource)
        {
            for (int i = 0; i < resource.Length; ++i)
                DestroyResource(ref resource[i]);
        }

        private const int MaximumBiasTextureWidth = 16;
        private const int MaximumBiasTextureHeight = 16;
        private static readonly float[] MaximumBias =  
        {
            2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	1.876f,	1.809f,	1.772f,	1.753f,	1.748f,
            2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	1.869f,	1.801f,	1.764f,	1.745f,	1.739f,
            2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	1.976f,	1.841f,	1.774f,	1.737f,	1.716f,	1.71f,
            2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	1.914f,	1.784f,	1.716f,	1.673f,	1.649f,	1.641f,
            2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	1.793f,	1.676f,	1.604f,	1.562f,	1.54f,	1.533f,
            2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	1.802f,	1.619f,	1.536f,	1.492f,	1.467f,	1.454f,	1.449f,
            2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	1.812f,	1.575f,	1.496f,	1.456f,	1.432f,	1.416f,	1.408f,	1.405f,
            2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	1.555f,	1.479f,	1.438f,	1.413f,	1.398f,	1.387f,	1.381f,	1.379f,
            2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	1.812f,	1.555f,	1.474f,	1.43f,	1.404f,	1.387f,	1.376f,	1.368f,	1.363f,	1.362f,
            2.0f,	2.0f,	2.0f,	2.0f,	2.0f,	1.802f,	1.575f,	1.479f,	1.43f,	1.401f,	1.382f,	1.369f,	1.36f,	1.354f,	1.351f,	1.35f,
            2.0f,	2.0f,	1.976f,	1.914f,	1.793f,	1.619f,	1.496f,	1.438f,	1.404f,	1.382f,	1.367f,	1.357f,	1.349f,	1.344f,	1.341f,	1.34f,
            1.876f,	1.869f,	1.841f,	1.784f,	1.676f,	1.536f,	1.456f,	1.413f,	1.387f,	1.369f,	1.357f,	1.347f,	1.341f,	1.336f,	1.333f,	1.332f,
            1.809f,	1.801f,	1.774f,	1.716f,	1.604f,	1.492f,	1.432f,	1.398f,	1.376f,	1.36f,	1.349f,	1.341f,	1.335f,	1.33f,	1.328f,	1.327f,
            1.772f,	1.764f,	1.737f,	1.673f,	1.562f,	1.467f,	1.416f,	1.387f,	1.368f,	1.354f,	1.344f,	1.336f,	1.33f,	1.326f,	1.323f,	1.323f,
            1.753f,	1.745f,	1.716f,	1.649f,	1.54f,	1.454f,	1.408f,	1.381f,	1.363f,	1.351f,	1.341f,	1.333f,	1.328f,	1.323f,	1.321f,	1.32f,
            1.748f,	1.739f,	1.71f,	1.641f,	1.533f,	1.449f,	1.405f,	1.379f,	1.362f,	1.35f,	1.34f,	1.332f,	1.327f,	1.323f,	1.32f,	1.319f,
        };
    }
}