﻿// Created by SHAJIKUworks
// ※Transform値によってうまくベイクできない場合もあります……許して……

using System.IO;
using System.Linq;
using System.Collections.Generic;
using UnityEngine;

#if UNITY_EDITOR
using UnityEditor;
#endif

namespace AvatarPoseBaker
{
    public class AvatarPoseBaker : MonoBehaviour
    {
        const string BAKED_PREFAB_DEST_DIRECTORY = "BakedAvatar/";
        static readonly string PROGRESS_TITLE = "AvatarPoseBaker";

        class BakedMeshWithRenderer
        {
            public readonly Renderer Renderer;
            public readonly Mesh Mesh;

            public BakedMeshWithRenderer(MeshRenderer meshRenderer)
            {
                Renderer = meshRenderer;

                // Meshのクローンと記憶
                Mesh = Instantiate(meshRenderer.GetComponent<MeshFilter>().sharedMesh);
                Mesh.name = "Mesh_" + meshRenderer.name;
            }

            public BakedMeshWithRenderer(SkinnedMeshRenderer skinRenderer)
            {
                Renderer = skinRenderer;

                // RootBoneの設定されているSkinnedMeshRendererはボーンに追随するので、Transform値を初期値にする
                var skinTransform = skinRenderer.transform;
                var beforeLocalPosition = skinTransform.localPosition;
                var beforeLocalRotation = skinTransform.localRotation;
                var beforeLocalScale = skinTransform.localScale;
                if (skinRenderer.rootBone != null)
                {
                    skinTransform.localPosition = Vector3.zero;
                    skinTransform.localRotation = Quaternion.identity;
                    skinTransform.localScale = Vector3.one;
                }

                // Meshのベイクと記憶
                Mesh = new Mesh { name = "Mesh_" + skinRenderer.name };
                skinRenderer.BakeMesh(Mesh);

                if (skinRenderer.rootBone != null)
                {
                    skinTransform.localPosition = beforeLocalPosition;
                    skinTransform.localRotation = beforeLocalRotation;
                    skinTransform.localScale = beforeLocalScale;
                }
            }
        }

        [SerializeField, Button("Bake", "ポーズをベイクする！", BAKED_PREFAB_DEST_DIRECTORY)] bool _bakeButton;

