/************************************************************************************

Copyright   :   Copyright 2014 Oculus VR, LLC. All Rights reserved.

Licensed under the Oculus VR Rift SDK License Version 3.3 (the "License");
you may not use the Oculus VR Rift SDK except in compliance with the License,
which is provided at the time of installation or download, or which
otherwise accompanies this software in either electronic or hard copy form.

You may obtain a copy of the License at

http://www.oculus.com/licenses/LICENSE-3.3

Unless required by applicable law or agreed to in writing, the Oculus VR SDK
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

************************************************************************************/

using UnityEngine;
using System;
using System.Collections;
using System.Runtime.InteropServices;
using VR = UnityEngine.VR;

/// <summary>
/// Add OVROverlay script to an object with an optional mesh primitive
/// rendered as a TimeWarp overlay instead by drawing it into the eye buffer.
/// This will take full advantage of the display resolution and avoid double
/// resampling of the texture.
/// 
/// If the texture is dynamically generated, as for an interactive GUI or
/// animation, it must be explicitly triple buffered to avoid flickering
/// when it is referenced asynchronously by TimeWarp, check OVRRTOverlayConnector.cs for triple buffers design
/// 
/// We support 3 types of Overlay shapes right now
///		1. Quad : This is most common overlay type , you render a quad in Timewarp space.
///		2. Cylinder: [Mobile Only][Experimental], Display overlay as partial surface of a cylinder
///			* The cylinder's center will be your game object's center
///			* We encoded the cylinder's parameters in transform.scale, 
///				**[scale.z] is the radius of the cylinder
///				**[scale.y] is the height of the cylinder
///				**[scale.x] is the length of the arc of cylinder
///		* Limitations
///				**Only the half of the cylinder can be displayed, which means the arc angle has to be smaller than 180 degree,  [scale.x] / [scale.z] <= PI
///				**Your camera has to be inside of the inscribed sphere of the cylinder, the overlay will be faded out automatically when the camera is close to the inscribed sphere's surface.
///				**Translation only works correctly with vrDriver 1.04 or above
///		3. Cubemap: Display overlay as a cube map
///		4. OffcenterCubemap: [Mobile Only] Display overlay as a cube map with a texture coordinate offset
///			* The actually sampling will looks like [color = texture(cubeLayerSampler, normalize(direction) + offset)] instead of [color = texture( cubeLayerSampler, direction )]
///			* The extra center offset can be feed from transform.position
///			* Note: if transform.position's magnitude is greater than 1, which will cause some cube map pixel always invisible 
///					Which is usually not what people wanted, we don't kill the ability for developer to do so here, but will warn out.
/// </summary>
public class OVROverlay : MonoBehaviour
{
#region Interface

	/// <summary>
	/// Determines the on-screen appearance of a layer.
	/// </summary>
	public enum OverlayShape
	{
		Quad = OVRPlugin.OverlayShape.Quad,
		Cylinder = OVRPlugin.OverlayShape.Cylinder,
		Cubemap = OVRPlugin.OverlayShape.Cubemap,
		OffcenterCubemap = OVRPlugin.OverlayShape.OffcenterCubemap,
	}

	/// <summary>
	/// Whether the layer appears behind or infront of other content in the scene.
	/// </summary>
	public enum OverlayType
	{
		None,
		Underlay,
		Overlay,
	};

	/// <summary>
	/// Specify overlay's type
	/// </summary>
	public OverlayType currentOverlayType = OverlayType.Overlay;

	/// <summary>
	/// If true, the texture's content is copied to the compositor each frame.
	/// </summary>
	public bool isDynamic = false;

	/// <summary>
	/// Specify overlay's shape
	/// </summary>
	public OverlayShape currentOverlayShape = OverlayShape.Quad;
	private OverlayShape prevOverlayShape = OverlayShape.Quad;

