using System.Collections.Generic;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Entities;
using Unity.Mathematics;
using Unity.Transforms;
using UnityEngine;
using UnityEngine.Assertions;
using UnityEngine.Experimental.PlayerLoop;
using Unity.Transforms2D;

namespace toinfiniityandbeyond.Rendering
{
    /// <summary>
    /// Renders all Entities containing both SpriteInstanceRenderer & TransformMatrix components.
    /// </summary>
    // [UpdateAfter(typeof(PreLateUpdate.ParticleSystemBeginUpdateAll))]
    // [UnityEngine.ExecuteInEditMode]
    public class SpriteInstanceRendererSystem : ComponentSystem
	{
        private readonly Dictionary<SpriteInstanceRenderer, Material> _cachedMaterialDictionary =
            new Dictionary<SpriteInstanceRenderer, Material>();

        private readonly Dictionary<SpriteInstanceRenderer, Mesh> _cachedMeshDictionary =
            new Dictionary<SpriteInstanceRenderer, Mesh>();

        // Instance renderer takes only batches of 1023
        Matrix4x4[]                 m_MatricesArray = new Matrix4x4[1023];
	    List<SpriteInstanceRenderer>  m_CacheduniqueRendererTypes = new List<SpriteInstanceRenderer>(10);
	    ComponentGroup              m_InstanceRendererGroup;

	    // This is the ugly bit, necessary until Graphics.DrawMeshInstanced supports NativeArrays pulling the data in from a job.
        public unsafe static void CopyMatrices(ComponentDataArray<TransformMatrix> transforms, int beginIndex, int length, Matrix4x4[] outMatrices)
        {
	        // @TODO: This is using unsafe code because the Unity DrawInstances API takes a Matrix4x4[] instead of NativeArray.
	        // We want to use the ComponentDataArray.CopyTo method
	        // because internally it uses memcpy to copy the data,
	        // if the nativeslice layout matches the layout of the component data. It's very fast...
            fixed (Matrix4x4* matricesPtr = outMatrices)
            {
                Assert.AreEqual(sizeof(Matrix4x4), sizeof(TransformMatrix));
	            var matricesSlice = NativeSliceUnsafeUtility.ConvertExistingDataToNativeSlice<TransformMatrix>(matricesPtr, sizeof(Matrix4x4), length);
	            #if ENABLE_UNITY_COLLECTIONS_CHECKS
	            NativeSliceUnsafeUtility.SetAtomicSafetyHandle(ref matricesSlice, AtomicSafetyHandle.GetTempUnsafePtrSliceHandle());
	            #endif
                transforms.CopyTo(matricesSlice, beginIndex);
            }
        }

	    protected override void OnCreateManager(int capacity)
	    {
	        // We want to find all MeshInstanceRenderer & TransformMatrix combinations and render them
	        m_InstanceRendererGroup = GetComponentGroup(typeof(SpriteInstanceRenderer), typeof(TransformMatrix));
	    }

	    protected override void OnUpdate()
		{
		    // We want to iterate over all unique MeshInstanceRenderer shared component data,
		    // that are attached to any entities in the world
            EntityManager.GetAllUniqueSharedComponentDatas(m_CacheduniqueRendererTypes);
            for (int i = 0;i != m_CacheduniqueRendererTypes.Count;i++)
            {
                // For each unique MeshInstanceRenderer data, we want to get all entities with a TransformMatrix
                // SharedComponentData gurantees that all those entities are packed togehter in a chunk with linear memory layout.
                // As a result the copy of the matrices out is internally done via memcpy.
                var renderer = m_CacheduniqueRendererTypes[i];
                if (renderer.sprite == null)
                    continue;
                m_InstanceRendererGroup.SetFilter(renderer);
                var transforms = m_InstanceRendererGroup.GetComponentDataArray<TransformMatrix>();

                Mesh mesh;
                Material material;
                var size = math.max(renderer.sprite.width, renderer.sprite.height) / (float) renderer.pixelsPerUnit;
                float2 meshPivot = renderer.pivot * size;
                if (!_cachedMeshDictionary.TryGetValue(renderer, out mesh))
                {
                    mesh = MeshUtils.GenerateQuad(size, meshPivot);
                    _cachedMeshDictionary.Add(renderer, mesh);
                }

                if (!_cachedMaterialDictionary.TryGetValue(renderer, out material))
                {
                    material = new Material(Shader.Find("Sprites/Instanced"))
                    {
                        enableInstancing = true,
                        mainTexture = renderer.sprite
                    };
                    _cachedMaterialDictionary.Add(renderer, material);
                }

                // Graphics.DrawMeshInstanced has a set of limitations that are not optimal for working with ECS.
                // Specifically:
                // * No way to push the matrices from a job
                // * no NativeArray API, currently uses Matrix4x4[]
                // As a result this code is not yet jobified.
                // We are planning to adjust this API to make it more efficient for this use case.

                // For now, we have to copy our data into Matrix4x4[] with a specific upper limit of how many instances we can render in one batch.
                // So we just have a for loop here, representing each Graphics.DrawMeshInstanced batch
                int beginIndex = 0;
                while (beginIndex < transforms.Length)
                {
                    int length = math.min(m_MatricesArray.Length, transforms.Length - beginIndex);
                    CopyMatrices(transforms, beginIndex, length, m_MatricesArray);
                    Graphics.DrawMeshInstanced(mesh, 0, material, m_MatricesArray, length);

                    beginIndex += length;
                }
            }

		    m_CacheduniqueRendererTypes.Clear();
		}
	}
}