// %BANNER_BEGIN%
// ---------------------------------------------------------------------
// %COPYRIGHT_BEGIN%
// Copyright (c) 2022-2023 Magic Leap, Inc. All Rights Reserved.
// Use of this file is governed by the Magic Leap 2 Software License Agreement, located here: https://www.magicleap.com/software-license-agreement-ml2  
// Terms and conditions applicable to third-party materials accompanying this distribution may also be found in the top-level NOTICE file appearing herein.
// %COPYRIGHT_END%
// ---------------------------------------------------------------------
// %BANNER_END%

using System;
using System.Collections;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.IO;
using UnityEngine;
using UnityEngine.Pool;
using static MagicLeap.Spectator.SpectatorMessages;

namespace MagicLeap.Spectator
{
    public class MLSpectatorRenderer : MonoBehaviour
    {
        #region Inspector variables

        [SerializeField, Tooltip("The maximum number of requets we should queue before we start skipping")]
        private int maxRequests = 5;

        #endregion

        #region Private variables

        // Virtual camera for rendering our spectator content
        private Camera renderCam = null;

        // Target texture for our virtual camera
        private RenderTexture render = null;

        // Material for separating rgb and alpha values for the encoder
        private Material splitMat = null;

        // Encoder for encoding splitTex
        private H264Encoder encoder = null;

        // Quad for displaying nonlinear depth image
        private GameObject depthScreen = null;

        // Linear depth image texture
        private Texture2D depthTex = null;

        // Upscaled depth image texture
        private RenderTexture depthRender = null;

        // Material used to upscale our depth
        private Material depthUpscale = null;

        // Material applied to depth screen for converting linear depth values to nonlinear 
        private Material depthMat = null;

        // Whether or not we have already been initialized
        private bool initialized = false;

        // Whether or not our process loop is running
        private bool running = false;

        // A pool to draw requests from to limit buffer (re)allocation
        private IObjectPool<MLSpectatorRequest> requestPool = null;

        // Render request queue (thread safe)
        private ConcurrentQueue<MLSpectatorRequest> requestQueue = new();

        // Are we currently connected
        private bool connected = false;

        // All lights in the scene
        private List<Light> sceneLights = new();

        // Light attached to my gameObject
        private Light myLight = null;

        // The time of our last request (used in place of heartbeats)
        private float timeOfLastRequest = 0.0f;

        // Our editor preview
        private MLSpectatorPreview editorPreview = null;

        #endregion

        #region Public events

        // Callback for custom render event
        public event Action<Camera> onRenderRequest = null;

        #endregion

        #region MonoBehaviour methods

        private void OnEnable()
        {
            // Create our render camera
            SetupRenderCam();

            // Create our depth screen
            SetupDepthScreen();

            // Setup lighting elements
            SetupLights();

            // Subscribe to host connection
            Host.Instance.OnConnected += OnConnected;

            // Check if we are already connected
            if (Host.Instance.Connected) OnConnected();
        }

        private void OnDisable()
        {
            // Unsubscribe from host connection
            Host.Instance.OnConnected -= OnConnected;

            // Disconnect if we are connected
            if (connected) OnDisconnect();

            // Destroy render camera material
            Destroy(splitMat);

            // Destroy our depth screen material
            Destroy(depthMat);
        }

        #endregion

        #region Application methods

        private void OnApplicationPause(bool pause)
        {
            Debug.Log($"MLSpectatorRenderer: OnApplicationPause() - {(pause ? "Paused" : "Resumed")}");
            Host.Instance.Send(new Message(SpectatorFocus, BitConverter.GetBytes(!pause)));
        }

        #endregion

        #region Setup methods

        private void SetupRenderCam()
        {
            // Add a camera component
            renderCam = gameObject.AddComponent<Camera>();

            // Set our camera with the appropriate values
            renderCam.depthTextureMode = DepthTextureMode.Depth;
            renderCam.clearFlags = CameraClearFlags.Color;
            renderCam.backgroundColor = Color.clear;
            renderCam.forceIntoRenderTexture = true;
            renderCam.allowMSAA = true;
            renderCam.allowHDR = false;

            // Create our material for separating rgb and alpha
            splitMat = new Material(Shader.Find("Image/SplitAlpha"));

            // Set inactive
            renderCam.enabled = false;

            // Force our camera output into a render texture
            renderCam.forceIntoRenderTexture = true;

            // Use better clipping planes
            renderCam.nearClipPlane = 0.1f;
            renderCam.farClipPlane = 10.0f;
        }