	/// <summary>
	/// Try to avoid setting texture frequently when app is running, texNativePtr updating is slow since rendering thread synchronization
	/// Please cache your nativeTexturePtr and use  OverrideOverlayTextureInfo
	/// </summary>
	public Texture[] textures = new Texture[] { null, null };

	/// <summary>
	/// Use this function to set texture and texNativePtr when app is running 
	/// GetNativeTexturePtr is a slow behavior, the value should be pre-cached 
	/// </summary>
	public void OverrideOverlayTextureInfo(Texture srcTexture, IntPtr nativePtr, VR.VRNode node)
	{
		int index = (node == VR.VRNode.RightEye) ? 1 : 0;

		if (textures.Length <= index)
			return;

		stageCount = 3;
		CreateLayerTextures(true, true, new OVRPlugin.Sizei() {w = srcTexture.width, h = srcTexture.height}, false);

		textures[index] = srcTexture;
		layerTextures[index].appTexture = srcTexture;
		layerTextures[index].appTexturePtr = nativePtr;
	}

#if UNITY_ANDROID && !UNITY_EDITOR
	internal const int maxInstances = 3;
#else
	internal const int maxInstances = 15;
#endif

	internal static OVROverlay[] instances = new OVROverlay[maxInstances];

#endregion

	private static Material premultiplyMaterial;

	private OVRPlugin.LayerLayout layout = OVRPlugin.LayerLayout.Mono;

	private struct LayerTexture {
		public Texture appTexture;
		public IntPtr appTexturePtr;
		public Texture[] swapChain;
		public IntPtr[] swapChainPtr;
	};
	private LayerTexture[] layerTextures;

	private OVRPlugin.LayerDesc layerDesc;
	private int stageCount = -1;

	private int layerIndex = -1; // Controls the composition order based on wake-up time.

	private int layerId = 0; // The layer's internal handle in the compositor.
	private GCHandle layerIdHandle;
	private IntPtr layerIdPtr = IntPtr.Zero;

	private int frameIndex = 0;
	private int prevFrameIndex = -1;

	private Renderer rend;

	private int texturesPerStage { get { return (layout == OVRPlugin.LayerLayout.Stereo) ? 2 : 1; } }

	private bool CreateLayer(int mipLevels, int sampleCount, OVRPlugin.EyeTextureFormat etFormat, int flags, OVRPlugin.Sizei size, OVRPlugin.OverlayShape shape)
	{
		if (!layerIdHandle.IsAllocated || layerIdPtr == IntPtr.Zero)
		{
			layerIdHandle = GCHandle.Alloc(layerId, GCHandleType.Pinned);
			layerIdPtr = layerIdHandle.AddrOfPinnedObject();
		}

		if (layerIndex == -1)
		{
			for (int i = 0; i < maxInstances; ++i)
			{
				if (instances[i] == null || instances[i] == this)
				{
					layerIndex = i;
					instances[i] = this;
					break;
				}
			}
		}

		bool needsSetup = (
			layerDesc.MipLevels != mipLevels ||
			layerDesc.SampleCount != sampleCount ||
			layerDesc.Format != etFormat ||
			layerDesc.LayerFlags != flags ||
			!layerDesc.TextureSize.Equals(size) ||
			layerDesc.Shape != shape);

		if (!needsSetup)
			return false;

		OVRPlugin.LayerDesc desc = OVRPlugin.CalculateLayerDesc(shape, layout, size, mipLevels, sampleCount, etFormat, flags);
		OVRPlugin.EnqueueSetupLayer(desc, layerIdPtr);
		layerId = (int)layerIdHandle.Target;

		if (layerId > 0)
		{
			layerDesc = desc;
			stageCount = OVRPlugin.GetLayerTextureStageCount(layerId);
		}

		return true;
	}

