using System.Collections.Generic;
using System.IO;
using Sirenix.OdinInspector;
using UnityEngine;
#if UNITY_EDITOR
using UnityEditor;
using UnityEditor.SceneManagement;
#endif

namespace DanieleMarotta.RiversongCodeShowcase
{
    public class MeshBaker : MonoBehaviour
    {
        [SerializeField]
        private string _layer = "Default";

        [SerializeField]
        private List<GameObject> _doNotBake = new();

#if UNITY_EDITOR
        [Button(ButtonSizes.Large)]
        [GUIColor("cyan")]
        private void BakeMesh()
        {
            var excluded = GetExcluded();

            var combineInstancesByMaterial = CollectCombineInstances(excluded);
            if (combineInstancesByMaterial.ValueCount <= 0) return;

            var bakedName = $"{gameObject.name.Replace("_Authoring", string.Empty)}_Baked";
            var bakedMesh = CreateBakedMesh(bakedName, combineInstancesByMaterial);
            var parentFolder = GetBakeOutputFolder(out var isPrefab);

            AssetDatabase.CreateAsset(bakedMesh, $"{parentFolder}/{bakedName}.asset");

            if (isPrefab) SaveBakedPrefab(parentFolder, bakedName, bakedMesh, combineInstancesByMaterial);

            AssetDatabase.SaveAssets();
            AssetDatabase.Refresh();
        }

        private ListMultiDictionary<Material, CombineInstance> CollectCombineInstances(HashSet<Transform> excluded)
        {
            var meshRenderers = GetComponentsInChildren<MeshRenderer>();
            var combineInstancesByMaterial = new ListMultiDictionary<Material, CombineInstance>();

            foreach (var childRenderer in meshRenderers)
            {
                if (excluded.Contains(childRenderer.transform)) continue;
                AddRendererCombineInstances(childRenderer, combineInstancesByMaterial);
            }

            return combineInstancesByMaterial;
        }

        private void AddRendererCombineInstances(MeshRenderer childRenderer, ListMultiDictionary<Material, CombineInstance> combineInstancesByMaterial)
        {
            var meshFilter = childRenderer.GetComponent<MeshFilter>();
            if (!meshFilter || !meshFilter.sharedMesh) return;

            var mesh = meshFilter.sharedMesh;
            var materials = childRenderer.sharedMaterials;
            var subMeshCount = Mathf.Min(mesh.subMeshCount, materials.Length);

            for (var i = 0; i < subMeshCount; i++)
            {
                var material = materials[i];
                if (!material) continue;

                combineInstancesByMaterial.Add(
                    material,
                    new CombineInstance
                    {
                        mesh = mesh,
                        subMeshIndex = i,
                        transform = meshFilter.transform.localToWorldMatrix
                    });
            }
        }

        private Mesh CreateBakedMesh(string bakedName, ListMultiDictionary<Material, CombineInstance> combineInstancesByMaterial)
        {
            var bakedMesh = new Mesh { name = bakedName };

            if (combineInstancesByMaterial.KeyCount == 1)
            {
                CombineSingleMaterialMesh(bakedMesh, combineInstancesByMaterial);
                return bakedMesh;
            }

            CombineMultiMaterialMesh(bakedMesh, combineInstancesByMaterial);
            return bakedMesh;
        }

        private void CombineSingleMaterialMesh(Mesh bakedMesh, ListMultiDictionary<Material, CombineInstance> combineInstancesByMaterial)
        {
            foreach (var material in combineInstancesByMaterial.Keys)
            {
                var combineInstances = combineInstancesByMaterial[material];
                bakedMesh.CombineMeshes(combineInstances.ToArray(), true, true);
            }
        }

        private void CombineMultiMaterialMesh(Mesh bakedMesh, ListMultiDictionary<Material, CombineInstance> combineInstancesByMaterial)
        {
            var finalCombineInstances = new List<CombineInstance>();
            var tempMeshes = new List<Mesh>();

            foreach (var material in combineInstancesByMaterial.Keys)
            {
                var tempMesh = new Mesh();
                tempMesh.CombineMeshes(combineInstancesByMaterial[material].ToArray(), true, true);

                finalCombineInstances.Add(
                    new CombineInstance
                    {
                        mesh = tempMesh,
                        transform = Matrix4x4.identity
                    });

                tempMeshes.Add(tempMesh);
            }

            bakedMesh.CombineMeshes(finalCombineInstances.ToArray(), false, false);

            foreach (var mesh in tempMeshes) DestroyImmediate(mesh);
        }

        private string GetBakeOutputFolder(out bool isPrefab)
        {
            var selectedPrefab = PrefabUtility.GetPrefabAssetPathOfNearestInstanceRoot(gameObject);
            if (string.IsNullOrEmpty(selectedPrefab))
            {
                var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
                if (prefabStage && prefabStage.prefabContentsRoot == gameObject) selectedPrefab = prefabStage.assetPath;
            }

            isPrefab = !string.IsNullOrEmpty(selectedPrefab);
            return isPrefab ? Path.GetDirectoryName(selectedPrefab) : "Assets";
        }

        private void SaveBakedPrefab(string parentFolder, string bakedName, Mesh bakedMesh, ListMultiDictionary<Material, CombineInstance> combineInstancesByMaterial)
        {
            var baked = new GameObject(bakedName) { layer = LayerMask.NameToLayer(_layer) };
            baked.AddComponent<MeshFilter>().sharedMesh = bakedMesh;

            var bakedRenderer = baked.AddComponent<MeshRenderer>();
            bakedRenderer.sharedMaterials = new List<Material>(combineInstancesByMaterial.Keys).ToArray();

            AddExcluded(baked.transform);

            PrefabUtility.SaveAsPrefabAsset(baked, $"{parentFolder}/{bakedName}.prefab");
            DestroyImmediate(baked);
        }

        private HashSet<Transform> GetExcluded()
        {
            var excluded = new HashSet<Transform>();
            foreach (var obj in _doNotBake)
            {
                if (!obj) continue;
                ExcludeHierarchy(obj.transform, excluded);
            }

            return excluded;
        }

        private void ExcludeHierarchy(Transform root, HashSet<Transform> excluded)
        {
            excluded.Add(root);
            for (var i = 0; i < root.childCount; i++) ExcludeHierarchy(root.GetChild(i), excluded);
        }

        private void AddExcluded(Transform root)
        {
            foreach (var obj in _doNotBake)
            {
                if (!obj) continue;
                var clone = Instantiate(obj, root);
                clone.name = obj.name;
                ApplyRelativeTransform(obj.transform, clone.transform);
            }
        }

        private void ApplyRelativeTransform(Transform source, Transform target)
        {
            var relativeMatrix = transform.worldToLocalMatrix * source.localToWorldMatrix;

            var localPosition = relativeMatrix.MultiplyPoint3x4(Vector3.zero);

            var localForward = relativeMatrix.MultiplyVector(Vector3.forward);
            var localUp = relativeMatrix.MultiplyVector(Vector3.up);

            var scaleX = relativeMatrix.GetColumn(0).magnitude;
            var scaleY = relativeMatrix.GetColumn(1).magnitude;
            var scaleZ = relativeMatrix.GetColumn(2).magnitude;

            target.localPosition = localPosition;
            target.localRotation = Quaternion.LookRotation(localForward, localUp);
            target.localScale = new Vector3(scaleX, scaleY, scaleZ);
        }
#endif
    }
}