﻿/* Copyright (C) Itseez3D, Inc. - All Rights Reserved
* You may not use this file except in compliance with an authorized license
* Unauthorized copying of this file, via any medium is strictly prohibited
* Proprietary and confidential
* UNLESS REQUIRED BY APPLICABLE LAW OR AGREED BY ITSEEZ3D, INC. IN WRITING, SOFTWARE 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.
* Written by Itseez3D, Inc. <support@avatarsdk.com>, December 2020
*/

using System.Collections;
using System.Collections.Generic;
using System.IO;
using UnityEngine;
using UnityEngine.Rendering;
using static GLTF.CoroutineGLTFSceneImporter;

namespace ItSeez3D.AvatarSdk.Core
{
	public class FullbodyMaterialAdjuster
	{
		enum ShaderType
		{
			Standard,
			HaircutStrands,
			HaircutSolid
		}

		private static readonly string templateBodyMaterialName = "fullbody_materials/avatar_sdk_template_body_material";
		private static readonly string templateBodyPbrMaterialName = "fullbody_materials/avatar_sdk_template_body_pbr_material";
		private static readonly string templateOutfitMaterialName = "fullbody_materials/avatar_sdk_template_outfit_material";
		private static readonly string templateHaircutMaterialName = "fullbody_materials/avatar_sdk_template_haircut_material";

		private static Dictionary<string, ShaderType> haircutsShaders = new Dictionary<string, ShaderType>()
		{
			{ "plus_wavy_bob", ShaderType.HaircutStrands },
			{ "plus_very_long", ShaderType.HaircutStrands },
			{ "plus_shoulder_length", ShaderType.HaircutStrands },
			{ "plus_short_parted", ShaderType.HaircutStrands },
			{ "plus_short_curls", ShaderType.HaircutStrands },
			{ "plus_roman", ShaderType.HaircutStrands },
			{ "plus_rasta", ShaderType.HaircutStrands },
			{ "plus_mid_length_straight2", ShaderType.HaircutStrands },
			{ "plus_mid_length_ruffled", ShaderType.HaircutStrands },
			{ "plus_corkscrew_curls", ShaderType.HaircutStrands },
			{ "plus_bob_parted", ShaderType.HaircutStrands },

			{ "plus_short_slick", ShaderType.HaircutSolid },
			{ "plus_mid_length_wispy", ShaderType.HaircutSolid },
			{ "plus_long_crimped", ShaderType.HaircutSolid },
			{ "plus_balding", ShaderType.HaircutSolid },
			{ "base_short_disheveled", ShaderType.HaircutSolid },
			{ "base_ponytail_with_bangs", ShaderType.HaircutSolid },
			{ "base_mid_length_straight", ShaderType.HaircutSolid },
			{ "base_long_wavy", ShaderType.HaircutSolid },
			{ "base_long_disheveled", ShaderType.HaircutSolid },
			{ "base_short_simple", ShaderType.HaircutSolid },
			{ "base_generated", ShaderType.HaircutSolid }
		};

		private IPersistentStorage persistentStorage = null;

		public FullbodyMaterialAdjuster()
		{
			persistentStorage = AvatarSdkMgr.Storage();
		}

		public IEnumerator PrepareBodyMaterial(CoroutineResult<Material> executionResult, string avatarCode, string outfitName, bool withPbrTextures)
		{
			string templateMaterialName = withPbrTextures ? templateBodyPbrMaterialName : templateBodyMaterialName;
			Material templateMaterial = Resources.Load<Material>(templateMaterialName);
			if (templateMaterial == null)
			{
				Debug.LogError("Template body material isn't found!");
				executionResult.result = null;
				yield break;
			}
			Material bodyMaterial = new Material(templateMaterial);

			yield return UpdateBodyMainTexture(bodyMaterial, avatarCode, outfitName, withPbrTextures);

			if (withPbrTextures)
			{
				string metallnessTextureFilename = persistentStorage.GetAvatarFilename(avatarCode, AvatarFile.FULLBODY_METALLNESS_MAP);
				string roughnessTextureFilename = persistentStorage.GetAvatarFilename(avatarCode, AvatarFile.FULLBODY_ROUGHNESS_MAP);
				string normalTextureFilename = persistentStorage.GetAvatarFilename(avatarCode, AvatarFile.FULLBODY_NORMAL_MAP);

				yield return ConfigurePbrTextures(bodyMaterial, normalTextureFilename, metallnessTextureFilename, roughnessTextureFilename);
			}

			executionResult.result = bodyMaterial;
		}