        private void SetupDepthScreen()
        {
            // Create our quad
            depthScreen = GameObject.CreatePrimitive(PrimitiveType.Quad);
            depthScreen.transform.parent = renderCam.transform;
            depthScreen.transform.localPosition = Vector3.forward * (renderCam.nearClipPlane + 1e-6f);
            depthScreen.transform.localRotation = Quaternion.identity;
            depthScreen.transform.localScale = Vector3.one;

            // Destroy our collider
            var collider = depthScreen.GetComponent<Collider>();
            if (collider) Destroy(collider);

            // Create our depth material for converting linear to nonlinear values
            var mr = depthScreen.GetComponent<MeshRenderer>();
            depthMat = new Material(Shader.Find("Hidden/IndirectDepth"));
            depthMat.SetTexture("_Depth", Texture2D.whiteTexture);
            mr.material = depthMat;

            // Set inactive
            depthScreen.SetActive(false);
        }

        private void SetupLights()
        {
            // Detect attached light component
            myLight = GetComponentInChildren<Light>();
            // Disable our light if we have one
            if (myLight != null) myLight.enabled = false;

            // Don't control lights if we weren't told to
            if (!MLSpectator.Instance.ControlSceneLights)
            {
                myLight = null;
            }

            // Otherwise, if we don't have a light
            else if (myLight == null)
            {
                Debug.LogWarning($"MLSpectatorRenderer: SetupLights() - No light component found");
            }

            // Otherwise, we do have a light
            else
            {
                Debug.Log($"MLSpectatorRenderer: SetupLights() - Light component found");
                // Find all active lights in our scene in no particular order
                var lights = FindObjectsByType<Light>(FindObjectsInactive.Exclude, FindObjectsSortMode.None);
                // Make absolutely certain that myLight is not included (this happens for some reason)
                foreach (var light in lights) if (light != myLight) sceneLights.Add(light);
                Debug.Log($"MLSpectatorRenderer: SetupLights() - {sceneLights.Count} active lights in scene");
                // Can't control lights if there aren't any any lights in our scene
                if (sceneLights.Count == 0) myLight = null;
            }
        }

        #endregion

        #region Host callbacks

        private void OnConnected()
        {
            //Debug.Log("MLSpectatorRenderer: OnConnected()");

            // Subscribe to host disconnect
            Host.Instance.OnDisconnect += OnDisconnect;

            // Mark that we are connected
            connected = true;
        }

        private void OnDisconnect()
        {
            //Debug.Log("MLSpectatorRenderer: OnDisconnect()");

            // Unsubscribe from host disconnect
            Host.Instance.OnDisconnect -= OnDisconnect;

            // Reset our state
            Reset();

            // Mark that we are no longer connected
            connected = false;
        }

        private void HandleRequest(Message message)
        {
            // If we have too many requests queued
            if (requestQueue.Count + 1 > maxRequests)
            {
                // Skip the first in our queue
                if (requestQueue.TryDequeue(out var skip))
                {
                    Host.Instance.Send(new Message(RenderIgnored, BitConverter.GetBytes(skip.id)));
                    lock (requestPool)
                        requestPool.Release(skip);
                }
            }

            // Grab a request from our pool, load it, and push it onto our queue
            MLSpectatorRequest request = null;
            lock (requestPool)
                request = requestPool.Get();
            request.Load(message.data);
            requestQueue.Enqueue(request);
        }

        private void HandleKeyFrameRequest(Message message)
        {
            if (!encoder.TryRequestSyncFrame())
                Debug.LogWarning("MLSpectatorRenderer: HandleKeyFrameRequest() - Failed to provide key frame");
        }

        #endregion

        #region Private methods

        private Ping ping = null;

        private IEnumerator ProcessRequestLoop()
        {
            if (running)
            {
                Debug.Log("MLSpectatorRenderer: ProcessRequestLoop() - Already running");
                yield break;
            }
            running = true;

            Debug.Log("MLSpectatorRenderer: ProcessRequestLoop() - Begin");

            while (initialized)
            {
                //Debug.Log($"MLSpectatorRenderer: ProcessRequestLoop() - Requests: {requestQueue.Count}");

                // Respond to our requests one at a time
                if (requestQueue.TryDequeue(out var request))
                {
                    if (MLSpectator.Instance.Audio.Paused)
                        MLSpectator.Instance.Audio.Paused = false;

                    // Mark the time of our response
                    timeOfLastRequest = Time.time;

                    // Generate and send our render
                    SendRender(request);
                    lock (requestPool)
                        requestPool.Release(request);
                }

                // If we haven't had any requets in a while, check our connection with a heartbeat
                else if (Time.time - timeOfLastRequest > 0.5f)
                {
                    if (!MLSpectator.Instance.Audio.Paused)
                        MLSpectator.Instance.Audio.Paused = true;

                    // Make a new ping as appropriate
                    if (ping == null || ping.isDone)
                    {
                        if (ping != null)
                        {
                            ping.DestroyPing();
                            ping = null;
                        }

                        ping = new Ping(Host.Instance.RemoteAddress.ToString());
                    }

                    // If we haven't gotten a ping back in 0.5s, then disconnect
                    else
                    {
                        Host.Instance.enabled = false;
                        Host.Instance.enabled = true;
                        ping.DestroyPing();
                        ping = null;
                    }

                    //Debug.Log("MLSpectator: ProcessRequestLoop() - Sending heartbeat");
                    Host.Instance.Send(Heartbeat);

                    // Mark the time of our response
                    timeOfLastRequest = Time.time;
                }
                
                yield return null;
            }
            running = false;

            Debug.Log("MLSpectatorRenderer: ProcessRequestLoop() - End");
        }

