﻿using UnityEngine;
using System.Collections;

using System;
using System.IO;

namespace Baidu.VR
{
	using System.Collections.Generic;
#if UNITY_EDITOR
	using UnityEditor;
	[CustomEditor(typeof(VRCubeCapture))]
	public class VRCubeCaptureEditor : Editor
	{
		public override void OnInspectorGUI()
		{
			base.OnInspectorGUI();
			if (GUILayout.Button("Capture"))
			{
				((VRCubeCapture)target).StartCaptureImage();
			}
		}
	}
#endif

	[ExecuteInEditMode]
	[RequireComponent(typeof(Camera))]
    public class VRCubeCapture : MonoBehaviour
    {
#region public variable
		[Header("FilePath")]
		public string _fileName = "Image360";
		public string Folder = "CubePanoCapture";
		public bool customPath = false;
		public string customPathFolder = "D:/";

		[Space(15)]
		public KeyCode CaptureKey = KeyCode.Space;

		public enum VRFormatList { JPG, PNG };
        public VRFormatList ImageFormatType;

        public int FPS = 25;

		public float nearPlane = 0.001f;
		public float farPlane = 20.0f;

        public int resolution = 4096;
        public int resolutionH = 4096;

        public float IPDistance = 0.066f;

        public int jpgQuality = 100;
		
		[Space(15)]
		[Header("RenderTime/Quality Optimization")]
        [Range(1, 16)]
        public int renderQuality = 16;
        TextureFormat texFormat;
        RenderTextureFormat renderTexFormat;

		public bool depth = true;
		public bool useAlpha = false;

        public bool disableTracking = true;

		public bool hQ = false;
        public int steps = 10;
		public float smoothing = 1f;
#endregion

#region private variable
		private int step = 0;
		private bool CanDestroy = true;
		public bool CaptureTrigger{ get;private set;}
		private float restoreFOV;
		//private int qualityTemp;
		public bool waitRenderFinish{ get; private set;}
		private float EnvironmentDistance = 2;

#region Objects
		//		private GameObject renderHead;
		private GameObject rig;
		private GameObject cam;

		private string[] sideNames = new string[] { "f", "l", "r", "d", "b", "u" };

		// F, L, R, D, B, U
		private RenderTexture[] textures = null;

		// F, L, R, D, B, U
		private GameObject[] camObjs = null;

		// F, L, R, D, B, U
		private GameObject[] cloneCameras = null;
		
		public class StepInfo{
			public int CurStep = 0;
			public int TotalStep = 1;
			public float ProgressInd{get{ return TotalStep > 0 ? ((float)CurStep) / TotalStep : 0.0f;}}
		}
		public StepInfo stepInfo = new StepInfo();

#endregion
#endregion

#region private&tools
        void Start()
        {
            Application.runInBackground = true;

            if (!disableTracking)
                UnityEngine.XR.XRSettings.enabled = false;
        }

		void LateUpdate()
		{
			if (Input.GetKeyDown(CaptureKey))
				StartCaptureImage();

			/*switch (step)
			{
				case 1:
					{
						Step1();
					}
					break;
				case 2:
					{
						Step2();
					}
					break;
				case 3:
					{
						Step3();
					}
					break;
				default:
					break;
			}*/
		}

		private void Step1()
		{
			waitRenderFinish = true;
			PreparePano();
			RenderPano();

			step = 2;
		}

		private void Step2()
		{
			if (CaptureTrigger)
			{
				Vector3 angleCorection = gameObject.transform.rotation.eulerAngles;
				angleCorection = new Vector3(0, angleCorection.y, 0);
				gameObject.transform.rotation = Quaternion.Euler(angleCorection);
				RenderVRPanorama();
				print("Panorama Captured");
				CaptureTrigger = false;
			}

			if (CanDestroy)
			{
				Time.captureFramerate = 0;
				step = 3;
			}
		}

		private void Step3()
		{
			DestroyObj(rig);
			foreach (RenderTexture tex in textures)
			{
				RenderTexture.ReleaseTemporary(tex);
			}
			textures = null;

			cam.GetComponent<Camera>().fieldOfView = restoreFOV;

			step = 0;
			waitRenderFinish = false;
		}

		private void RenderPano()
		{
			float aAfactor = resolutionH;
			float aAfactorH = resolution;

			int depthBufferSize;
			if (depth == true)
				depthBufferSize = 24;
			else
				depthBufferSize = 0;

			textures = new RenderTexture[6];
			for (int i = 0; i < 6; i++)
			{
				textures[i] = RenderTexture.GetTemporary(resolution, resolutionH, depthBufferSize, renderTexFormat);
				cloneCameras[i].GetComponent<Camera>().targetTexture = textures[i];
			}
			
			if (Application.isPlaying)
			{
				Time.captureFramerate = FPS;
			}
			Directory.CreateDirectory(GetFullPath());

			foreach (GameObject c in cloneCameras)
			{
				c.GetComponent<Camera>().Render();
			}
		}

