/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * Licensed under the Oculus SDK License Agreement (the "License");
 * you may not use the Oculus 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
 *
 * https://developer.oculus.com/licenses/oculussdk/
 *
 * Unless required by applicable law or agreed to in writing, the Oculus 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.Collections;
using System.IO;

/// <summary>
/// Helper script for capture cubemap and save it into PNG or JPG file
/// </summary>

/// <description>
/// How it works:
/// 1) This script can be attached to a existing game object, you can also use prefab Assets\OVR\Prefabs\OVRCubemapCaptureProbe
///	There are 2 ways to trigger a capture if you attached this script to a game object.
///		* Automatic capturing: if [autoTriggerAfterLaunch] is true, a automatic capturing will be triggered after [autoTriggerDelay] seconds.
///		* Keyboard trigger: press key [triggeredByKey], a capturing will be triggered.
/// 2) If you like to trigger the screen capture in your code logic, just call static function [OVRCubemapCapture.TriggerCubemapCapture] with proper input arguments.
/// </description>

public class OVRCubemapCapture : MonoBehaviour
{
	/// <summary>
	/// Enable the automatic screenshot trigger, which will capture a cubemap after autoTriggerDelay (seconds)
	/// </summary>
	public bool autoTriggerAfterLaunch = true;
	public float autoTriggerDelay = 1.0f;
	private float autoTriggerElapse = 0.0f;

	/// <summary>
	/// Trigger cubemap screenshot if user pressed key triggeredByKey
	/// </summary>
	public KeyCode triggeredByKey = KeyCode.F8;

	/// <summary>
	/// The complete file path for saving the cubemap screenshot, including the filename and extension
	/// if pathName is blank, screenshots will be saved into %USERPROFILE%\Documents\OVR_ScreenShot360
	/// </summary>
	public string pathName;

	/// <summary>
	/// The cube face resolution
	/// </summary>
	public int cubemapSize = 2048;

	// Update is called once per frame
	void Update()
	{
		// Trigger after autoTriggerDelay
		if (autoTriggerAfterLaunch)
		{
			autoTriggerElapse += Time.deltaTime;
			if (autoTriggerElapse >= autoTriggerDelay)
			{
				autoTriggerAfterLaunch = false;
				TriggerCubemapCapture(transform.position, cubemapSize, pathName);
			}
		}

		//todo: enable for Unity Input System
#if ENABLE_LEGACY_INPUT_MANAGER
		// Trigger by press triggeredByKey
		if ( Input.GetKeyDown( triggeredByKey ) )
		{
			TriggerCubemapCapture(transform.position, cubemapSize, pathName);
		}
#endif
	}

	/// <summary>
	/// Generate unity cubemap at specific location and save into JPG/PNG
	/// </summary>
	/// <description>
	/// Default save folder: your app's persistentDataPath
	/// Default file name: using current time OVR_hh_mm_ss.png
	/// Note1: this will take a few seconds to finish
	/// Note2: if you only want to specify path not filename, please end [pathName] with "/"
	/// </description>

	public static void TriggerCubemapCapture(Vector3 capturePos, int cubemapSize = 2048, string pathName = null)
	{
		GameObject ownerObj = new GameObject("CubemapCamera", typeof(Camera));
		ownerObj.hideFlags = HideFlags.HideAndDontSave;
		ownerObj.transform.position = capturePos;
		ownerObj.transform.rotation = Quaternion.identity;
		Camera camComponent = ownerObj.GetComponent<Camera>();
		camComponent.farClipPlane = 10000.0f;
		camComponent.enabled = false;

		Cubemap cubemap = new Cubemap(cubemapSize, TextureFormat.RGB24, false);
		RenderIntoCubemap(camComponent, cubemap);
		SaveCubemapCapture(cubemap, pathName);
		DestroyImmediate(cubemap);
		DestroyImmediate(ownerObj);
	}


