#if UNITY_POST_PROCESSING_STACK_V2
using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Rendering;
using UnityEngine.Rendering.PostProcessing;

namespace FidelityFX
{
    [Serializable]
    [PostProcess(typeof(Fsr2PostProcessRenderer), PostProcessEvent.BeforeStack, "FidelityFX/FSR 2", false)]
    public class Fsr2PostProcessEffect : PostProcessEffectSettings
    {
        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 override bool IsEnabledAndSupported(PostProcessRenderContext context)
        {
            return base.IsEnabledAndSupported(context) && SystemInfo.supportsComputeShaders;
        }
    }

    [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;
    }

    public class Fsr2PostProcessRenderer : PostProcessEffectRenderer<Fsr2PostProcessEffect>
    {
        private Fsr2Context _fsrContext;
        private Fsr2PostProcessHelper _helper;

        private ref Vector2Int DisplaySize => ref _helper.DisplaySize;
        private ref Vector2Int RenderSize => ref _helper.RenderSize;
        
        private readonly Fsr2.DispatchDescription _dispatchDescription = new Fsr2.DispatchDescription();
        private readonly Fsr2.GenerateReactiveDescription _genReactiveDescription = new Fsr2.GenerateReactiveDescription();
        
        private Fsr2.QualityMode _prevQualityMode;
        private Vector2Int _prevDisplaySize;
        
        public override void Init()
        {
            base.Init();
            
            Debug.Log($"Initializing FSR2 post-process effect, quality = {settings.qualityMode.value}");
        }

        public override void Release()
        {
            Debug.Log("Releasing FSR2 post-process effect");
            
            base.Release();
        }

        public override DepthTextureMode GetCameraFlags()
        {
            return DepthTextureMode.Depth | DepthTextureMode.MotionVectors;
        }

        public override 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;
            }

            // Ensure that the helper script exists and is enabled
            if (_helper == null || !_helper.enabled)
            {
                if (_helper == null)
                    Inject(context);
                
                _helper.enabled = settings.IsEnabledAndSupported(context);
                
                // The injected script won't come into effect until the next frame, so just do a simple blit this frame
                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 || settings.qualityMode != _prevQualityMode)
            {
                DestroyFsrContext();
                CreateFsrContext(context);
                
                _prevQualityMode = settings.qualityMode;
            }

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

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

        /// <summary>
        /// Inject ourselves into the camera's render loop at points where PPV2 normally wouldn't allow us.
        /// This also sets up PPV2's own stack to use FSR2's upscaled output as its source. 
        /// </summary>
        private void Inject(PostProcessRenderContext context)
        {
            var go = context.camera.gameObject;
            _helper = go.AddComponent<Fsr2PostProcessHelper>();
            
            _helper.Settings = settings;
            // TODO: create output RT, set it as PPV2's post-transparency stack source
        }
        
        private void CreateFsrContext(PostProcessRenderContext context)
        {
            _prevDisplaySize = DisplaySize;
            
            Fsr2.GetRenderResolutionFromQualityMode(out var renderWidth, out var renderHeight, DisplaySize.x, DisplaySize.y, settings.qualityMode);
            RenderSize = new Vector2Int(renderWidth, renderHeight);
            
            Fsr2.InitializationFlags flags = 0;
            if (context.camera.allowHDR) flags |= Fsr2.InitializationFlags.EnableHighDynamicRange;
            if (settings.enableFP16) flags |= Fsr2.InitializationFlags.EnableFP16Usage;
            if (settings.enableAutoExposure) flags |= Fsr2.InitializationFlags.EnableAutoExposure;

#if UNITY_EDITOR || DEVELOPMENT_BUILD
            flags |= Fsr2.InitializationFlags.EnableDebugChecking;
#endif
            
            _fsrContext = Fsr2.CreateContext(DisplaySize, RenderSize, flags);

            // Apply a mipmap bias so that textures retain their sharpness
            float biasOffset = Fsr2.GetMipmapBiasOffset(RenderSize.x, DisplaySize.x);
            Fsr2.GlobalCallbacks.ApplyMipmapBias(biasOffset);
        }
        
        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);
            }
        }

        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);

            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 = true;
        }

        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 (!settings.enableAutoExposure && settings.exposure.value != null) _dispatchDescription.Exposure = settings.exposure.value;
            if (settings.reactiveMask.value != null) _dispatchDescription.Reactive = settings.reactiveMask.value;
            if (settings.transparencyAndCompositionMask.value != null) _dispatchDescription.TransparencyAndComposition = settings.transparencyAndCompositionMask.value;
            
            _dispatchDescription.Output = null;
            _dispatchDescription.PreExposure = settings.preExposure;
            _dispatchDescription.EnableSharpening = settings.performSharpenPass;
            _dispatchDescription.Sharpness = settings.sharpness;
            _dispatchDescription.MotionVectorScale.x = -RenderSize.x;
            _dispatchDescription.MotionVectorScale.y = -RenderSize.y;
            _dispatchDescription.RenderSize = RenderSize;
            _dispatchDescription.InputResourceSize = new Vector2Int(context.width, context.height);
            _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 = m_ResetHistory;
            
            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);
            }
        }
    }

    internal class Fsr2PostProcessHelper : MonoBehaviour
    {
        internal Fsr2PostProcessEffect Settings;
        internal Vector2Int DisplaySize;
        internal Vector2Int RenderSize;

        private Camera _camera;
        private Rect _originalRect;

        private void Awake()
        {
            _camera = GetComponent<Camera>();
            // TODO: inject opaque-only command buffer & anything else?

            DisplaySize.x = _camera.pixelWidth;
            DisplaySize.y = _camera.pixelHeight;
            RenderSize = DisplaySize;
        }

        private void Start()
        {
            StartCoroutine(CResetCamera());
        }

        /// <summary>
        /// This needs to run before PostProcessLayer's OnPreCull, hence why we place this code in LateUpdate.
        /// </summary>
        private void LateUpdate()
        {
            if (!Settings.enabled)
            {
                enabled = false;
                return;
            }

            // Relay information about the camera's output size before rescaling
            DisplaySize.x = _camera.pixelWidth;
            DisplaySize.y = _camera.pixelHeight;
            
            // Render to a smaller portion of the screen by manipulating the camera's viewport rect
            _originalRect = _camera.rect;
            _camera.aspect = (_camera.pixelWidth * _originalRect.width) / (_camera.pixelHeight * _originalRect.height);
            _camera.rect = new Rect(0, 0, _originalRect.width * RenderSize.x / DisplaySize.x, _originalRect.height * RenderSize.y / DisplaySize.y);
        }

        // private void OnPreCull()    // TODO: may need to do this in LateUpdate instead? So we execute before PPV2's OnPreCull
        // {
        //     // TODO: check if FSR2 is still enabled; if not: reset PPV2 source texture & disable self
        //     // TODO: fiddle with the camera parameters, rect, jitter, etc
        // }

        private IEnumerator CResetCamera()
        {
            while (true)
            {
                yield return new WaitForEndOfFrame();

                _camera.rect = _originalRect;
                _camera.ResetProjectionMatrix();
            }
        }
    }
}
#endif  // UNITY_POST_PROCESSING_STACK_V2