		private void RenderVRPanorama()
		{
			for (int i = 0; i < cloneCameras.Length; i++)
			{
				SaveScreenshot(cloneCameras[i], "_" + sideNames[i], resolution, resolutionH);
			}
		}

		private string GetFullPath()
		{
			if (customPath) return customPathFolder + Folder + "/";
			else return Path.GetFullPath(string.Format(@"{0}/", Folder));

		}

		private void SaveScreenshot(GameObject panoCamObj, string suffix, int width, int height)
        {
            if (ImageFormatType == VRFormatList.JPG)
			{
				Texture2D screenShot = GetScreenshot(true, panoCamObj, width, height);
				string filePathf = GetFullPath() + _fileName + suffix + ".jpg";
				SaveFileJPG(filePathf, screenShot, jpgQuality);
				DestroyObj(screenShot);
            }

            else if (ImageFormatType == VRFormatList.PNG)
            {
				Texture2D screenShot = GetScreenshot(true, panoCamObj, width, height);
				string filePathf = GetFullPath() + _fileName + suffix + ".png";
				SaveFilePNG(filePathf, screenShot);
				DestroyObj(screenShot);
			}
        }


        private static void SaveFileJPG(string filePathd, Texture2D SShot, int jpg)
        {
            byte[] bytes = SShot.EncodeToJPG(jpg);
            File.WriteAllBytes(filePathd, bytes);
        }

        private static void SaveFilePNG(string filePathd, Texture2D SShot)
        {
            byte[] bytes = SShot.EncodeToPNG();
            File.WriteAllBytes(filePathd, bytes);
        }

		private Texture2D GetScreenshot(bool eye, GameObject panoCamObj, int width, int height)
		{
			RenderTexture rt = RenderTexture.GetTemporary(width, height, 0, renderTexFormat);

			Texture2D tex2D = new Texture2D(width, height, texFormat, false);

			Camera VRCam = panoCamObj.GetComponent<Camera>();
			VRCam.targetTexture = rt;
			VRCam.Render();
			RenderTexture.active = rt;

			tex2D.ReadPixels(new Rect(0, 0, width, height), 0, 0);
			RenderTexture.active = null;
			VRCam.targetTexture = null;
			RenderTexture.ReleaseTemporary(rt);
			return tex2D;
		}

		private void DestroyObj(UnityEngine.Object obj)
		{
			if (Application.isPlaying)
			{
				Destroy(obj);
			}
			else
			{
				DestroyImmediate(obj);
			}
		}
		
		private void PreparePano()
		{
			if (!useAlpha)
				texFormat = TextureFormat.RGB24;
			else texFormat = TextureFormat.ARGB32;

			//qualityTemp = resolution / 32 * renderQuality;
			rig = Instantiate(Resources.Load<GameObject>("VRCubeRig"));
			rig.hideFlags = HideFlags.HideInHierarchy;
			rig.name = "VRCubeRig";
			rig.transform.SetParent(transform, false);
			cam = gameObject;
			restoreFOV = cam.GetComponent<Camera>().fieldOfView;
			cam.GetComponent<Camera>().fieldOfView = 90;

			float IPDistanceV = -IPDistance;

			camObjs = new GameObject[6];
			for (int i = 0; i < 6; i++)
			{
				camObjs[i] = rig.transform.Find(sideNames[i]).gameObject;
			}

			cloneCameras = new GameObject[6];
			for (int i = 0; i < 6; i++)
			{
				cloneCameras[i] = new GameObject("cloneCam");
				Camera c = cloneCameras[i].AddComponent<Camera>();
				c.CopyFrom(cam.GetComponent<Camera>());
				c.farClipPlane = farPlane;
				c.nearClipPlane = nearPlane;
				cloneCameras[i].transform.SetParent(camObjs[i].transform, false);
				cloneCameras[i].transform.localPosition = Vector3.zero;
				cloneCameras[i].transform.localRotation = Quaternion.identity;
			}
		}
#endregion

#region interface
		public bool StartCaptureImage(){
			if (waitRenderFinish == false) {
				CanDestroy = true;
				CaptureTrigger = true;
				step = 1;

				Step1();
				Step2();
				Step3();

				return true;
			} else
				return false;
		}

#endregion
    }
}