	private bool CreateLayerTextures(bool isSrgb, bool useMipmaps, OVRPlugin.Sizei size, bool isHdr)
	{
		bool needsCopy = false;

		if (stageCount <= 0)
			return false;

		// For newer SDKs, blit directly to the surface that will be used in compositing.

		if (layerTextures == null)
		{
			frameIndex = 0;
			layerTextures = new LayerTexture[texturesPerStage];
		}

		for (int eyeId = 0; eyeId < texturesPerStage; ++eyeId)
		{
			if (layerTextures[eyeId].swapChain == null)
				layerTextures[eyeId].swapChain = new Texture[stageCount];
			
			if (layerTextures[eyeId].swapChainPtr == null)
				layerTextures[eyeId].swapChainPtr = new IntPtr[stageCount];

			for (int stage = 0; stage < stageCount; ++stage)
			{
				Texture sc = layerTextures[eyeId].swapChain[stage];
				IntPtr scPtr = layerTextures[eyeId].swapChainPtr[stage];

				if (sc != null && scPtr != IntPtr.Zero)
					continue;

				if (scPtr == IntPtr.Zero)
					scPtr = OVRPlugin.GetLayerTexture(layerId, stage, (OVRPlugin.Eye)eyeId);

				if (scPtr == IntPtr.Zero)
					continue;

				var txFormat = (isHdr) ? TextureFormat.RGBAHalf : TextureFormat.RGBA32;

				if (currentOverlayShape != OverlayShape.Cubemap && currentOverlayShape != OverlayShape.OffcenterCubemap)
					sc = Texture2D.CreateExternalTexture(size.w, size.h, txFormat, useMipmaps, isSrgb, scPtr);
#if UNITY_2017_1_OR_NEWER
				else
					sc = Cubemap.CreateExternalTexture(size.w, txFormat, useMipmaps, scPtr);
#endif

				layerTextures[eyeId].swapChain[stage] = sc;
				layerTextures[eyeId].swapChainPtr[stage] = scPtr;

				needsCopy = true;
			}
		}

		return needsCopy;
	}

	private void DestroyLayerTextures()
	{
		for (int eyeId = 0; layerTextures != null && eyeId < texturesPerStage; ++eyeId)
		{
			if (layerTextures[eyeId].swapChain != null)
			{
				for (int stage = 0; stage < stageCount; ++stage)
					DestroyImmediate(layerTextures[eyeId].swapChain[stage]);
			}
		}

		layerTextures = null;
	}

	private void DestroyLayer()
	{
		if (layerIndex != -1)
		{
			// Turn off the overlay if it was on.
			OVRPlugin.EnqueueSubmitLayer(true, false, IntPtr.Zero, IntPtr.Zero, -1, 0, OVRPose.identity.ToPosef(), Vector3.one.ToVector3f(), layerIndex, (OVRPlugin.OverlayShape)prevOverlayShape);
			instances[layerIndex] = null;
			layerIndex = -1;
		}

		if (layerIdPtr != IntPtr.Zero)
		{
			OVRPlugin.EnqueueDestroyLayer(layerIdPtr);
			layerIdPtr = IntPtr.Zero;
			layerIdHandle.Free();
			layerId = 0;
		}

		layerDesc = new OVRPlugin.LayerDesc();
	}

	private bool LatchLayerTextures()
	{
		for (int i = 0; i < texturesPerStage; ++i)
		{
			if (textures[i] != layerTextures[i].appTexture || layerTextures[i].appTexturePtr == IntPtr.Zero)
			{
				if (textures[i] != null)
				{
					layerTextures[i].appTexturePtr = textures[i].GetNativeTexturePtr();
					if (layerTextures[i].appTexturePtr != IntPtr.Zero)
						layerTextures[i].appTexture = textures[i];
				}
			}

			if (currentOverlayShape == OverlayShape.Cubemap)
			{
				if (textures[i] as Cubemap == null)
				{
					Debug.LogError("Need Cubemap texture for cube map overlay");
					return false;
				}
			}
		}

#if !UNITY_ANDROID || UNITY_EDITOR
		if (currentOverlayShape == OverlayShape.OffcenterCubemap)
		{
			Debug.LogWarning("Overlay shape " + currentOverlayShape + " is not supported on current platform");
			return false;
		}
#endif

		if (layerTextures[0].appTexture == null || layerTextures[0].appTexturePtr == IntPtr.Zero)
			return false;

		return true;
	}

