﻿/* 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;

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 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 void ConfigureBodyMaterial(SkinnedMeshRenderer bodyRenderer, string avatarCode, string outfitName)
		{
			Material templateMaterial = Resources.Load<Material>(templateBodyMaterialName);
			if (templateMaterial == null)
			{
				Debug.LogError("Template body material isn't found!");
				return;
			}

			Material currentBodyMaterial = bodyRenderer.sharedMaterial;
			Material bodyMaterial = new Material(templateMaterial);

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

			Texture2D mainTexture = currentBodyMaterial.GetTexture("_MainTex") as Texture2D;
			Texture2D dstMainTexture = new Texture2D(mainTexture.width, mainTexture.height);
			Color32[] mainTexColors = mainTexture.GetPixels32();
			Color32[] dstMainTexColors = dstMainTexture.GetPixels32();
			
			for (int i = 0; i < mainTexColors.Length; i++)
			{
				dstMainTexColors[i] = mainTexColors[i];
				dstMainTexColors[i].a = withAlphaMask ? alphaMaskTexColors[i].r : (byte)255;
			}
			dstMainTexture.SetPixels32(dstMainTexColors);
			dstMainTexture.Apply();
			bodyMaterial.SetTexture("_MainTex", dstMainTexture);
			bodyMaterial.SetTexture("_EmissionMap", dstMainTexture);

			Object.DestroyImmediate(mainTexture);

			bodyRenderer.sharedMaterial = bodyMaterial;

			Resources.UnloadUnusedAssets();
		}

		public void ConfigureHairMaterial(SkinnedMeshRenderer hairRenderer, string haircutName)
		{
			if (haircutsShaders.ContainsKey(haircutName))
			{
				if (haircutsShaders[haircutName] == ShaderType.HaircutSolid)
					ConfigureHairWithAvatarSdkSolidShader(hairRenderer);
				else
					ConfigureHairWithAvatarSdkStrandsShader(hairRenderer);
			}
			else
				ConfigureHairWithStandardShader(hairRenderer);
		}

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

			Material currentOutfitMaterial = outfitRenderer.sharedMaterial;
			Material outfitMaterial = new Material(templateMaterial);

			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);

			Texture2D metallnessTexture = new Texture2D(0, 0);
			metallnessTexture.LoadImage(File.ReadAllBytes(metallnessTextureFilename));
			Color32[] metallnessColors = metallnessTexture.GetPixels32();
			int metallnessTextureWidth = metallnessTexture.width;
			int metallnessTextureHeight = metallnessTexture.height;
			Object.DestroyImmediate(metallnessTexture);
			yield return null;

			Texture2D roughnessTexture = new Texture2D(0, 0);
			roughnessTexture.LoadImage(File.ReadAllBytes(roughnessTextureFilename));
			Color32[] roughnessColors = roughnessTexture.GetPixels32();
			Object.DestroyImmediate(roughnessTexture);
			yield return null;

			Texture2D normalTexture = new Texture2D(0, 0);
			normalTexture.LoadImage(File.ReadAllBytes(normalTextureFilename));
			normalTexture.Apply(true, true);

			Texture2D mergedTexture = new Texture2D(metallnessTextureWidth, metallnessTextureHeight);
			Color32[] mergedColors = mergedTexture.GetPixels32();
			for (int i = 0; i < mergedColors.Length; i++)
			{
				mergedColors[i].r = metallnessColors[i].r;
				mergedColors[i].g = metallnessColors[i].r;
				mergedColors[i].b = metallnessColors[i].r;
				mergedColors[i].a = (byte)(255 - roughnessColors[i].r);
			}
			mergedTexture.SetPixels32(mergedColors);
			mergedTexture.Apply(true, true);

			outfitMaterial.SetTexture("_MainTex", currentOutfitMaterial.GetTexture("_MainTex"));
			outfitMaterial.SetTexture("_MetallicGlossMap", mergedTexture);
			outfitMaterial.SetTexture("_BumpMap", normalTexture);

			outfitRenderer.sharedMaterial = outfitMaterial;

			Resources.UnloadUnusedAssets();
			yield return null;
		}

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

			Material currentHairMaterial = hairRenderer.sharedMaterial;
			Material hairMaterial = new Material(templateMaterial);

			hairMaterial.SetTexture("_MainTex", currentHairMaterial.GetTexture("_MainTex"));
			hairMaterial.SetTexture("_EmissionMap", currentHairMaterial.GetTexture("_MainTex"));

			hairRenderer.sharedMaterial = hairMaterial;
		}

		private void ConfigureHairWithAvatarSdkSolidShader(SkinnedMeshRenderer hairRenderer)
		{
			Shader solidShader = Shader.Find(ShadersUtils.haircutSolidLitShaderName);
			if (solidShader == null)
			{
				Debug.LogErrorFormat("{0} shader wasn't found. Use Standard shader for haircut.", ShadersUtils.haircutSolidLitShaderName);
				ConfigureHairWithStandardShader(hairRenderer);
				return;
			}

			hairRenderer.sharedMaterial.shader = solidShader;
			hairRenderer.sharedMaterial.renderQueue = (int)RenderQueue.Transparent;
		}

		private void ConfigureHairWithAvatarSdkStrandsShader(SkinnedMeshRenderer hairRenderer)
		{
			Shader strandShader = Shader.Find(ShadersUtils.haircutStrandLitShaderName);
			if (strandShader == null)
			{
				Debug.LogErrorFormat("{0} shader wasn't found. Use Standard shader for haircut.", ShadersUtils.haircutStrandLitShaderName);
				ConfigureHairWithStandardShader(hairRenderer);
				return;
			}

			hairRenderer.sharedMaterial.shader = strandShader;
		}
	}
}
