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

namespace FidelityFX
{
    /// <summary>
    /// This class is responsible for hooking into various Unity events and translating them to the FSR2 subsystem.
    /// This includes creation and destruction of the FSR2 context, as well as dispatching commands at the right time.
    /// This component also exposes various FSR2 parameters to the Unity inspector.
    /// </summary>
    [RequireComponent(typeof(Camera))]
    public class Fsr2Controller : MonoBehaviour
    {
        [SerializeField] private Fsr2.QualityMode qualityMode;
        
        [SerializeField] private bool performSharpenPass = true;

        [SerializeField, Range(0, 1)] private float sharpness = 0.8f;

        private Fsr2Context _context;
        
        private Camera _renderCamera;
        private RenderTexture _originalRenderTarget;
        private DepthTextureMode _originalDepthTextureMode;

        private GameObject _displayCameraObject;
        private Camera _displayCamera;
        private Fsr2Dispatcher _dispatcher;
        private readonly Fsr2.DispatchDescription _dispatchDescription = new Fsr2.DispatchDescription();

        private Fsr2.QualityMode _prevQualityMode;
        private Vector2Int _prevScreenSize;

        private CommandBuffer _opaqueOnlyCommandBuffer;
        private CommandBuffer _inputsCommandBuffer;
        
        private void OnEnable()
        {
            if (_displayCameraObject == null)
            {
                // Create a helper object with a camera that outputs at screen resolution
                _displayCameraObject = new GameObject("FSR2 Camera Object");
                _displayCameraObject.transform.SetParent(transform);
                _displayCameraObject.transform.SetPositionAndRotation(Vector3.zero, Quaternion.identity);
                //outputCameraObject.transform.hideFlags = HideFlags.HideInHierarchy | HideFlags.HideInInspector;

                // Create a camera that does nothing except output the upscaled image
                _displayCamera = _displayCameraObject.AddComponent<Camera>();
                _displayCamera.backgroundColor = Color.clear;
                _displayCamera.clearFlags = CameraClearFlags.Nothing;
                _displayCamera.eventMask = 0;
                _displayCamera.cullingMask = 0;
                _displayCamera.useOcclusionCulling = false;
                _displayCamera.orthographic = true;
                _displayCamera.allowMSAA = false;
                _displayCamera.renderingPath = RenderingPath.Forward;

                _dispatcher = _displayCameraObject.AddComponent<Fsr2Dispatcher>();
                _dispatcher.DispatchDescription = _dispatchDescription;
            }
            
            Fsr2.GetRenderResolutionFromQualityMode(out var renderWidth, out var renderHeight, Screen.width, Screen.height, qualityMode);
            _context = Fsr2.CreateContext(new Vector2Int(Screen.width, Screen.height), new Vector2Int(renderWidth, renderHeight), Fsr2.InitializationFlags.EnableMotionVectorsJitterCancellation);

            _dispatcher.Context = _context;
            _dispatcher.enabled = true;

            // Set up the original camera to output all of the required FSR2 input resources at the desired resolution
            _renderCamera = GetComponent<Camera>();
            _originalRenderTarget = _renderCamera.targetTexture; // TODO: if this isn't null, could maybe reuse this for the output texture?
            _originalDepthTextureMode = _renderCamera.depthTextureMode;

            _renderCamera.targetTexture = new RenderTexture(
                renderWidth, renderHeight,
                _originalRenderTarget != null ? _originalRenderTarget.depth : 32,
                _originalRenderTarget != null ? _originalRenderTarget.format : RenderTextureFormat.ARGBHalf) { name = "FSR2 Input Texture" };

            _renderCamera.targetTexture.Create();

            _renderCamera.depthTextureMode |= DepthTextureMode.Depth | DepthTextureMode.MotionVectors;

            // Create command buffers to bind the camera's output at the right moments in the rendering pipeline
            _opaqueOnlyCommandBuffer = new CommandBuffer { name = "FSR2 Opaque Input" };
            // TODO: may need to copy the opaque-only render buffer to a temp RT here, in which case we'll need an additional CommandBuffer to release the temp RT again
            _opaqueOnlyCommandBuffer.SetGlobalTexture(Fsr2Pipeline.SrvOpaqueOnly, BuiltinRenderTextureType.CameraTarget, RenderTextureSubElement.Color);
            _renderCamera.AddCommandBuffer(CameraEvent.BeforeImageEffectsOpaque, _opaqueOnlyCommandBuffer);

            _inputsCommandBuffer = new CommandBuffer { name = "FSR2 Inputs" };
            _inputsCommandBuffer.SetGlobalTexture(Fsr2Pipeline.SrvInputColor, BuiltinRenderTextureType.CameraTarget, RenderTextureSubElement.Color);
            _inputsCommandBuffer.SetGlobalTexture(Fsr2Pipeline.SrvInputDepth, BuiltinRenderTextureType.CameraTarget, RenderTextureSubElement.Depth);
            _inputsCommandBuffer.SetGlobalTexture(Fsr2Pipeline.SrvInputMotionVectors, BuiltinRenderTextureType.MotionVectors);
            _renderCamera.AddCommandBuffer(CameraEvent.BeforeImageEffects, _inputsCommandBuffer);

            // Adjust texture mipmap LOD bias by log2(renderResolution/displayResolution) - 1.0;
            // May need to leave this to the game dev, as we don't know which textures do and don't belong to the 3D scene
            float biasOffset = Fsr2.GetMipmapBiasOffset(renderWidth, Screen.width);
            Fsr2.GlobalCallbacks.ApplyMipmapBias(biasOffset);

            _prevScreenSize = new Vector2Int(Screen.width, Screen.height);
            _prevQualityMode = qualityMode;
        }