		public IEnumerator UpdateBodyMainTexture(Material bodyMaterial, string avatarCode, string outfitName, bool withPbrTextures)
		{
			string bodyTextureFilename = persistentStorage.GetAvatarTextureFilename(avatarCode);
			Texture2D bodyTexture = new Texture2D(0, 0);
			bodyTexture.LoadImage(File.ReadAllBytes(bodyTextureFilename));
			yield return null;

			if (!string.IsNullOrEmpty(outfitName))
			{
				Texture2D alphaMaskTexture = new Texture2D(2, 2);
				string alphaMaskTextureFilename = persistentStorage.GetOutfitFilename(avatarCode, outfitName, OutfitFile.OUTFIT_BODY_VISIBILITY_MASK_TEXTURE);
				alphaMaskTexture.LoadImage(File.ReadAllBytes(alphaMaskTextureFilename));
				Color32[] alphaMaskTexColors = alphaMaskTexture.GetPixels32();
				Object.DestroyImmediate(alphaMaskTexture);
				yield return null;

				Color32[] bodyTexColors = bodyTexture.GetPixels32();
				for (int i = 0; i < bodyTexColors.Length; i++)
					bodyTexColors[i].a = alphaMaskTexColors[i].r;
				yield return null;

				bodyTexture = new Texture2D(bodyTexture.width, bodyTexture.height, TextureFormat.ARGB32, true);
				bodyTexture.SetPixels32(bodyTexColors);
				bodyTexture.Apply();

				Resources.UnloadUnusedAssets();
				yield return null;
			}
			bodyMaterial.SetTexture("_MainTex", bodyTexture);
			if (!withPbrTextures)
				bodyMaterial.SetTexture("_EmissionMap", bodyTexture);
		}

		public IEnumerator PrepareHairMaterial(CoroutineResult<Material> executionResult, string avatarCode, string haircutName)
		{
			HaircutMetadata haircutMetadata = HaircutsPersistentStorage.Instance.GetHaircutMetadata(haircutName, avatarCode);
			Texture2D haircutTexture = new Texture2D(0, 0);
			haircutTexture.LoadImage(File.ReadAllBytes(haircutMetadata.Texture));
			yield return null;

			if (haircutsShaders.ContainsKey(haircutName))
			{
				if (haircutsShaders[haircutName] == ShaderType.HaircutSolid)
					executionResult.result = PrepareHairMaterialWithAvatarSdkSolidShader(haircutTexture);
				else
					executionResult.result = PrepareHairMaterialWithAvatarSdkStrandsShader(haircutTexture);
			}
			else
				executionResult.result = PrepareHairMaterialWithStandardShader(haircutTexture);
		}

		public IEnumerator PrepareOutfitMaterial(CoroutineResult<Material> executionResult, string avatarCode, string outfitName, bool withPbrTextures)
		{
			Material templateMaterial = Resources.Load<Material>(templateOutfitMaterialName);
			if (templateMaterial == null)
			{
				Debug.LogError("Template outfit material isn't found!");
				yield break;
			}

			Material outfitMaterial = new Material(templateMaterial);

			string outfitTextureFilename = persistentStorage.GetOutfitFilename(avatarCode, outfitName, OutfitFile.OUTFIT_TEXTURE);
			Texture2D outfitTexture = new Texture2D(0, 0);
			outfitTexture.LoadImage(File.ReadAllBytes(outfitTextureFilename));
			outfitMaterial.SetTexture("_MainTex", outfitTexture);
			yield return null;

			if (withPbrTextures)
			{
				string metallnessTextureFilename = persistentStorage.GetOutfitFilename(avatarCode, outfitName, OutfitFile.OUTFIT_METALLIC_MAP_TEXTURE);
				string roughnessTextureFilename = persistentStorage.GetOutfitFilename(avatarCode, outfitName, OutfitFile.OUTFIT_ROUGHNESS_MAP_TEXTURE);
				string normalTextureFilename = persistentStorage.GetOutfitFilename(avatarCode, outfitName, OutfitFile.OUTFIT_NORMAL_MAP_TEXTURE);

				yield return ConfigurePbrTextures(outfitMaterial, normalTextureFilename, metallnessTextureFilename, roughnessTextureFilename);
			}

			executionResult.result = outfitMaterial;
		}

		private Material PrepareHairMaterialWithStandardShader(Texture2D mainTexture)
		{
			Material templateMaterial = Resources.Load<Material>(templateHaircutMaterialName);
			if (templateMaterial == null)
			{
				Debug.LogError("Template haircut material isn't found!");
				return null;
			}

			Material hairMaterial = new Material(templateMaterial);

			hairMaterial.SetTexture("_MainTex", mainTexture);
			hairMaterial.SetTexture("_EmissionMap", mainTexture);

			return hairMaterial;
		}