	private OVRPlugin.LayerDesc GetCurrentLayerDesc()
	{
		OVRPlugin.LayerDesc newDesc = new OVRPlugin.LayerDesc() {
			Format = OVRPlugin.EyeTextureFormat.R8G8B8A8_sRGB,
			LayerFlags = (int)OVRPlugin.LayerFlags.TextureOriginAtBottomLeft,
			Layout = layout,
			MipLevels = 1,
			SampleCount = 1,
			Shape = (OVRPlugin.OverlayShape)currentOverlayShape,
			TextureSize = new OVRPlugin.Sizei() { w = textures[0].width, h = textures[0].height }
		};

		var tex2D = textures[0] as Texture2D;
		if (tex2D != null)
		{
			if (tex2D.format == TextureFormat.RGBAHalf || tex2D.format == TextureFormat.RGBAFloat)
				newDesc.Format = OVRPlugin.EyeTextureFormat.R16G16B16A16_FP;
			newDesc.MipLevels = tex2D.mipmapCount;
		}

		var texCube = textures[0] as Cubemap;
		if (texCube != null)
			newDesc.MipLevels = texCube.mipmapCount;

		var rt = textures[0] as RenderTexture;
		if (rt != null)
		{
			isDynamic = true;

			newDesc.SampleCount = rt.antiAliasing;

			if (rt.format == RenderTextureFormat.ARGBHalf)
				newDesc.Format = OVRPlugin.EyeTextureFormat.R16G16B16A16_FP;
		}

		return newDesc;
	}