        private void SendRender(MLSpectatorRequest request)
        {
            // Update our transform
            renderCam.transform.SetLocalPositionAndRotation(request.position, request.rotation);

            // Update our projection matrix and depth screen
            if (renderCam.projectionMatrix != request.projectionMatrix)
            {
                renderCam.projectionMatrix = request.projectionMatrix;

                // Calculate our depth screen size to completely cover the camera view
                float distance = depthScreen.transform.localPosition.z;
                float tanFov = 1.0f / renderCam.projectionMatrix.m11;
                float aspect = (1.0f / renderCam.projectionMatrix.m00) / tanFov;
                var height = 2.0f * distance * tanFov;
                var width = height * aspect;
                depthScreen.transform.localScale = new Vector3(width, -height, 1);

                depthMat.SetFloat("_Near", renderCam.nearClipPlane);
                depthMat.SetFloat("_Far", renderCam.farClipPlane);
            }

            // Apply depth data if our request came with any
            if (request.HasDepth && depthTex)
            {
                depthTex.LoadRawTextureData(request.DecompressedData);
                depthTex.Apply();

                Graphics.Blit(depthTex, depthRender, depthUpscale);
            }

            // Update our msaa level
            if (request.msaaLevel != render.antiAliasing)
            {
                // Save our old resolution
                Resolution res = new Resolution() { width = render.width, height = render.height };

                // Remove our old texture from our render camera
                renderCam.targetTexture = null;

                // Remove our old texture from our editor preview
                if (editorPreview != null) editorPreview.SetPreview(null);

                // Get rid of our old texture
                render.Release();
                Destroy(render);
                render = null;

                // Create a new texture with the requested msaa level
                render = new RenderTexture(res.width, res.height, depthTex ? 32 : 0, RenderTextureFormat.ARGB32, 0);
                render.antiAliasing = request.msaaLevel;
                render.Create();

                // Assign the texture to our render camera
                renderCam.targetTexture = render;

                // Assign the texture to our editor preview
                if (editorPreview != null) editorPreview.SetPreview(render);
            }

            // Copy our main camera's culling mask so we see what it sees
            if (!MLSpectator.Instance.UseCustomLayerMask)
                renderCam.cullingMask = MLSpectator.Instance.mainCam.cullingMask;
            else renderCam.cullingMask = MLSpectator.Instance.CustomLayerMask;

            // Clear our target texture before rendering (prevents glitching)
            Graphics.Blit(Texture2D.blackTexture, render);

            // If we have our own light, enable it
            if (myLight != null) myLight.enabled = true;
            // Disable all lights in our scene
            foreach (var light in sceneLights) light.enabled = false;

            // Perform custom render events
            onRenderRequest?.Invoke(renderCam);

            // Set our frame number
            splitMat.SetInt("_FrameNum", (int)request.id);

            // Capture the scene with our render camera
            renderCam.Render();

            encoder.Push(render, splitMat);

            // If we have our own light, disable it
            if (myLight != null) myLight.enabled = false;
            // Enable all lights in our scene
            foreach (var light in sceneLights) light.enabled = true;
        }