        void Bake(string saveDirectory = "")
        {
#if UNITY_EDITOR
            EditorUtility.DisplayProgressBar(PROGRESS_TITLE, "準備をしています...", 0f);

            var meshes = GetComponentsInChildren<MeshRenderer>().Select(m => new BakedMeshWithRenderer(m));
            var skins = GetComponentsInChildren<SkinnedMeshRenderer>().Select(s => new BakedMeshWithRenderer(s));
            var allMeshes = meshes.Union(skins);

            var bakedAvatarObj = new GameObject("Baked_" + gameObject.name);
            var bakedAvatarTransform = bakedAvatarObj.transform;
            bakedAvatarTransform.position = transform.position;

            var destDirectory = Path.Combine(Path.Combine("Assets", saveDirectory), bakedAvatarObj.name);
            if (!Directory.Exists(destDirectory))
                Directory.CreateDirectory(destDirectory);

            EditorUtility.DisplayProgressBar(PROGRESS_TITLE, "オブジェクトのコピー中...", 0f);

            // 表示物のクローン
            var bakedObjects = new List<GameObject>();
            foreach (var mesh in allMeshes)
            {
                var newObj = new GameObject(mesh.Renderer.name);
                var objTransform = newObj.transform;

                var meshTransform = mesh.Renderer.transform;
                objTransform.SetParent(bakedAvatarTransform);
                objTransform.localScale = meshTransform.lossyScale;
                objTransform.position = meshTransform.position;
                objTransform.rotation = meshTransform.rotation;

                newObj.AddComponent<MeshFilter>().sharedMesh = mesh.Mesh;
                var newRenderer = newObj.AddComponent<MeshRenderer>();
                CopyRendererParameters(mesh.Renderer, newRenderer);

                bakedObjects.Add(newObj);
            }

            // 使用アセット辞書
            var dicDestMaterials = new Dictionary<Material, Material>();
            var dicDestTextures = new Dictionary<Texture, Texture>();

            // アセットの保存
            var meshIndex = 0;
            foreach (var element in bakedObjects.Select((obj, index) => new { obj, index }))
            {
                EditorUtility.DisplayProgressBar(PROGRESS_TITLE, "[" + element.obj.name + "]の処理中...", 1f * element.index / bakedObjects.Count());

                // Meshの保存
                var mesh = element.obj.GetComponent<MeshFilter>().sharedMesh;
                SaveAsset(mesh, Path.Combine(destDirectory, mesh.name + "_asset" + (meshIndex++).ToString() + ".asset"));

                var materialIndex = 0;
                var materials = new List<Material>();
                var meshRenderer = element.obj.GetComponent<MeshRenderer>();
                foreach (var sourceMaterial in meshRenderer.sharedMaterials)
                {
                    // コピー済みMaterialならリストに追加するのみ
                    if (dicDestMaterials.ContainsKey(sourceMaterial))
                    {
                        materials.Add(dicDestMaterials[sourceMaterial]);
                        continue;
                    }

                    // 未コピーMaterialのクローンと差し替え
                    var destMaterial = Instantiate(sourceMaterial);
                    materials.Add(destMaterial);

                    // Shaderの差し替え
                    var shader = destMaterial.shader;

                    // Textureの保存と差し替え
                    var textureIndex = 0;
                    var textureProperties = GetTextureProperties(shader);
                    foreach (var texName in textureProperties)
                    {
                        // 使用していないTextureはスキップ
                        var sourceTexture = destMaterial.GetTexture(texName);
                        if (sourceTexture == null)
                            continue;

                        // コピー済みTextureなら差し替えするのみ
                        if (dicDestTextures.ContainsKey(sourceTexture))
                        {
                            destMaterial.SetTexture(texName, dicDestTextures[sourceTexture]);
                            continue;
                        }

                        // Textureのコピーと差し替え
                        var sourceTexturePath = AssetDatabase.GetAssetPath(sourceTexture);
                        var sourceTextureExtension = Path.GetExtension(sourceTexturePath);
                        var newTexturePath = Path.Combine(destDirectory, sourceTexture.name + "_asset" + (textureIndex++).ToString() + sourceTextureExtension);
                        var destTexture = CopyAsset(sourceTexture, newTexturePath) as Texture;
                        destMaterial.SetTexture(texName, destTexture);
                        dicDestTextures.Add(sourceTexture, destTexture);
                    }

                    // Materialの保存と差し替え
                    AssetDatabase.CreateAsset(destMaterial, Path.Combine(destDirectory, destMaterial.name + "_asset" + (materialIndex++).ToString() + ".mat"));
                    dicDestMaterials.Add(sourceMaterial, destMaterial);
                }
                meshRenderer.sharedMaterials = materials.ToArray();
            }

            bakedAvatarTransform.position = Vector3.zero;

            EditorUtility.DisplayProgressBar(PROGRESS_TITLE, "Prefabの保存中...", 1f);

            var prefabPath = Path.Combine(destDirectory, bakedAvatarObj.name + ".prefab").Replace("\\", "/");
            PrefabUtility.CreatePrefab(prefabPath, bakedAvatarObj);
            AssetDatabase.SaveAssets();

            EditorUtility.ClearProgressBar();
            EditorUtility.DisplayDialog(PROGRESS_TITLE, "ポーズのベイクが完了しました！\n" + prefabPath, "OK");
#endif
        }

        Object SaveAsset(Object asset, string path)
        {
#if UNITY_EDITOR
            var dir = Path.GetDirectoryName(path);
            if (!Directory.Exists(dir))
                Directory.CreateDirectory(dir);
            AssetDatabase.CreateAsset(asset, path);
            AssetDatabase.SaveAssets();
#endif
            return asset;
        }

        Object CopyAsset(Object asset, string destPath)
        {
#if UNITY_EDITOR
            var destDir = Path.GetDirectoryName(destPath);
            if (!Directory.Exists(destDir))
                Directory.CreateDirectory(destDir);
            var sourcePath = AssetDatabase.GetAssetPath(asset);
            AssetDatabase.CopyAsset(sourcePath, destPath);
            AssetDatabase.SaveAssets();

            return AssetDatabase.LoadMainAssetAtPath(destPath);
#else
            return asset;
#endif
        }

        List<string> GetTextureProperties(Shader shader)
        {
#if UNITY_EDITOR
            var textureProperties = new List<string>();

            for (var i = 0; i < ShaderUtil.GetPropertyCount(shader); i++)
                if (ShaderUtil.GetPropertyType(shader, i) == ShaderUtil.ShaderPropertyType.TexEnv)
                    textureProperties.Add(ShaderUtil.GetPropertyName(shader, i));

            return textureProperties;
#else
            return new List<string>();
#endif
        }

        void CopyRendererParameters(Renderer source, Renderer dest)
        {
            dest.lightProbeUsage = source.lightProbeUsage;
            dest.reflectionProbeUsage = source.reflectionProbeUsage;
            dest.probeAnchor = source.probeAnchor;
            dest.shadowCastingMode = source.shadowCastingMode;
            dest.receiveShadows = source.receiveShadows;
            dest.motionVectorGenerationMode = source.motionVectorGenerationMode;

            dest.sharedMaterials = source.sharedMaterials;
            dest.allowOcclusionWhenDynamic = source.allowOcclusionWhenDynamic;
        }
    }
}