        private void OnDisable()
        {
            float biasOffset = Fsr2.GetMipmapBiasOffset(_renderCamera.targetTexture.width, _prevScreenSize.x);
            Fsr2.GlobalCallbacks.ApplyMipmapBias(-biasOffset);

            _renderCamera.RemoveCommandBuffer(CameraEvent.BeforeImageEffectsOpaque, _opaqueOnlyCommandBuffer);
            _opaqueOnlyCommandBuffer.Release();
            _opaqueOnlyCommandBuffer = null;

            _renderCamera.RemoveCommandBuffer(CameraEvent.BeforeImageEffects, _inputsCommandBuffer);
            _inputsCommandBuffer.Release();
            _inputsCommandBuffer = null;

            _renderCamera.targetTexture.Release();
            _renderCamera.targetTexture = _originalRenderTarget;
            _renderCamera.depthTextureMode = _originalDepthTextureMode;

            _dispatcher.Context = null;
            _dispatcher.enabled = false;
            _displayCamera.enabled = false;
            
            if (_context != null)
            {
                _context.Destroy();
                _context = null;
            }
        }

        private void Update()
        {
            _displayCamera.enabled = _renderCamera.enabled;

            if (Screen.width != _prevScreenSize.x || Screen.height != _prevScreenSize.y || qualityMode != _prevQualityMode)
            {
                // Force all resources to be destroyed and recreated with the new settings
                OnDisable();
                OnEnable();
            }
        }

        public void Reset()
        {
            _reset = true;
        }

        private Rect _tempRect;
        private bool _reset;
        
