using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Rendering;

namespace pdinklag
{
    public static class MeshBake
    {
        /// <summary>
        /// Bakes all meshes found in the given object hierarchy into a single static mesh.
        /// </summary>
        /// <param name="root">the root transform</param>
        /// <returns>the baked mesh</returns>
        public static Mesh BakeMeshes(this Transform root)
        {
            var mesh = new Mesh();
            mesh.indexFormat = IndexFormat.UInt32;

            var skinnedMeshRenderers = root.GetComponentsInChildren<SkinnedMeshRenderer>();
            var meshRenderers = root.GetComponentsInChildren<MeshRenderer>();

            var numMeshRenderers = meshRenderers.Length;
            var numSkinnedMeshRenderers = skinnedMeshRenderers.Length;

            if (numMeshRenderers + numSkinnedMeshRenderers > 0)
            {
                var combineInstances = new List<CombineInstance>();

                for (var i = 0; i < numSkinnedMeshRenderers; i++)
                {
                    if (skinnedMeshRenderers[i].enabled)
                    {
                        var baked = new Mesh();
                        skinnedMeshRenderers[i].BakeMesh(baked);

                        var numSubmeshes = baked.subMeshCount;
                        for (var j = 0; j < numSubmeshes; j++)
                        {
                            var combineInstance = new CombineInstance
                            {
                                mesh = baked,
                                subMeshIndex = j,
                                transform = Matrix4x4.TRS(skinnedMeshRenderers[i].transform.position, skinnedMeshRenderers[i].transform.rotation, Vector3.one),
                                lightmapScaleOffset = Vector4.zero,
                                realtimeLightmapScaleOffset = Vector4.zero
                            };
                            combineInstances.Add(combineInstance);
                        }
                    }
                }

                for (var i = 0; i < numMeshRenderers; i++)
                {
                    if (meshRenderers[i].enabled && meshRenderers[i].TryGetComponent(out MeshFilter mf))
                    {
                        var mfMesh = mf.sharedMesh;
                        var numSubmeshes = mfMesh.subMeshCount;
                        for (var j = 0; j < numSubmeshes; j++)
                        {
                            var combineInstance = new CombineInstance
                            {
                                mesh = mfMesh,
                                subMeshIndex = j,
                                transform = Matrix4x4.TRS(meshRenderers[i].transform.position, meshRenderers[i].transform.rotation, meshRenderers[i].transform.lossyScale),
                                lightmapScaleOffset = Vector4.zero,
                                realtimeLightmapScaleOffset = Vector4.zero
                            };
                            combineInstances.Add(combineInstance);
                        }
                    }
                }

                mesh.CombineMeshes(combineInstances.ToArray(), true, true, false);
            }
            return mesh;
        }
    }
}
