// Copyright (c) 2023 Nico de Poel
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Rendering;
using UnityEngine.Rendering.PostProcessing;
using FidelityFX;

namespace UnityEngine.Rendering.PostProcessing
{
    [Serializable]
    public class Fsr2QualityModeParameter : ParameterOverride<Fsr2.QualityMode>
    {
    }

    [Serializable]
    public class Fsr2GenerateReactiveParameters: ParameterOverride<Fsr2GenerateReactiveParams>
    {
    }

    [Serializable]
    public class Fsr2GenerateReactiveParams
    {
        [Range(0, 2)] public float scale = 0.5f;
        [Range(0, 1)] public float cutoffThreshold = 0.2f;
        [Range(0, 1)] public float binaryValue = 0.9f;
        public Fsr2.GenerateReactiveFlags flags = Fsr2.GenerateReactiveFlags.ApplyTonemap | Fsr2.GenerateReactiveFlags.ApplyThreshold | Fsr2.GenerateReactiveFlags.UseComponentsMax;
    }
    
    [UnityEngine.Scripting.Preserve]
    [Serializable]
    public class SuperResolution
    {
        public Fsr2QualityModeParameter qualityMode = new Fsr2QualityModeParameter() { value = Fsr2.QualityMode.Quality };

        public BoolParameter performSharpenPass = new BoolParameter() { value = true };
        [Range(0, 1)] public FloatParameter sharpness = new FloatParameter() { value = 0.8f };
        
        [Tooltip("Allow the use of half precision compute operations, potentially improving performance")]
        public BoolParameter enableFP16 = new BoolParameter() { value = false };

        [Header("Exposure")]
        public BoolParameter enableAutoExposure = new BoolParameter() { value = true };
        public FloatParameter preExposure = new FloatParameter() { value = 1.0f };
        public TextureParameter exposure = new TextureParameter() { value = null };

        [Header("Reactivity, Transparency & Composition")]
        public TextureParameter reactiveMask = new TextureParameter() { value = null };
        public TextureParameter transparencyAndCompositionMask = new TextureParameter() { value = null };
        public BoolParameter autoGenerateReactiveMask = new BoolParameter() { value = true };
        public Fsr2GenerateReactiveParameters generateReactiveParameters = new Fsr2GenerateReactiveParameters();

        public Vector2 jitter { get; private set; }

        private Fsr2Context _fsrContext;
        private Vector2Int _renderSize;
        private Vector2Int _displaySize;
        private bool _reset;

        private readonly Fsr2.DispatchDescription _dispatchDescription = new Fsr2.DispatchDescription();
        private readonly Fsr2.GenerateReactiveDescription _genReactiveDescription = new Fsr2.GenerateReactiveDescription();

        private RenderTexture _upscaledOutput;
        
        private Fsr2.QualityMode _prevQualityMode;
        private Vector2Int _prevDisplaySize;

        public bool IsSupported()
        {
            return SystemInfo.supportsComputeShaders && SystemInfo.supportsMotionVectors;
        }
        
        public DepthTextureMode GetCameraFlags()
        {
            return DepthTextureMode.Depth | DepthTextureMode.MotionVectors;
        }

        public void ResetHistory()
        {
            _reset = true;
        }

        public void ConfigureJitteredProjectionMatrix(PostProcessRenderContext context)
        {
            ApplyJitter(context.camera);
        }

        public void Render(PostProcessRenderContext context)
        {
            var cmd = context.command;

            if (!Application.isPlaying)
            {
                // We don't want this effect to start injecting scripts in edit mode, so just blit and skip the rest entirely
                cmd.BlitFullscreenTriangle(context.source, context.destination);
                return;
            }

            // Monitor for any resolution changes and recreate the FSR2 context if necessary
            // We can't create an FSR2 context without info from the post-processing context, so delay the initial setup until here
            if (_fsrContext == null || _displaySize.x != _prevDisplaySize.x || _displaySize.y != _prevDisplaySize.y || qualityMode != _prevQualityMode)
            {
                DestroyFsrContext();
                CreateFsrContext(context);
                
                _prevQualityMode = qualityMode;
            }

            // Effects rendering happens in OnPreCull, so this is the right place to apply camera jittering
            ApplyJitter(context.camera);
            
            cmd.SetGlobalTexture(Fsr2ShaderIDs.SrvInputColor, BuiltinRenderTextureType.CameraTarget, RenderTextureSubElement.Color);
            cmd.SetGlobalTexture(Fsr2ShaderIDs.SrvInputDepth, BuiltinRenderTextureType.CameraTarget, RenderTextureSubElement.Depth);
            cmd.SetGlobalTexture(Fsr2ShaderIDs.SrvInputMotionVectors, BuiltinRenderTextureType.MotionVectors);
            
            SetupDispatchDescription(context);

            if (autoGenerateReactiveMask)
            {
                // TODO: auto-generate reactive mask
            }
            
            cmd.GetTemporaryRT(Fsr2ShaderIDs.UavUpscaledOutput, _displaySize.x, _displaySize.y, 0, default, context.sourceFormat, default, 1, true);
            // _dispatchDescription.Output = _upscaledOutput;
            
            _fsrContext.Dispatch(_dispatchDescription, cmd);
            
            cmd.BlitFullscreenTriangle(Fsr2ShaderIDs.UavUpscaledOutput, context.destination);
            cmd.ReleaseTemporaryRT(Fsr2ShaderIDs.UavUpscaledOutput);
        }