	public static void RenderIntoCubemap(Camera ownerCamera, Cubemap outCubemap)
	{
		int width = (int)outCubemap.width;
		int height = (int)outCubemap.height;

		CubemapFace[] faces = new CubemapFace[] { CubemapFace.PositiveX, CubemapFace.NegativeX, CubemapFace.PositiveY, CubemapFace.NegativeY, CubemapFace.PositiveZ, CubemapFace.NegativeZ };
		Vector3[] faceAngles = new Vector3[] { new Vector3(0.0f, 90.0f, 0.0f), new Vector3(0.0f, -90.0f, 0.0f), new Vector3(-90.0f, 0.0f, 0.0f), new Vector3(90.0f, 0.0f, 0.0f), new Vector3(0.0f, 0.0f, 0.0f), new Vector3(0.0f, 180.0f, 0.0f) };

		// Backup states
		RenderTexture backupRenderTex = RenderTexture.active;
		float backupFieldOfView = ownerCamera.fieldOfView;
		float backupAspect = ownerCamera.aspect;
		Quaternion backupRot = ownerCamera.transform.rotation;
		//RenderTexture backupRT = ownerCamera.targetTexture;

		// Enable 8X MSAA
		RenderTexture faceTexture = new RenderTexture(width, height, 24);
		faceTexture.antiAliasing = 8;
		faceTexture.dimension = UnityEngine.Rendering.TextureDimension.Tex2D;
		faceTexture.hideFlags = HideFlags.HideAndDontSave;

		// For intermediate saving
		Texture2D swapTex = new Texture2D(width, height, TextureFormat.RGB24, false);
		swapTex.hideFlags = HideFlags.HideAndDontSave;

		// Capture 6 Directions
		ownerCamera.targetTexture = faceTexture;
		ownerCamera.fieldOfView = 90;
		ownerCamera.aspect = 1.0f;

		Color[] mirroredPixels = new Color[swapTex.height * swapTex.width];
		for (int i = 0; i < faces.Length; i++)
		{
			ownerCamera.transform.eulerAngles = faceAngles[i];
			ownerCamera.Render();
			RenderTexture.active = faceTexture;
			swapTex.ReadPixels(new Rect(0, 0, width, height), 0, 0);

			// Mirror vertically to meet the standard of unity cubemap
			Color[] OrignalPixels = swapTex.GetPixels();
			for (int y1 = 0; y1 < height; y1++)
			{
				for (int x1 = 0; x1 < width; x1++)
				{
					mirroredPixels[y1 * width + x1] = OrignalPixels[((height - 1 - y1) * width) + x1];
				}
			};
			outCubemap.SetPixels(mirroredPixels, faces[i]);
		}

		outCubemap.SmoothEdges();

		// Restore states
		RenderTexture.active = backupRenderTex;
		ownerCamera.fieldOfView = backupFieldOfView;
		ownerCamera.aspect = backupAspect;
		ownerCamera.transform.rotation = backupRot;
		ownerCamera.targetTexture = backupRenderTex;

		DestroyImmediate(swapTex);
		DestroyImmediate(faceTexture);

		}


	/// <summary>
	/// Save unity cubemap into NPOT 6x1 cubemap/texture atlas in the following format PX NX PY NY PZ NZ
	/// </summary>
	/// <description>
	/// Supported format: PNG/JPG
	/// Default file name: using current time OVR_hh_mm_ss.png
	/// </description>

	public static bool SaveCubemapCapture(Cubemap cubemap, string pathName = null)
	{
		string fileName;
		string dirName;
		int width = cubemap.width;
		int height = cubemap.height;
		int x = 0;
		int y = 0;
		bool saveToPNG = true;

		if (string.IsNullOrEmpty(pathName))
		{
			dirName = Application.persistentDataPath + "/OVR_ScreenShot360/";
			fileName = null;
		}
		else
		{
			dirName = Path.GetDirectoryName(pathName);
			fileName = Path.GetFileName(pathName);

			if (dirName[dirName.Length - 1] != '/' || dirName[dirName.Length - 1] != '\\')
				dirName += "/";
		}

		if (string.IsNullOrEmpty(fileName))
			fileName = "OVR_" + System.DateTime.Now.ToString("hh_mm_ss") + ".png";

		string extName = Path.GetExtension(fileName);
		if (extName == ".png")
		{
			saveToPNG = true;
		}
		else if (extName == ".jpg")
		{
			saveToPNG = false;
		}
		else
		{
			Debug.LogError("Unsupported file format" + extName);
			return false;
		}

		// Validate path
		try
		{
			System.IO.Directory.CreateDirectory(dirName);
		}
		catch (System.Exception e)
		{
			Debug.LogError("Failed to create path " + dirName + " since " + e.ToString());
			return false;
		}


		// Create the new texture
		Texture2D tex = new Texture2D(width * 6, height, TextureFormat.RGB24, false);
		if (tex == null)
		{
			Debug.LogError("[OVRScreenshotWizard] Failed creating the texture!");
			return false;
		}

		// Merge all the cubemap faces into the texture
		// Reference cubemap format: http://docs.unity3d.com/Manual/class-Cubemap.html
		CubemapFace[] faces = new CubemapFace[] { CubemapFace.PositiveX, CubemapFace.NegativeX, CubemapFace.PositiveY, CubemapFace.NegativeY, CubemapFace.PositiveZ, CubemapFace.NegativeZ };
		for (int i = 0; i < faces.Length; i++)
		{
			// get the pixels from the cubemap
			Color[] srcPixels = null;
			Color[] pixels = cubemap.GetPixels(faces[i]);
			// if desired, flip them as they are ordered left to right, bottom to top
			srcPixels = new Color[pixels.Length];
			for (int y1 = 0; y1 < height; y1++)
			{
				for (int x1 = 0; x1 < width; x1++)
				{
					srcPixels[y1 * width + x1] = pixels[((height - 1 - y1) * width) + x1];
				}
			}
			// Copy them to the dest texture
			tex.SetPixels(x, y, width, height, srcPixels);
			x += width;
		}

		try
		{
			// Encode the texture and save it to disk
			byte[] bytes = saveToPNG ? tex.EncodeToPNG() : tex.EncodeToJPG();

			System.IO.File.WriteAllBytes(dirName + fileName, bytes);
			Debug.Log("Cubemap file created " + dirName + fileName);
		}
		catch (System.Exception e)
		{
			Debug.LogError("Failed to save cubemap file since " + e.ToString());
			return false;
		}

		DestroyImmediate(tex);
		return true;
	}

}