	private bool PopulateLayer(int mipLevels, bool isHdr, OVRPlugin.Sizei size, int sampleCount)
	{
		bool ret = false;

		RenderTextureFormat rtFormat = (isHdr) ? RenderTextureFormat.ARGBHalf : RenderTextureFormat.ARGB32;

		for (int eyeId = 0; eyeId < texturesPerStage; ++eyeId)
		{
			int dstElement = (layout == OVRPlugin.LayerLayout.Array) ? eyeId : 0;

			int stage = frameIndex % stageCount;
			Texture et = layerTextures[eyeId].swapChain[stage];
			if (et == null)
				continue;

			for (int mip = 0; mip < mipLevels; ++mip)
			{
#if UNITY_2017_1_1 || UNITY_2017_2_OR_NEWER
				int width = size.w >> mip;
				if (width < 1) width = 1;
				int height = size.h >> mip;
				if (height < 1) height = 1;

				RenderTextureDescriptor descriptor = new RenderTextureDescriptor(width, height, rtFormat, 0);
				descriptor.msaaSamples = sampleCount;
				descriptor.useMipMap = true;
				descriptor.autoGenerateMips = false;

				var tempRTDst = RenderTexture.GetTemporary(descriptor);
#else
				var tempRTDst = RenderTexture.GetTemporary(size.w >> mip, size.h >> mip, 0, rtFormat, RenderTextureReadWrite.Default, sampleCount);
#endif

				if (!tempRTDst.IsCreated())
					tempRTDst.Create();

				tempRTDst.DiscardContents();

				if (currentOverlayShape != OverlayShape.Cubemap && currentOverlayShape != OverlayShape.OffcenterCubemap)
				{
#if UNITY_ANDROID && !UNITY_EDITOR
					if (((textures[eyeId] as Cubemap) != null)
						&& ((et as Cubemap) != null)
						&& ((textures[eyeId] as Cubemap).format == (et as Cubemap).format))
					{
						Graphics.CopyTexture(textures[eyeId], 0, mip, et, 0, mip);
					}
					else
					{
						Graphics.Blit(textures[eyeId], tempRTDst); //Resolve, decompress, swizzle, etc not handled by simple CopyTexture.
						Graphics.CopyTexture(tempRTDst, 0, 0, et, dstElement, mip);
					}
#else
					// The PC compositor uses premultiplied alpha, so multiply it here.
					Graphics.Blit(textures[eyeId], tempRTDst, premultiplyMaterial);
					Graphics.CopyTexture(tempRTDst, 0, 0, et, dstElement, mip);
#endif
				}
#if UNITY_2017_1_OR_NEWER
				else // Cubemap
				{
					var tempRTSrc = RenderTexture.GetTemporary(size.w >> mip, size.h >> mip, 0, rtFormat, RenderTextureReadWrite.Default, sampleCount);

					if (!tempRTSrc.IsCreated())
						tempRTSrc.Create();

					tempRTSrc.DiscardContents();

					for (int face = 0; face < 6; ++face)
					{
#if UNITY_ANDROID && !UNITY_EDITOR
						if ((textures[eyeId] as Cubemap).format == (et as Cubemap).format)
						{
							Graphics.CopyTexture(textures[eyeId], face, mip, et, 0, mip);
						}
						else
						{
							//HACK: It would be much more efficient to blit directly from textures[eyeId] to et, but Unity's API doesn't support that.
							//Suggest using a native plugin to render directly to a cubemap layer for 360 video, etc.
							Graphics.CopyTexture(textures[eyeId], face, mip, tempRTSrc, 0, 0);
							Graphics.Blit(tempRTSrc, tempRTDst);
							Graphics.CopyTexture(tempRTDst, 0, 0, et, face, mip);
						}
#else
						//HACK: It would be much more efficient to blit directly from textures[eyeId] to et, but Unity's API doesn't support that.
						//Suggest using a native plugin to render directly to a cubemap layer for 360 video, etc.
						Graphics.CopyTexture(textures[eyeId], face, mip, tempRTSrc, 0, 0);
						// The PC compositor uses premultiplied alpha, so multiply it here.
						Graphics.Blit(tempRTSrc, tempRTDst, premultiplyMaterial);
						Graphics.CopyTexture(tempRTDst, 0, 0, et, face, mip);
#endif
					}
					RenderTexture.ReleaseTemporary(tempRTSrc);
				}
#endif
				RenderTexture.ReleaseTemporary(tempRTDst);

				ret = true;
			}
		}

		return ret;
	}

	private bool SubmitLayer(bool overlay, bool headLocked, OVRPose pose, Vector3 scale)
	{
		int rightEyeIndex = (texturesPerStage >= 2) ? 1 : 0;
		bool isOverlayVisible = OVRPlugin.EnqueueSubmitLayer(overlay, headLocked, layerTextures[0].appTexturePtr, layerTextures[rightEyeIndex].appTexturePtr, layerId, frameIndex, pose.flipZ().ToPosef(), scale.ToVector3f(), layerIndex, (OVRPlugin.OverlayShape)currentOverlayShape);
		if (isDynamic)
			++frameIndex;

		prevOverlayShape = currentOverlayShape;

		return isOverlayVisible;
	}

#region Unity Messages

	void Awake()
	{
		Debug.Log("Overlay Awake");

		if (premultiplyMaterial == null)
			premultiplyMaterial = new Material(Shader.Find("Oculus/Alpha Premultiply"));

		rend = GetComponent<Renderer>();

		if (textures.Length == 0)
			textures = new Texture[] { null };

		// Backward compatibility
		if (rend != null && textures[0] == null)
			textures[0] = rend.material.mainTexture;

#if UNITY_ANDROID && !UNITY_EDITOR
		if (textures.Length == 2 && textures[1] != null)
			layout = OVRPlugin.LayerLayout.Stereo;
#endif
	}

	void OnEnable()
	{
		if (!OVRManager.isHmdPresent)
		{
			enabled = false;
			return;
		}
	}

	void OnDisable()
	{
		DestroyLayerTextures();
		DestroyLayer();
	}