        private void CreateFsrContext(PostProcessRenderContext context)
        {
            _displaySize = new Vector2Int(context.width, context.height);
            
            _prevDisplaySize = _displaySize;
            
            // TODO: re-enable actual resolution scaling
            // Fsr2.GetRenderResolutionFromQualityMode(out var renderWidth, out var renderHeight, _displaySize.x, _displaySize.y, qualityMode);
            // _renderSize = new Vector2Int(renderWidth, renderHeight);
            _renderSize = _displaySize;
            
            Fsr2.InitializationFlags flags = 0;
            if (context.camera.allowHDR) flags |= Fsr2.InitializationFlags.EnableHighDynamicRange;
            if (enableFP16) flags |= Fsr2.InitializationFlags.EnableFP16Usage;
            if (enableAutoExposure) flags |= Fsr2.InitializationFlags.EnableAutoExposure;

            _fsrContext = Fsr2.CreateContext(_displaySize, _renderSize, new Callbacks(context.resources), flags);

            // Apply a mipmap bias so that textures retain their sharpness
            float biasOffset = Fsr2.GetMipmapBiasOffset(_renderSize.x, _displaySize.x);
            //Fsr2.GlobalCallbacks.ApplyMipmapBias(biasOffset);

            _upscaledOutput = new RenderTexture(_displaySize.x, _displaySize.y, 0, context.sourceFormat) { name = "FSR2 Upscaled Output", enableRandomWrite = true };
            _upscaledOutput.Create();
        }
        
        private void DestroyFsrContext()
        {
            if (_fsrContext != null)
            {
                _fsrContext.Destroy();
                _fsrContext = null;
                
                // Undo the previous mipmap bias adjustment
                float biasOffset = Fsr2.GetMipmapBiasOffset(_renderSize.x, _prevDisplaySize.x);
                //Fsr2.GlobalCallbacks.ApplyMipmapBias(-biasOffset);
            }

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

        private void ApplyJitter(Camera camera)
        {
            // Perform custom jittering of the camera's projection matrix according to FSR2's recipe
            int jitterPhaseCount = Fsr2.GetJitterPhaseCount(_renderSize.x, _displaySize.x);
            Fsr2.GetJitterOffset(out float jitterX, out float jitterY, Time.frameCount, jitterPhaseCount);
            
            _dispatchDescription.JitterOffset = new Vector2(jitterX, jitterY);

            jitterX = 2.0f * jitterX / _renderSize.x;
            jitterY = 2.0f * jitterY / _renderSize.y;

            var jitterTranslationMatrix = Matrix4x4.Translate(new Vector3(jitterX, jitterY, 0));
            camera.nonJitteredProjectionMatrix = camera.projectionMatrix;
            camera.projectionMatrix = jitterTranslationMatrix * camera.nonJitteredProjectionMatrix;
            camera.useJitteredProjectionMatrixForTransparentRendering = false;

            jitter = new Vector2(jitterX, jitterY);
        }

        private void SetupDispatchDescription(PostProcessRenderContext context)
        {
            var camera = context.camera;
            
            // Set up the main FSR2 dispatch parameters
            // The input and output textures are left blank here, as they are already being bound elsewhere in this source file
            _dispatchDescription.Color = null;
            _dispatchDescription.Depth = null;
            _dispatchDescription.MotionVectors = null;
            _dispatchDescription.Exposure = null;
            _dispatchDescription.Reactive = null;
            _dispatchDescription.TransparencyAndComposition = null;
            
            if (!enableAutoExposure && exposure.value != null) _dispatchDescription.Exposure = exposure.value;
            if (reactiveMask.value != null) _dispatchDescription.Reactive = reactiveMask.value;
            if (transparencyAndCompositionMask.value != null) _dispatchDescription.TransparencyAndComposition = transparencyAndCompositionMask.value;
            
            _dispatchDescription.Output = null;
            _dispatchDescription.PreExposure = preExposure;
            _dispatchDescription.EnableSharpening = performSharpenPass;
            _dispatchDescription.Sharpness = sharpness;
            _dispatchDescription.MotionVectorScale.x = -_renderSize.x;
            _dispatchDescription.MotionVectorScale.y = -_renderSize.y;
            _dispatchDescription.RenderSize = _renderSize;
            _dispatchDescription.FrameTimeDelta = Time.unscaledDeltaTime;
            _dispatchDescription.CameraNear = camera.nearClipPlane;
            _dispatchDescription.CameraFar = camera.farClipPlane;
            _dispatchDescription.CameraFovAngleVertical = camera.fieldOfView * Mathf.Deg2Rad;
            _dispatchDescription.ViewSpaceToMetersFactor = 1.0f; // 1 unit is 1 meter in Unity
            _dispatchDescription.Reset = _reset;
            _reset = false;
            
            if (SystemInfo.usesReversedZBuffer)
            {
                // Swap the near and far clip plane distances as FSR2 expects this when using inverted depth
                (_dispatchDescription.CameraNear, _dispatchDescription.CameraFar) = (_dispatchDescription.CameraFar, _dispatchDescription.CameraNear);
            }
        }

        public void Release()
        {
            DestroyFsrContext();
        }

        private class Callbacks : Fsr2CallbacksBase
        {
            private readonly PostProcessResources _resources;
            
            public Callbacks(PostProcessResources resources)
            {
                _resources = resources;
            }
            
            public override ComputeShader LoadComputeShader(string name)
            {
                return _resources.computeShaders.FindComputeShader(name);
            }

            public override void UnloadComputeShader(ComputeShader shader)
            {
            }
        }
    }
}