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

namespace UltraCombos
{
    public class MeshVertexCollect : MonoBehaviour
    {
        public ComputeShader shader;
        ComputeBuffer args_buffer = null;
        [SerializeField] StructuredBuffer positionBuffer;
        [SerializeField] StructuredBuffer velocityBuffer;

        List<SkinnedMeshRenderer> renderers;
        List<ComputeBuffer> vertex_buffers;

        Mesh bake_mesh;

        private void Start()
        {
            renderers = new List<SkinnedMeshRenderer>();
            renderers.AddRange(GetComponentsInChildren<SkinnedMeshRenderer>());
            vertex_buffers = new List<ComputeBuffer>();

            int sum = 0;
            foreach (var rdr in renderers)
            {
                var m = rdr.sharedMesh;
                var ptr = m.GetNativeVertexBufferPtr(0);
                int size = Utility.Buffer.GetSize(ptr);
                int count = m.vertexCount;
                int stride = size / m.vertexCount;
                sum += count;

                ComputeBuffer buf = new ComputeBuffer(count, stride);
                Utility.Buffer.Copy(buf.GetNativeBufferPtr(), m.GetNativeVertexBufferPtr(0));
                vertex_buffers.Add(buf);
            }

            positionBuffer.Allocate(sum, sizeof(float) * 4);
            velocityBuffer.Allocate(sum, sizeof(float) * 4);

            args_buffer = shader.CreateIndirectComputeArgsBuffer(positionBuffer.count, 1, 1);

            bake_mesh = new Mesh();

        }

        private void FixedUpdate()
        {
            for (int i = 0; i < renderers.Count; i++)
            {
                renderers[i].BakeMesh(bake_mesh);
                Utility.Buffer.Copy(vertex_buffers[i].GetNativeBufferPtr(), bake_mesh.GetNativeVertexBufferPtr(0));
            }


            int index = 0;
            shader.SetBuffer(0, positionBuffer.bufferName, positionBuffer.obj);
            shader.SetBuffer(0, velocityBuffer.bufferName, velocityBuffer.obj);
            for (int i = 0; i < renderers.Count; i++)
            //foreach (var buf in vertex_buffers)
            {
                var buf = vertex_buffers[i];
                shader.SetInt("index", index);
                shader.SetInt("count", buf.count);
                shader.SetMatrix("model_matrix", renderers[i].transform.localToWorldMatrix);
                shader.SetBuffer(0, "vertex_buffer", buf);
                shader.DispatchIndirect(0, args_buffer);
                index += buf.count;
            }


        }

        private void OnDestroy()
        {
            if (args_buffer != null)
            {
                args_buffer.Release();
                args_buffer = null;
            }

            foreach (var buf in vertex_buffers)
            {
                buf.Release();
            }
        }
    }

}