/*
 * Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
 *
 * NVIDIA CORPORATION and its licensors retain all intellectual property
 * and proprietary rights in and to this software, related documentation
 * and any modifications thereto.  Any use, reproduction, disclosure or
 * distribution of this software and related documentation without an express
 * license agreement from NVIDIA CORPORATION is strictly prohibited.
 */
using System;
using System.Collections;
using System.Runtime.InteropServices;
using UnityEngine;
using UnityEngine.Rendering;
using NVStreamline;

namespace NVStreamline
{
    [AddComponentMenu("NVIDIA/Streamline Reflex Builtin & HDRP")]
    public class StreamlineReflex : MonoBehaviour
    {
        [Tooltip("Is this the first camera to be rendered")]
        public bool isFirstCamera = true;
        private bool isAttachedBeforeRender = false;
        [Tooltip("Is this the last camera to be rendered")]
        public bool isLastCamera = true;
        private bool isAttachedAfterRender = false;
        [Tooltip("Minimum interval in microseconds for frame synchronization. 0 = no frame rate limit")]
        public uint intervalUs = 0;
        [Tooltip("NVIDIA Reflex Low Latency mode.")]
        public bool isLowLatencyMode = true;
        [Tooltip("NVIDIA Reflex Low Latency Boost")]
        public bool isLowLatencyBoost = false;
        [Tooltip("NVIDIA Reflex Low Latency Boost")]
        public bool useMarkersToOptimize = false;
        private static uint previousIntervalUs = 0;
        private static bool previousIsLowLatencyMode = false;
        private static bool previousIsLowLatencyBoost = true;
        private static bool previousUseMarkersToOptimize = false;
        private bool hasSimulationStarted = false;
        private CommandBuffer renderSubmitStart_CommandBuffer;
        private CommandBuffer renderSubmitEnd_CommandBuffer;
        public bool isIgnored = false;                          // Used by DLSS plugin for downsampled camera.

        private uint m_id;
        private bool m_reflexEnabled = false;

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