        private void Reset()
        {
            if (!initialized) return;

            Debug.Log("MLSpectatorRenderer: Reset()");

            // Unregister from relevant messages
            Host.Instance.UnregisterCallback(RenderRequest, HandleRequest);
            Host.Instance.UnregisterCallback(RequestKeyFrame, HandleKeyFrameRequest);

            // Start checking our render queue
            StopCoroutine(ProcessRequestLoop());

            // Hide our depth screen
            depthScreen.SetActive(false);

            // Remove target texture from render camera
            renderCam.targetTexture = null;

            // Dispose of our editor preview
            if (editorPreview != null) editorPreview.Dispose();

            // Destroy our render texture
            if (render != null)
            {
                render.Release();
                Destroy(render);
                render = null;
            }

            // Destroy our encoder
            if (encoder != null)
            {
                encoder.Dispose();
                encoder = null;
            }

            // Destroy our depth texture
            if (depthTex != null)
            {
                Destroy(depthTex);
                depthTex = null;
            }

            // Destroy our depth render
            if (depthRender != null)
            {
                depthMat?.SetTexture("_Depth", Texture2D.whiteTexture);
                depthRender.Release();
                Destroy(depthRender);
                depthRender = null;
            }

            // Destroy our upscale material
            if (depthUpscale != null)
            {
                Destroy(depthUpscale);
                depthUpscale = null;
            }

            // Clear all requests
            requestQueue.Clear();
            lock (requestPool)
                requestPool?.Clear();
            requestPool = null;

            // Mark that we are no longer initialized
            initialized = false;

            // Mark that we are no longer running
            running = false;
        }

        #endregion

        #region Public methods

        public void Initialize(byte[] data, bool renderPreview)
        {
            // If we have already initialized, answer and break
            if (initialized)
            {
                Debug.Log("MLSpectatorRenderer: Initialize() - Already initialized");
                return;
            }
            initialized = true;

            Debug.Log("MLSpectatorRenderer: Initialize() - Initializing");

            // Extract all necessary information
            int msaaLevel = 0;
            Resolution colorResolution, depthResolution;
            Pose startPos;
            using (MemoryStream stream = new MemoryStream(data))
            {
                colorResolution = StreamHelper.ReadResolution(stream);
                depthResolution = StreamHelper.ReadResolution(stream);
                msaaLevel = StreamHelper.ReadInt32(stream);
                startPos = StreamHelper.ReadPose(stream);
            }

            // Set our pose
            transform.SetPositionAndRotation(startPos.position, startPos.rotation);

            // Create a request pool (number of depth pixels may be zero)
            var numDepthPixels = depthResolution.width * depthResolution.height;
            requestPool = new LinkedPool<MLSpectatorRequest>(() => new(numDepthPixels), null, null, null, true, maxRequests);

            // If our depth resolution is valid
            if (depthResolution.width > 2 && depthResolution.height > 2)
            {
                // Create a depth texture for rendering
                depthTex = new Texture2D(depthResolution.width, depthResolution.height, TextureFormat.R8, false);
                depthTex.filterMode = FilterMode.Bilinear;

                // Create a render texture for upscaling our depth texture
                depthRender = new RenderTexture(colorResolution.width, colorResolution.height, 0, RenderTextureFormat.R8, 0);
                depthRender.Create();

                // Clear our depth render
                Graphics.Blit(Texture2D.whiteTexture, depthRender);

                // Create material for upscaling our depth
                depthUpscale = new Material(Shader.Find("Depth/Upscale"));
                
                // Assign our depth render to our depth material
                depthMat?.SetTexture("_Depth", depthRender);

                // Activate our depth screen
                depthScreen.SetActive(true);
            }

            // If our color resolution is valid
            if (colorResolution.width > 2 && colorResolution.height > 2)
            {
                // Create our render texture
                render = new RenderTexture(colorResolution.width, colorResolution.height, depthTex ? 32 : 0, RenderTextureFormat.ARGB32, 0);
                render.antiAliasing = msaaLevel;
                render.Create();

                // Assign the texture to our render camera
                renderCam.targetTexture = render;

                // Pack alpha into green channel if our phone version can handle it
                splitMat.SetInt("_PackGreen", MLSpectator.Instance.PackGreen ? 1 : 0);
                
                // Create our encoder and attach it to our camera
                encoder = H264Encoder.Create(colorResolution.width, colorResolution.height * 2, 5000000);
                encoder.onFrameEncoded += (data, length) => Host.Instance.Send(new Message(RenderResponse, data, length));
                //encoder.AttachToCamera(renderCam, splitMat);

#if UNITY_EDITOR_WIN || UNITY_EDITOR_OSX
                if (renderPreview) editorPreview = new MLSpectatorPreview(render);
#endif
            }

            // Register to relevant messages
            Host.Instance.RegisterCallback(RenderRequest, HandleRequest);
            Host.Instance.RegisterCallback(RequestKeyFrame, HandleKeyFrameRequest);

            // Set time of last request to now
            timeOfLastRequest = Time.time;

            // Start checking our render queue
            StartCoroutine(ProcessRequestLoop());
        }

        #endregion
    }
}