using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using Unity.Mathematics;
using UnityEngine.Assertions;

namespace UnityEngine.Rendering.RadeonRays
{

    internal class AABB
    {
        public float3 Min;
        public float3 Max;

        public AABB()
        {
            Min = new float3(float.PositiveInfinity, float.PositiveInfinity, float.PositiveInfinity);
            Max = new float3(float.NegativeInfinity, float.NegativeInfinity, float.NegativeInfinity);
        }

        public AABB(float3 min, float3 max)
        {
            Min = min;
            Max = max;
        }

        public void Encapsulate(AABB aabb)
        {
            Min = math.min(Min, aabb.Min);
            Max = math.max(Max, aabb.Max);
        }

        public void Encapsulate(float3 point)
        {
            Min = math.min(Min, point);
            Max = math.max(Max, point);
        }

        public bool Contains(AABB rhs)
        {
            return rhs.Min.x >= Min.x && rhs.Min.y >= Min.y && rhs.Min.z >= Min.z &&
                   rhs.Max.x <= Max.x && rhs.Max.y <= Max.y && rhs.Max.z <= Max.z;
        }

        public bool IsValid()
        {
            return Min.x <= Max.x && Min.y <= Max.y && Min.z <= Max.z;
        }
    }


    internal class BvhCheck
    {
        const uint kInvalidID = ~0u;

        public class VertexBuffers
        {
            public GraphicsBuffer vertices;
            public GraphicsBuffer indices;
            public uint vertexBufferOffset = 0;
            public uint vertexCount;
            public uint vertexStride = 3;
            public uint indexBufferOffset = 0;
            public IndexFormat indexFormat = IndexFormat.Int32;
            public uint indexCount;

        };

        public static VertexBuffers Convert(MeshBuildInfo info)
        {
            var res = new VertexBuffers()
            {
                vertices = info.vertices,
                indices = info.triangleIndices,
                vertexBufferOffset = (uint)info.verticesStartOffset,
                vertexCount = info.vertexCount,
                vertexStride = info.vertexStride,
                indexBufferOffset = (uint)info.indicesStartOffset,
                indexCount = info.triangleCount * 3,
                indexFormat = info.indexFormat
            };

            return res;
        }

        public static double SurfaceArea(AABB aabb)
        {
            float3 edges = aabb.Max - aabb.Min;
            return 2.0f * (edges.x * edges.y + edges.x * edges.z + edges.z * edges.y);
        }

        public static double NodeSahCost(uint nodeAddr, AABB nodeAabb, AABB parentAabb)
        {
            double cost = IsLeafNode(nodeAddr) ? GetLeafNodePrimCount(nodeAddr) : 1.2f;
            var a = SurfaceArea(nodeAabb);
            var b = SurfaceArea(parentAabb);
            return cost *  a/ b;
        }

        public static double CheckConsistency(VertexBuffers bvhVertexBuffers, BottomLevelLevelAccelStruct bvh, uint primitiveCount)
        {
            return CheckConsistency(bvhVertexBuffers, bvh.bvh, bvh.bvhOffset, bvh.bvhLeaves, bvh.bvhLeavesOffset, primitiveCount);
        }

        public static double CheckConsistency(GraphicsBuffer bvhBuffer, uint bvhBufferOffset, uint primitiveCount)
        {
            return CheckConsistency(null, bvhBuffer, bvhBufferOffset, null, 0, primitiveCount);
        }

        static double CheckConsistency(
            VertexBuffers bvhVertexBuffers,
            GraphicsBuffer bvhBuffer, uint bvhBufferOffset, GraphicsBuffer bvhLeavesBuffer, uint bvhLeavesBufferOffset,
            uint primitiveCount)
        {
            var header = new BvhHeader[1];
            bvhBuffer.GetData(header, 0, (int)bvhBufferOffset, 1);

            return CheckConsistency(bvhVertexBuffers, bvhBuffer, bvhBufferOffset + 1, bvhLeavesBuffer, bvhLeavesBufferOffset, header[0], primitiveCount);
        }

        public static int ExtractBits(uint value, int startBit, int count)
        {
            uint mask = (uint)(((1 << count) - 1) << startBit);
            return ((int)(mask & value)) >> startBit;
        }

        public static bool IsLeafNode(uint nodeAddr)
        {
            return (nodeAddr & (1 << 31)) != 0;
        }

        public static uint GetLeafNodeFirstPrim(uint nodeAddr)
        {
            return (nodeAddr & ~0xE0000000);
        }

        public static uint GetLeafNodePrimCount(uint nodeAddr)
        {
            return (uint)ExtractBits(nodeAddr, 29, 2) + 1;
        }

