using UnityEngine;
using System.Runtime.InteropServices;

public class VoxelRaycastGpuManager
{
    ComputeShader raycastShader;
    OctreeNode root;                 // assign your built octree root
    public VoxelRaycastGpuManager(ComputeShader computeShader, OctreeNode octreeRoot)
    {
        raycastShader = computeShader;
        root = octreeRoot;
    }

    ComputeBuffer nodeBuffer;

    public LinearTree linearTree;

    int kernel;

    ComputeBuffer hitCounterBuffer = null;
    ComputeBuffer rayBuffer = null;

    int raysPerBatch;

    int batchDataClassSize = Marshal.SizeOf(typeof(VoxelRaycastGPU.BatchData));

    int groupsX;

    int maxRaycastPerIteration;


    public VoxelRaycastGPU.BatchData[] Raycast(in VoxelRaycastGPU.BatchData[] batchData, int datasLenght)
    {
        ComputeBuffer hitBuffer = new ComputeBuffer(5000, batchDataClassSize, ComputeBufferType.Append);
        ComputeBuffer datasBuffer = new ComputeBuffer(datasLenght, batchDataClassSize, ComputeBufferType.Default);
        ComputeBuffer countBuffer = new ComputeBuffer(1, sizeof(int), ComputeBufferType.Raw);

        int iteration = 0;
        int currentCount = datasLenght;
        int previousCount = datasLenght;

        datasBuffer.SetData(batchData, 0, 0, currentCount);

        while (iteration < 4 && currentCount > 0)
        {
            previousCount = currentCount;

            raycastShader.SetBuffer(kernel, "batchDatas", datasBuffer);
            raycastShader.SetBuffer(kernel, "hits", hitBuffer);

            int threadsY = 8;

            int groupsY = Mathf.CeilToInt((float)currentCount / threadsY);

            raycastShader.Dispatch(kernel, groupsX, groupsY, 1);

            ComputeBuffer.CopyCount(hitBuffer, countBuffer, 0);
            int[] countArr = new int[1];
            countBuffer.GetData(countArr);
            currentCount = countArr[0];

            /*VoxelRaycastGPU.BatchData[] hits = new VoxelRaycastGPU.BatchData[currentCount];
            hitBuffer.GetData(hits, 0, 0, currentCount);

            for (int i = 0; i < currentCount; i++)
            {
                GameObject sphere = GameObject.CreatePrimitive(PrimitiveType.Sphere);
                sphere.transform.position = hits[i].origin;
                sphere.transform.localScale = Vector3.one * 0.5f;
            }*/

            if (currentCount > 0)
            {
                (datasBuffer, hitBuffer) = (hitBuffer, datasBuffer);

                hitBuffer.Release();
                hitBuffer = new ComputeBuffer(5000, batchDataClassSize, ComputeBufferType.Append);
            }

            iteration++;
        }

        VoxelRaycastGPU.BatchData[] result = new VoxelRaycastGPU.BatchData[previousCount];
        if (currentCount == 0)
            datasBuffer.GetData(result, 0, 0, previousCount);
        else
            hitBuffer.GetData(result, 0, 0, previousCount);

        hitBuffer.Release();
        datasBuffer.Release();
        countBuffer.Release();

        return result;
    }

    public void Init( int nbRaysPerBatch, in VoxelRaycastGPU.Ray[] rays )
    {
        // Flatten octree
        linearTree = OctreeGpuHelpers.FlattenOctree(root);
        int nodeStride = Marshal.SizeOf(typeof(LinearNode)); // should be 64

        rayBuffer = new ComputeBuffer(rays.Length, Marshal.SizeOf(typeof(VoxelRaycastGPU.Ray)), ComputeBufferType.Default);
        rayBuffer.SetData(rays, 0, 0, rays.Length);

        // Create GPU buffer for nodes
        nodeBuffer = new ComputeBuffer(linearTree.nodes.Length, nodeStride, ComputeBufferType.Default);
        nodeBuffer.SetData(linearTree.nodes);

        hitCounterBuffer = new ComputeBuffer(1, sizeof(int), ComputeBufferType.Raw);
        uint[] counterInit = { 0 };
        counterInit[0] = 0;
        hitCounterBuffer.SetData(counterInit);
        
        kernel = raycastShader.FindKernel("CSMain");

        raycastShader.SetBuffer(kernel, "nodes", nodeBuffer);
        raycastShader.SetBuffer(kernel, "hitCount", hitCounterBuffer);
        raycastShader.SetBuffer(kernel, "rays", rayBuffer);

        raycastShader.SetInt("raysPerBatch", nbRaysPerBatch);
        raycastShader.SetInt("rootIndex", linearTree.rootIndex);
        raycastShader.SetInt("nodeCount", linearTree.nodes.Length);

        raycastShader.SetFloat("rootHalfSize", root.bounds.size.x / 2f);
        raycastShader.SetFloats("rootCenter", new float[3] { root.bounds.center.x, root.bounds.center.y, root.bounds.center.z });

        raysPerBatch = nbRaysPerBatch;

        groupsX = Mathf.CeilToInt((float)raysPerBatch / 8);

        maxRaycastPerIteration = 5000 / raysPerBatch;
    }

    ~VoxelRaycastGpuManager()
    {
        if (hitCounterBuffer != null)
            hitCounterBuffer.Release();

        if (rayBuffer != null)
            rayBuffer.Release();
    }
}