using UnityEngine;
using Unity.Jobs;
using Unity.Burst;
using Unity.Mathematics;
using Unity.Collections;
using UnityEngine.Rendering;
using Pathfinding;

namespace M7 {
    [RequireComponent(typeof(MeshFilter))]
    public class EntityManagerRVOExample : MonoBehaviour {
        [System.Runtime.InteropServices.StructLayout(System.Runtime.InteropServices.LayoutKind.Sequential)]
        struct Vertex {
            public float3 position;
            public Color32 color;
            public float2 uv;
        }

        public CollisionAvoidanceProvider colaProvider;
        public bool spawnOnEdge;                                // If spawning inside the circle or on it's perimiter
        public float spawnAreaScale = 100;

        [Space]
        public int poolSize = 10000;
        public int numAgentsToSpawn = 10000;
        public float agentRadius = 1.5f;                        // Agent radius
        public float agentMaxSpeed = 5;                         // Max speed for an agent
		public Vector3 renderingOffset = Vector3.up * 0.1f;     // Offset from the agent position. Used to get rid of z-buffer issues


        NativeArray<Vertex> nativeVerts = default;
        NativeArray<int> nativeTris = default;
        NativeArray<Color> agentColors = default;

        Mesh mesh;                                              // Mesh used to render agents
        int vertexCount;
        VertexAttributeDescriptor[] layout;
        JobHandle meshJob;

        void OnDestroy() {
            if (nativeVerts != default) nativeVerts.Dispose();
            if (nativeTris != default) nativeTris.Dispose();
            if (agentColors != default) agentColors.Dispose();
        }

        void Start() {
            mesh = new Mesh();
            GetComponent<MeshFilter>().mesh = mesh;
            vertexCount = poolSize * 4;

            colaProvider.OnJobScheduled += () => GenerateMeshJob(colaProvider.job);
            colaProvider.OnJobCompleted += () => UpdateMesh();
            colaProvider.Initialize(poolSize, agentRadius, agentMaxSpeed);  // Must add agents to colaP before initializing it

            Initialize();

            colaProvider.ActivateAgents();     // Aftter adding agents to colaP we need to enable them
        }

        internal void Initialize() {
            float agentArea = poolSize * agentRadius * agentRadius * Mathf.PI;
            const float EmptyFraction = 0.7f;
            const float PackingDensity = 0.9f;
            float innerCircleRadius = Mathf.Sqrt(agentArea / (Mathf.PI * (1 - EmptyFraction * EmptyFraction)));
            float outerCircleRadius = Mathf.Sqrt(innerCircleRadius * innerCircleRadius + poolSize * agentRadius * agentRadius / PackingDensity);
            Vector3 pos;

            Pathfinding.Util.Memory.Realloc(ref agentColors, poolSize, Allocator.Persistent);

            // Set agent position and add them to the simulation
            if (spawnOnEdge) {
                // Point
                for (int i = 0; i < numAgentsToSpawn; i++) {
                    pos = new Vector3(Mathf.Cos(i * Mathf.PI * 2.0f / poolSize), 0, Mathf.Sin(i * Mathf.PI * 2.0f / poolSize)) * spawnAreaScale;
                    colaProvider.AddAgent(i, pos, float3.zero);
                    agentColors[i] = AstarMath.HSVToRGB(i * 360.0f / poolSize, 0.8f, 0.6f);
                }
            } else {
                // Circle
                for (int i = 0; i < numAgentsToSpawn; i++) {
                    pos = new Vector3(Mathf.Cos(i * Mathf.PI * 2.0f / poolSize), 0, Mathf.Sin(i * Mathf.PI * 2.0f / poolSize)) * math.lerp(innerCircleRadius, outerCircleRadius, UnityEngine.Random.value);
                    colaProvider.AddAgent(i, pos, -pos);
                    agentColors[i] = AstarMath.HSVToRGB(i * 360.0f / poolSize, 0.8f, 0.6f);
                }
            }

            // Make sure the arrays are large enough
            layout = new[] {
                new VertexAttributeDescriptor(VertexAttribute.Position, VertexAttributeFormat.Float32, 3),
                new VertexAttributeDescriptor(VertexAttribute.Color, VertexAttributeFormat.UNorm8, 4),
                new VertexAttributeDescriptor(VertexAttribute.TexCoord0, VertexAttributeFormat.Float32, 2),
            };

            mesh.SetVertexBufferParams(vertexCount, layout);
            mesh.SetIndexBufferParams(poolSize * 6, IndexFormat.UInt32);    // To allow for more than ≈16k agents we need to use a 32 bit format for the mesh

            nativeVerts = new NativeArray<Vertex>(poolSize * 4, Allocator.Persistent, NativeArrayOptions.UninitializedMemory);
            nativeTris = new NativeArray<int>(poolSize * 6, Allocator.Persistent, NativeArrayOptions.UninitializedMemory);

            //GenerateMeshJob(colaProvider.job).Complete();
        }