        private void OnPreRender()
        {
            _tempRect = _renderCamera.rect;
            _renderCamera.aspect = (Screen.width * _tempRect.width) / (Screen.height * _tempRect.height);
            _renderCamera.rect = new Rect(0, 0, 1, 1);

            var targetTexture = _renderCamera.targetTexture;
            
            _dispatchDescription.Color = null;
            _dispatchDescription.Depth = null;
            _dispatchDescription.MotionVectors = null;
            _dispatchDescription.Output = null;
            _dispatchDescription.Exposure = null;
            _dispatchDescription.Reactive = null;
            _dispatchDescription.PreExposure = 0;
            _dispatchDescription.EnableSharpening = performSharpenPass;
            _dispatchDescription.Sharpness = sharpness;
            _dispatchDescription.MotionVectorScale.x = -_renderCamera.pixelWidth;
            _dispatchDescription.MotionVectorScale.y = -_renderCamera.pixelHeight;
            _dispatchDescription.RenderSize = new Vector2Int(targetTexture.width, targetTexture.height);
            _dispatchDescription.FrameTimeDelta = Time.unscaledDeltaTime;
            _dispatchDescription.CameraNear = _renderCamera.nearClipPlane;
            _dispatchDescription.CameraFar = _renderCamera.farClipPlane;
            _dispatchDescription.CameraFovAngleVertical = _renderCamera.fieldOfView * Mathf.Deg2Rad;
            _dispatchDescription.ViewSpaceToMetersFactor = 1.0f; // 1 unit is 1 meter in Unity
            _dispatchDescription.Reset = _reset;
            _reset = false;

            // Perform custom jittering of the camera's projection matrix according to FSR2's instructions
            int jitterPhaseCount = Fsr2.GetJitterPhaseCount(targetTexture.width, Screen.width);
            Fsr2.GetJitterOffset(out float jitterX, out float jitterY, Time.frameCount, jitterPhaseCount);

            _dispatchDescription.JitterOffset = new Vector2(jitterX, jitterY);

            jitterX = 2.0f * jitterX / targetTexture.width;
            jitterY = 2.0f * jitterY / targetTexture.height;

            var jitterTranslationMatrix = Matrix4x4.Translate(new Vector3(jitterX, jitterY, 0));
            _renderCamera.projectionMatrix = jitterTranslationMatrix * _renderCamera.nonJitteredProjectionMatrix;
        }

        private void OnPostRender()
        {
            _renderCamera.rect = _tempRect;
            _renderCamera.ResetProjectionMatrix();
        }
    }

    /// <summary>
    /// Helper class to dispatch FSR2 commands on the display camera object, and render the final output texture.
    /// </summary>
    internal class Fsr2Dispatcher : MonoBehaviour
    {
        public Fsr2Context Context;
        public Fsr2.DispatchDescription DispatchDescription;
        
        private CommandBuffer _commandBuffer;

        private void OnEnable()
        {
            _commandBuffer = new CommandBuffer { name = "FSR2 Dispatch" };
        }

        private void OnDisable()
        {
            if (_commandBuffer != null)
            {
                _commandBuffer.Release();
                _commandBuffer = null;
            }
        }

        private void OnRenderImage(RenderTexture src, RenderTexture dest)
        {
            _commandBuffer.Clear();

            if (dest != null)
            {
                // We have more image effects lined up after this, so FSR2 can output straight to the intermediate render texture
                // TODO: we should probably use a shader to include depth & motion vectors into the output
                _commandBuffer.SetGlobalTexture(Fsr2Pipeline.UavUpscaledOutput, dest);
            }
            else
            {
                // We are rendering to the backbuffer, so we need a temporary render texture for FSR2 to output to
                _commandBuffer.GetTemporaryRT(Fsr2Pipeline.UavUpscaledOutput, Screen.width, Screen.height, 0, default, GraphicsFormat.R16G16B16A16_SFloat, 1, true);
            }

            Context.Dispatch(DispatchDescription, _commandBuffer);

            // Output upscaled image to screen
            // TODO: if `dest` is null, we likely don't care about the depth & motion vectors anymore
            if (dest == null)
            {
                _commandBuffer.Blit(Fsr2Pipeline.UavUpscaledOutput, dest);
                _commandBuffer.ReleaseTemporaryRT(Fsr2Pipeline.UavUpscaledOutput);
            }

            Graphics.ExecuteCommandBuffer(_commandBuffer);

            // Shut up the Unity warning about not writing to the destination texture 
            Graphics.SetRenderTarget(dest);
        }
    }
}