        static double CheckConsistency(
            VertexBuffers bvhVertexBuffers,
            GraphicsBuffer bvhBuffer, uint bvhBufferOffset, GraphicsBuffer bvhLeavesBuffer, uint bvhLeavesBufferOffset,
            BvhHeader header, uint primitiveCount)
        {
            uint leafCount = header.leafNodeCount;
            uint rootAddr = header.root;
            var nodeCount = HlbvhBuilder.GetBvhNodeCount(leafCount);
            bool isTopLevel = bvhVertexBuffers == null;

            var bvhNodes = new BvhNode[nodeCount];
            bvhBuffer.GetData(bvhNodes, 0, (int)bvhBufferOffset, (int)nodeCount);

            VertexBuffersCPU vertexBuffers = null;
            uint4[] bvhLeafNodes = null;
            if (!isTopLevel)
            {
                vertexBuffers = DownloadVertexData(bvhVertexBuffers);
                bvhLeafNodes = new uint4[primitiveCount];
                bvhLeavesBuffer.GetData(bvhLeafNodes, 0, (int)bvhLeavesBufferOffset, (int)primitiveCount);
            }

            uint countedPrimitives = 0;

            var rootAabb = GetAabb(vertexBuffers, bvhNodes, bvhLeafNodes, rootAddr, isTopLevel);
            double sahCost = 0.0f;

            var q = new Queue<(uint Addr, uint Parent)>();
            q.Enqueue((Addr: rootAddr, Parent: kInvalidID));
            while (q.Count != 0)
            {
                var current = q.Dequeue();
                uint addr = current.Addr;
                uint parent = current.Parent;

                AABB aabb = GetAabb(vertexBuffers, bvhNodes, bvhLeafNodes, addr, isTopLevel);
                sahCost += NodeSahCost(addr, aabb, rootAabb);

                if (!(isTopLevel && IsLeafNode(addr)))
                    Assert.IsTrue(aabb.IsValid());

                if (IsLeafNode(addr))
                {
                    countedPrimitives += isTopLevel ? 1 : GetLeafNodePrimCount(addr);
                }
                else // internal node
                {
                    var node = bvhNodes[addr];
                    Assert.AreEqual(parent, node.parent);
                    var leftAabb = GetAabb(vertexBuffers, bvhNodes, bvhLeafNodes, node.child0, isTopLevel);
                    var rightAabb = GetAabb(vertexBuffers, bvhNodes, bvhLeafNodes, node.child1, isTopLevel);

                    bool leftOk = (aabb.Contains(leftAabb));
                    bool rightOk = (aabb.Contains(rightAabb));

                    Assert.IsTrue(leftOk);
                    Assert.IsTrue(rightOk);

                    q.Enqueue((Addr: node.child0, Parent: addr));
                    q.Enqueue((Addr: node.child1, Parent: addr));
                }
            }

            Assert.AreEqual(countedPrimitives, primitiveCount);

            return sahCost;
        }

        private sealed class VertexBuffersCPU
        {
            public float[] vertices;
            public uint[] indices;
            public uint vertexStride;
        };


        static uint3 GetFaceIndices(uint[] indices, uint triangleIdx)
        {
            return new uint3(
                indices[3 * triangleIdx],
                indices[3 * triangleIdx + 1],
                indices[3 * triangleIdx + 2]);
        }

        static float3 GetVertex(float[] vertices, uint stride, uint idx)
        {
            uint indexInFloats = idx * stride;
            return new float3(
                vertices[indexInFloats],
                vertices[indexInFloats + 1],
                vertices[indexInFloats + 2]);
        }

        struct Triangle
        {
            public float3 v0;
            public float3 v1;
            public float3 v2;
        };

        static Triangle GetTriangle(float[] vertices, uint stride, uint3 idx)
        {
            Triangle tri;
            tri.v0 = GetVertex(vertices, stride, idx.x);
            tri.v1 = GetVertex(vertices, stride, idx.y);
            tri.v2 = GetVertex(vertices, stride, idx.z);
            return tri;
        }

        static VertexBuffersCPU DownloadVertexData(VertexBuffers vertexBuffers)
        {
            var result = new VertexBuffersCPU();
            result.vertices = new float[vertexBuffers.vertexCount * vertexBuffers.vertexStride];
            result.indices = new uint[vertexBuffers.indexCount];
            result.vertexStride = vertexBuffers.vertexStride;

            if (vertexBuffers.indexFormat == IndexFormat.Int32)
            {
                vertexBuffers.indices.GetData(result.indices, 0, (int)vertexBuffers.indexBufferOffset, (int)vertexBuffers.indexCount);
            }
            else
            {
                var tmp = new ushort[vertexBuffers.indexCount];
                vertexBuffers.indices.GetData(tmp, 0, (int)vertexBuffers.indexBufferOffset, (int)vertexBuffers.indexCount);
                for (int i = 0; i < vertexBuffers.indexCount; ++i)
                    result.indices[i] = tmp[i];
            }

            vertexBuffers.vertices.GetData(result.vertices, 0, (int)vertexBuffers.vertexBufferOffset, (int)(vertexBuffers.vertexCount * vertexBuffers.vertexStride));

            return result;
        }

        static AABB GetAabb(VertexBuffersCPU bvhVertexBuffers, BvhNode[] bvhNodes, uint4[] bvhLeafNodes, uint nodeAddr, bool isTopLevel)
        {
            var aabb = new AABB();

            if (!IsLeafNode(nodeAddr))
            {
                var node = bvhNodes[nodeAddr];
                AABB left = new AABB(node.aabb0_min, node.aabb0_max);
                aabb.Encapsulate(left);

                AABB right = new AABB(node.aabb1_min, node.aabb1_max);
                aabb.Encapsulate(right);
            }
            else if (!isTopLevel)
            {
                int firstIndex = (int)GetLeafNodeFirstPrim(nodeAddr);
                int triangleCount = (int)GetLeafNodePrimCount(nodeAddr);
                for (int i = 0; i < triangleCount; ++i)
                {
                    uint index = (uint)(i + firstIndex);
                    uint3 triangleIndices = bvhLeafNodes[index].xyz;
                    uint3 meshTriangleindices = GetFaceIndices(bvhVertexBuffers.indices, bvhLeafNodes[index].w);

                    Assert.AreEqual(meshTriangleindices, triangleIndices);

                    var triangle = GetTriangle(bvhVertexBuffers.vertices, bvhVertexBuffers.vertexStride, triangleIndices);

                    aabb.Encapsulate(triangle.v0);
                    aabb.Encapsulate(triangle.v1);
                    aabb.Encapsulate(triangle.v2);
                }
            }

            return aabb;
        }
    }
}