        void Update() {
            colaProvider.UpdateLoop();
        }

        private void UpdateMesh() {
            meshJob.Complete();

            mesh.SetVertexBufferData(nativeVerts, 0, 0, nativeVerts.Length);
            mesh.SetIndexBufferData(nativeTris, 0, 0, nativeTris.Length);
            mesh.subMeshCount = 1;
            mesh.SetSubMesh(0, new SubMeshDescriptor(0, nativeTris.Length, MeshTopology.Triangles), MeshUpdateFlags.DontRecalculateBounds);
            mesh.RecalculateBounds();
        }

        JobHandle GenerateMeshJob(JobHandle dependsOn) {
            // Create a job to generate a mesh (for entity display)
            return meshJob = new JobGenerateMesh {
                interpolatedRotations = colaProvider.interpolatedRotations,
                agentPositions = colaProvider.burstSim.simulationData.position,
                agentRadii = colaProvider.burstSim.simulationData.radius,
                agentColors = agentColors,
                verts = nativeVerts,
                tris = nativeTris,
                agentCount = colaProvider.agentCount,
                renderingOffset = renderingOffset,
            }.Schedule(dependsOn);
        }

        [BurstCompile(FloatMode = FloatMode.Fast)]
        struct JobGenerateMesh : IJob {
            [ReadOnly] public NativeArray<float2> interpolatedRotations;
            [ReadOnly] public NativeArray<float3> agentPositions;
            [ReadOnly] public NativeArray<float> agentRadii;
            [ReadOnly] public NativeArray<Color> agentColors;
            [WriteOnly] public NativeArray<Vertex> verts;
            [WriteOnly] public NativeArray<int> tris;
            [ReadOnly] public Vector3 renderingOffset;
            [ReadOnly] public int agentCount;

            int vc;
            int tc;
            Color32 color;
            float3 orig;
            float3 right;
            float3 forward;

            public void Execute() {
                for (int i = 0; i < agentCount; i++) {

                    // Create a square with the "forward" direction along the agent's velocity
                    forward = math.normalizesafe(new float3(interpolatedRotations[i].x, 0, interpolatedRotations[i].y)) * agentRadii[i];
                    if (math.all(forward == 0)) forward = new float3(0, 0, agentRadii[i]);
                    right = math.cross(new float3(0, 1, 0), forward);


                    orig = agentPositions[i] + (float3)renderingOffset;

                    vc = 4 * i;
                    tc = 2 * 3 * i;

                    color = agentColors[i];
                    verts[vc + 0] = new Vertex {
                        position = orig + forward - right,
                        uv = new float2(0, 1),
                        color = color,
                    };

                    verts[vc + 1] = new Vertex {
                        position = orig + forward + right,
                        uv = new float2(1, 1),
                        color = color,
                    };

                    verts[vc + 2] = new Vertex {
                        position = orig - forward + right,
                        uv = new float2(1, 0),
                        color = color,
                    };

                    verts[vc + 3] = new Vertex {
                        position = orig - forward - right,
                        uv = new float2(0, 0),
                        color = color,
                    };

                    tris[tc + 0] = vc + 0;
                    tris[tc + 1] = vc + 1;
                    tris[tc + 2] = vc + 2;

                    tris[tc + 3] = vc + 0;
                    tris[tc + 4] = vc + 2;
                    tris[tc + 5] = vc + 3;
                }
            }
        }
    }
}