		private Material PrepareHairMaterialWithAvatarSdkSolidShader(Texture2D mainTexture)
		{
			Shader solidShader = Shader.Find(ShadersUtils.haircutSolidLitShaderName);
			if (solidShader == null)
			{
				Debug.LogErrorFormat("{0} shader wasn't found. Use Standard shader for haircut.", ShadersUtils.haircutSolidLitShaderName);
				return PrepareHairMaterialWithStandardShader(mainTexture);
			}

			Material hairMaterial = new Material(solidShader);
			hairMaterial.shader = solidShader;
			hairMaterial.renderQueue = (int)RenderQueue.Transparent;
			hairMaterial.SetTexture("_MainTex", mainTexture);
			return hairMaterial;
		}

		private Material PrepareHairMaterialWithAvatarSdkStrandsShader(Texture2D mainTexture)
		{
			Shader strandShader = Shader.Find(ShadersUtils.haircutStrandLitShaderName);
			if (strandShader == null)
			{
				Debug.LogErrorFormat("{0} shader wasn't found. Use Standard shader for haircut.", ShadersUtils.haircutStrandLitShaderName);
				return PrepareHairMaterialWithStandardShader(mainTexture);
			}

			Material hairMaterial = new Material(strandShader);
			hairMaterial.shader = strandShader;
			hairMaterial.SetTexture("_MainTex", mainTexture);
			return hairMaterial;
		}

		private IEnumerator ConfigurePbrTextures(Material material, string normalMapFilename, string metallicMapFilename, string roughnessMapFilename)
		{
			int metallicWithRoughnessTextureWidth = 0;
			int metallicWithRoughnessTextureHeight = 0;

			Color32[] metallnessColors = null;
			if (File.Exists(metallicMapFilename))
			{
				Texture2D metallnessTexture = new Texture2D(0, 0);
				metallnessTexture.LoadImage(File.ReadAllBytes(metallicMapFilename));
				metallicWithRoughnessTextureWidth = metallnessTexture.width;
				metallicWithRoughnessTextureHeight = metallnessTexture.height;
				metallnessColors = metallnessTexture.GetPixels32();
				Object.DestroyImmediate(metallnessTexture);
				yield return null;
			}
			else
				Debug.LogWarningFormat("Texture not found: {0}", metallicMapFilename);

			Color32[] roughnessColors = null;
			if (File.Exists(roughnessMapFilename))
			{
				Texture2D roughnessTexture = new Texture2D(0, 0);
				roughnessTexture.LoadImage(File.ReadAllBytes(roughnessMapFilename));
				metallicWithRoughnessTextureWidth = roughnessTexture.width;
				metallicWithRoughnessTextureHeight = roughnessTexture.height;
				roughnessColors = roughnessTexture.GetPixels32();
				Object.DestroyImmediate(roughnessTexture);
				yield return null;
			}
			else
				Debug.LogWarningFormat("Texture not found: {0}", roughnessMapFilename);

			if (metallicWithRoughnessTextureWidth > 0 && metallicWithRoughnessTextureHeight > 0)
			{
				Texture2D metallicWithRoughnessTexture = new Texture2D(metallicWithRoughnessTextureWidth, metallicWithRoughnessTextureHeight);
				Color32[] metallicWithRoughnessTextureColors = metallicWithRoughnessTexture.GetPixels32();
				for (int i = 0; i < metallicWithRoughnessTextureColors.Length; i++)
				{
					byte metallValue = metallnessColors == null ? (byte)0 : metallnessColors[i].r;
					metallicWithRoughnessTextureColors[i].r = metallValue;
					metallicWithRoughnessTextureColors[i].g = metallValue;
					metallicWithRoughnessTextureColors[i].b = metallValue;

					metallicWithRoughnessTextureColors[i].a = 255;
					if (roughnessColors != null)
						metallicWithRoughnessTextureColors[i].a -= roughnessColors[i].r;
				}
				metallicWithRoughnessTexture.SetPixels32(metallicWithRoughnessTextureColors);
				metallicWithRoughnessTexture.Apply(true, true);
				yield return null;

				material.SetTexture("_MetallicGlossMap", metallicWithRoughnessTexture);
			}

			Texture2D normalTexture = new Texture2D(0, 0);
			normalTexture.LoadImage(File.ReadAllBytes(normalMapFilename));
			yield return null;
			material.SetTexture("_BumpMap", normalTexture);

			Resources.UnloadUnusedAssets();
			yield return null;
		}
	}
}