                if (StreamlineReflexCore.NvReflex_Plugin_HasReceivedRenderEvent() != 0)
                {
                    if (isLastCamera)
                    {
                        if ((previousIsLowLatencyMode != isLowLatencyMode) || (previousIntervalUs != intervalUs) || (previousIsLowLatencyBoost != isLowLatencyBoost) || (previousUseMarkersToOptimize != useMarkersToOptimize))
                        {
                            StreamlineReflexCore.NvReflex_Plugin_SetSleepMode(isLowLatencyMode, intervalUs, isLowLatencyBoost, useMarkersToOptimize);
                            previousIsLowLatencyMode = isLowLatencyMode;
                            previousIsLowLatencyBoost = isLowLatencyBoost;
                            previousIntervalUs = intervalUs;
                            previousUseMarkersToOptimize = useMarkersToOptimize;
                        }
                        if (isLowLatencyMode)
                        {
                            GL.Flush();
                            StreamlineReflexCore.NvReflex_Plugin_Sleep((ulong)Time.frameCount);
                        }
                    }
                }
            }
        }

        private void Awake()
        {
            if (!StreamlineCore.IsFeatureSupported(StreamlineCore.Feature.eFeatureReflex))
            {
                Debug.Log("Reflex is not supported");
                return;
            }

            int id = StreamlineCore.CreateFeature(StreamlineCore.Feature.eFeatureReflex);
            if (id != 1)  //CreateFeature must return 1 for Reflex
            {
                Debug.Log("Reflex Awake FAILED to get feature id");
                m_reflexEnabled = false;
                return;
            }
            m_id = (uint)id;
            m_reflexEnabled = true;

            renderSubmitStart_CommandBuffer = new CommandBuffer();
            renderSubmitStart_CommandBuffer.name = "Reflex RENDERSUBMIT_START";
            renderSubmitEnd_CommandBuffer = new CommandBuffer();
            renderSubmitEnd_CommandBuffer.name = "Reflex RENDERSUBMIT_END";
        }

        void OnDestroy()
        {
            StreamlineCore.DestroyFeature((int)m_id);
        }

        private void OnBeginFrameRendering(ScriptableRenderContext context, Camera[] camera)
        {
            context.ExecuteCommandBuffer(renderSubmitStart_CommandBuffer);
        }

        private void OnEndFrameRendering(ScriptableRenderContext context, Camera[] camera)
        {
            context.ExecuteCommandBuffer(renderSubmitEnd_CommandBuffer);
            context.Submit();

        }

        private void AfterPresent()
        {
            if (hasSimulationStarted == false)
            {
                int simStartEventID = StreamlineReflexCore.NvReflex_Plugin_GetEventIDForEvent((int)StreamlineReflexCore.NvReflex_LATENCY_MARKER_TYPE.SIMULATION_START);
                StreamlineReflexCore.SetLatencyMarker(simStartEventID, (ulong)Time.frameCount);

                hasSimulationStarted = true;
            }
        }

        public void FixedUpdate()
        {
            if (m_reflexEnabled)
            {
                if (isFirstCamera &&
                 !isIgnored &&
                 StreamlineReflexCore.NvReflex_Plugin_HasReceivedRenderEvent() != 0)
                {
                    AfterPresent();
                }
            }
        }

        public void Update()
        {
            if (m_reflexEnabled)
            {
                if (isFirstCamera &&
                 !isIgnored &&
                  StreamlineReflexCore.NvReflex_Plugin_HasReceivedRenderEvent() != 0)
                {

                    AfterPresent();
                    int inputSampleEventID = StreamlineReflexCore.NvReflex_Plugin_GetEventIDForEvent((int)StreamlineReflexCore.NvReflex_LATENCY_MARKER_TYPE.INPUT_SAMPLE);
                    StreamlineReflexCore.SetLatencyMarker(inputSampleEventID, (ulong)Time.frameCount);

                    if (Input.GetKeyUp(KeyCode.F13))
                    {
                        StreamlineReflexCore.PCLatencyPing();
                    }
                    if (Input.GetMouseButtonDown(0))
                    {
                        StreamlineReflexCore.TriggerFlash();
                    }
                }
            }
        }

        private void AttachDetachCommandBuffers()
        {
            if ((renderSubmitStart_CommandBuffer != null) & (renderSubmitEnd_CommandBuffer != null))
            {
                Camera camComponent = this.GetComponent<Camera>();
                bool shouldFirstCameraBeAttached = (isFirstCamera & !isIgnored & enabled);
                bool shouldLastCameraBeAttached = (isLastCamera & /*!isIgnored &*/ enabled);

                if (shouldFirstCameraBeAttached != isAttachedBeforeRender)
                {
                    if (shouldFirstCameraBeAttached)
                    {
                        // Send renderSubmitStart_CommandBuffer on BeforeDepthTexture, BeforeForwardOpaque and BeforeGBuffer events. This ensures whatever configuration the Legacy pipeline is using
                        // with regard to 2D/3D and Deferred/Rendering the event will be sent.  The plugin filters the SubmitStart event and only processes the first that occurs.
                        var oldBDCmdBuffers = camComponent.GetCommandBuffers(CameraEvent.BeforeDepthTexture);
                        var oldBOCmdBuffers = camComponent.GetCommandBuffers(CameraEvent.BeforeForwardOpaque);
                        var oldBDGmdBuffers = camComponent.GetCommandBuffers(CameraEvent.BeforeGBuffer);
                        camComponent.RemoveCommandBuffers(CameraEvent.BeforeDepthTexture);
                        camComponent.RemoveCommandBuffers(CameraEvent.BeforeForwardOpaque);
                        camComponent.RemoveCommandBuffers(CameraEvent.BeforeGBuffer);
                        camComponent.AddCommandBuffer(CameraEvent.BeforeDepthTexture, renderSubmitStart_CommandBuffer);
                        camComponent.AddCommandBuffer(CameraEvent.BeforeForwardOpaque, renderSubmitStart_CommandBuffer);
                        camComponent.AddCommandBuffer(CameraEvent.BeforeGBuffer, renderSubmitStart_CommandBuffer);
                        foreach (var cmdBuff in oldBDCmdBuffers)
                        {
                            camComponent.AddCommandBuffer(CameraEvent.BeforeDepthTexture, cmdBuff);
                        }
                        foreach (var cmdBuff in oldBOCmdBuffers)
                        {
                            camComponent.AddCommandBuffer(CameraEvent.BeforeForwardOpaque, cmdBuff);
                        }
                        foreach (var cmdBuff in oldBDGmdBuffers)
                        {
                            camComponent.AddCommandBuffer(CameraEvent.BeforeGBuffer, cmdBuff);
                        }
#if UNITY_2019_1_OR_NEWER
                        RenderPipelineManager.beginFrameRendering += OnBeginFrameRendering;
#endif
                    }
                    else
                    {
                        camComponent.RemoveCommandBuffer(CameraEvent.BeforeDepthTexture, renderSubmitStart_CommandBuffer);
                        camComponent.RemoveCommandBuffer(CameraEvent.BeforeForwardOpaque, renderSubmitStart_CommandBuffer);
                        camComponent.RemoveCommandBuffer(CameraEvent.BeforeGBuffer, renderSubmitStart_CommandBuffer);
#if UNITY_2019_1_OR_NEWER
                        RenderPipelineManager.beginFrameRendering -= OnBeginFrameRendering;
#endif
                    }
                    isAttachedBeforeRender = shouldFirstCameraBeAttached;
                }

                if (shouldLastCameraBeAttached != isAttachedAfterRender)
                {
                    if (shouldLastCameraBeAttached)
                    {
                        camComponent.AddCommandBuffer(CameraEvent.AfterEverything, renderSubmitEnd_CommandBuffer);
#if UNITY_2019_1_OR_NEWER
                        RenderPipelineManager.endFrameRendering += OnEndFrameRendering;
#endif
                    }
                    else
                    {
                        camComponent.RemoveCommandBuffer(CameraEvent.AfterEverything, renderSubmitEnd_CommandBuffer);
#if UNITY_2019_1_OR_NEWER
                        RenderPipelineManager.endFrameRendering -= OnEndFrameRendering;
#endif
                    }
                    isAttachedAfterRender = shouldLastCameraBeAttached;
                }
            }
        }

        void OnEnable()
        {
            if (m_reflexEnabled)
            {
                AttachDetachCommandBuffers();
                StartCoroutine(SleepCoroutine());
            }
        }

        void OnDisable()
        {
            if (m_reflexEnabled)
            {
                AttachDetachCommandBuffers();
                StopCoroutine(SleepCoroutine());
            }
        }


        private void LateUpdate()
        {
            if (m_reflexEnabled)
            {
                AttachDetachCommandBuffers();

                if (isFirstCamera & !isIgnored)
                {
                    if (hasSimulationStarted &&
                     StreamlineReflexCore.NvReflex_Plugin_HasReceivedRenderEvent() != 0)
                    {
                        int simEndEventID = StreamlineReflexCore.NvReflex_Plugin_GetEventIDForEvent((int)StreamlineReflexCore.NvReflex_LATENCY_MARKER_TYPE.SIMULATION_END);
                        StreamlineReflexCore.SetLatencyMarker(simEndEventID, (ulong)Time.frameCount);
                        hasSimulationStarted = false;
                    }
                    int startEventID = StreamlineReflexCore.NvReflex_Plugin_GetEventIDForEvent((int)StreamlineReflexCore.NvReflex_LATENCY_MARKER_TYPE.RENDERSUBMIT_START);
                    renderSubmitStart_CommandBuffer.Clear();
                    renderSubmitStart_CommandBuffer.IssuePluginEventAndData(StreamlineReflexCore.NvReflex_Plugin_GetRenderEventAndDataFunc(), startEventID, (IntPtr)(ulong)Time.frameCount);
                }
                if (isLastCamera /*& !isIgnored*/)
                {
                    int endEventID = StreamlineReflexCore.NvReflex_Plugin_GetEventIDForEvent((int)StreamlineReflexCore.NvReflex_LATENCY_MARKER_TYPE.RENDERSUBMIT_END);
                    renderSubmitEnd_CommandBuffer.Clear();
                    renderSubmitEnd_CommandBuffer.IssuePluginEventAndData(StreamlineReflexCore.NvReflex_Plugin_GetRenderEventAndDataFunc(), endEventID, (IntPtr)(ulong)Time.frameCount);

                    endEventID = StreamlineReflexCore.NvReflex_Plugin_GetEventIDForEvent((int)StreamlineReflexCore.NvReflex_LATENCY_MARKER_TYPE.PRESENT_START);
                    renderSubmitEnd_CommandBuffer.IssuePluginEventAndData(StreamlineReflexCore.NvReflex_Plugin_GetRenderEventAndDataFunc(), endEventID, (IntPtr)(ulong)Time.frameCount);

                    endEventID = StreamlineReflexCore.NvReflex_Plugin_GetEventIDForEvent((int)StreamlineReflexCore.NvReflex_LATENCY_MARKER_TYPE.PRESENT_END);
                    renderSubmitEnd_CommandBuffer.IssuePluginEventAndData(StreamlineReflexCore.NvReflex_Plugin_GetRenderEventAndDataFunc(), endEventID, (IntPtr)(ulong)Time.frameCount);
                }
            }
        }

    }
}