	void OnDestroy()
	{
		DestroyLayerTextures();
		DestroyLayer();
	}

	bool ComputeSubmit(ref OVRPose pose, ref Vector3 scale, ref bool overlay, ref bool headLocked)
	{
		overlay = (currentOverlayType == OverlayType.Overlay);
		headLocked = false;
		for (var t = transform; t != null && !headLocked; t = t.parent)
			headLocked |= (t == Camera.current.transform);

		pose = (headLocked) ? transform.ToHeadSpacePose() : transform.ToTrackingSpacePose();
		scale = transform.lossyScale;
		for (int i = 0; i < 3; ++i)
			scale[i] /= Camera.current.transform.lossyScale[i];

		if (currentOverlayShape == OverlayShape.Cubemap)
			pose.position = Camera.current.transform.position;

		// Pack the offsetCenter directly into pose.position for offcenterCubemap
		if (currentOverlayShape == OverlayShape.OffcenterCubemap)
		{
			pose.position = transform.position;
			if (pose.position.magnitude > 1.0f)
			{
				Debug.LogWarning("Your cube map center offset's magnitude is greater than 1, which will cause some cube map pixel always invisible .");
				return false;
			}
		}

		// Cylinder overlay sanity checking
		if (currentOverlayShape == OverlayShape.Cylinder)
		{
			float arcAngle = scale.x / scale.z / (float)Math.PI * 180.0f;
			if (arcAngle > 180.0f)
			{
				Debug.LogWarning("Cylinder overlay's arc angle has to be below 180 degree, current arc angle is " + arcAngle + " degree." );
				return false;
			}
		}

		return true;
	}

	void OnRenderObject()
	{
		// The overlay must be specified every eye frame, because it is positioned relative to the
		// current head location.  If frames are dropped, it will be time warped appropriately,
		// just like the eye buffers.
		if (!Camera.current.CompareTag("MainCamera") || Camera.current.cameraType != UnityEngine.CameraType.Game)
			return;

		if (currentOverlayType == OverlayType.None || textures.Length < texturesPerStage)
			return;

		// Don't submit the same frame twice.
		if (Time.frameCount <= prevFrameIndex)
			return;
		prevFrameIndex = Time.frameCount;

		OVRPose pose = OVRPose.identity;
		Vector3 scale = Vector3.one;
		bool overlay = false;
		bool headLocked = false;
		if (!ComputeSubmit(ref pose, ref scale, ref overlay, ref headLocked))
			return;

		OVRPlugin.LayerDesc newDesc = GetCurrentLayerDesc();
		bool isHdr = (newDesc.Format == OVRPlugin.EyeTextureFormat.R16G16B16A16_FP);

		bool createdLayer = CreateLayer(newDesc.MipLevels, newDesc.SampleCount, newDesc.Format, newDesc.LayerFlags, newDesc.TextureSize, newDesc.Shape);

		if (layerIndex == -1 || layerId <= 0)
			return;
	
		bool isSrgb = (newDesc.Format == OVRPlugin.EyeTextureFormat.B8G8R8A8_sRGB || newDesc.Format == OVRPlugin.EyeTextureFormat.R8G8B8A8_sRGB);
		bool useMipmaps = (newDesc.MipLevels > 1);

		createdLayer |= CreateLayerTextures(isSrgb, useMipmaps, newDesc.TextureSize, isHdr);

		if (layerTextures[0].appTexture as RenderTexture != null)
			isDynamic = true;

		if (!isDynamic && !createdLayer)
			return;

		if (!LatchLayerTextures())
			return;
		
		if (!PopulateLayer(newDesc.MipLevels, isHdr, newDesc.TextureSize, newDesc.SampleCount))
			return;

		bool isOverlayVisible = SubmitLayer(overlay, headLocked, pose, scale);

		// Backward compatibility: show regular renderer if overlay isn't visible.
		if (rend)
			rend.enabled = !isOverlayVisible;
	}

#endregion
}
