using System;
using System.Collections.Generic;
using Unity.Burst;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;

namespace UnityEngine.Rendering.UnifiedRayTracing
{
    internal static class TerrainToMesh
    {
		static Dictionary<Terrain, Mesh> cache = new(); //WW1MOD Ok we will keep these, i dont care, its still faster and less crappy then re-doing this all the time for multiple bakes in a row

        static private AsyncTerrainToMeshRequest MakeAsyncTerrainToMeshRequest(int width, int height, Vector3 heightmapScale, float[,] heightmap, bool[,] holes)
        {
			int vertexCount = width * height;
            var job = new ComputeTerrainMeshJob();

			//WW1MOD This crashes when baking 4k terrains, as it tries to allocate a *2GB* Buffer!.. which is fine *here* but when it gets uploaded to GPU, it goes crashy or error, depending on dx-API/Driver
			//Note heightmap size is always pow2+1
			//the "else" part is the original code
			if (height == 4097)
			{
				const int halfSize = 2049;
				int _vertexCount = halfSize * halfSize;

				Debug.Log("Downsizing Terrain heightmap mesh output to 2k to prevent crashes");
				
				job.heightmap = new NativeArray<float>(_vertexCount, Allocator.Persistent);
				for (int i = 0; i < _vertexCount; ++i)
				{
					float h0 = heightmap[((i*2)+0) / (width), ((i*2)+0) % (width)];
					float h1 = heightmap[((i*2)+1) / (width), ((i*2)+0) % (width)];
					float h2 = heightmap[((i*2)+0) / (width), ((i*2)+1) % (width)];
					float h3 = heightmap[((i*2)+1) / (width), ((i*2)+1) % (width)];

					job.heightmap[i] = (h0 + h1 + h2 + h3) / 4f;
				}

				job.holes = new NativeArray<bool>((halfSize - 1) * (halfSize - 1), Allocator.Persistent);
				for (int i = 0; i < (halfSize - 1) * (halfSize - 1); ++i)
				{
					bool h0 = holes[((i*2)+0) / (width - 1), ((i*2)+0) % (width - 1)];
					bool h1 = holes[((i*2)+1) / (width - 1), ((i*2)+0) % (width - 1)];
					bool h2 = holes[((i*2)+0) / (width - 1), ((i*2)+1) % (width - 1)];
					bool h3 = holes[((i*2)+1) / (width - 1), ((i*2)+1) % (width - 1)];

					job.holes[i] = h0 || h1 || h2 || h3;
				}

				height = width = halfSize;
				vertexCount = _vertexCount;
			}
			else
			{
				job.heightmap = new NativeArray<float>(vertexCount, Allocator.Persistent);
				for (int i = 0; i < vertexCount; ++i)
					job.heightmap[i] = heightmap[i / (width), i % (width)];

				job.holes = new NativeArray<bool>((width - 1) * (height - 1), Allocator.Persistent);
				for (int i = 0; i < (width - 1) * (height - 1); ++i)
					job.holes[i] = holes[i / (width - 1), i % (width - 1)];
			}

            job.width = width;
            job.height = height;
            job.heightmapScale = heightmapScale;

            job.positions = new NativeArray<float3>(vertexCount, Allocator.Persistent);
            job.uvs = new NativeArray<float2>(vertexCount, Allocator.Persistent);
            job.normals = new NativeArray<float3>(vertexCount, Allocator.Persistent);
            job.indices = new NativeArray<int>((width - 1) * (height - 1) * 6, Allocator.Persistent);

            JobHandle jobHandle = job.Schedule(vertexCount, math.max(width, 128));

			GC.Collect(GC.MaxGeneration, GCCollectionMode.Forced); //WW1MOD force GC as ive seen in log ends up with like 12GBs of it at this point...

            return new AsyncTerrainToMeshRequest(job, jobHandle);
        }

        static public AsyncTerrainToMeshRequest ConvertAsync(Terrain terrain)
        {
            TerrainData terrainData = terrain.terrainData;
            int width = terrainData.heightmapTexture.width;
            int height = terrainData.heightmapTexture.height;
            float[,] heightmap = terrain.terrainData.GetHeights(0, 0, width, height);
            bool[,] holes = terrain.terrainData.GetHoles(0, 0, width - 1, height - 1);
            return MakeAsyncTerrainToMeshRequest(width, height, terrainData.heightmapScale, heightmap, holes);
        }

		static public AsyncTerrainToMeshRequest ConvertAsync(int heightmapWidth, int heightmapHeight, short[] heightmapData, Vector3 heightmapScale, int holeWidth, int holeHeight, byte[] holedata)
		{
			float[,] heightmap = new float[heightmapWidth, heightmapHeight];
			for (int y = 0; y < heightmapHeight; ++y)
				for (int x = 0; x < heightmapWidth; ++x)
					heightmap[y, x] = (float)heightmapData[y * heightmapWidth + x] / (float)32766;

			bool[,] holes = new bool[heightmapWidth - 1, heightmapHeight - 1];
			if (holedata != null)
			{
				for (int y = 0; y < heightmapHeight - 1; ++y)
					for (int x = 0; x < heightmapWidth - 1; ++x)
						holes[y, x] = holedata[y * holeWidth + x] != 0;
			}
			else
			{
				for (int y = 0; y < heightmapHeight - 1; ++y)
					for (int x = 0; x < heightmapWidth - 1; ++x)
						holes[x, y] = true;
			}
			return MakeAsyncTerrainToMeshRequest(heightmapWidth, heightmapHeight, heightmapScale, heightmap, holes);
		}

		//WW1MOD this function, Add cache and mesh-name to output
		static public Mesh Convert(Terrain terrain)
		{
			if (cache.TryGetValue(terrain, out Mesh mesh) && mesh != null && mesh.vertexCount > 0)
				return mesh;

			var request = ConvertAsync(terrain);
			request.WaitForCompletion();

			mesh = request.GetMesh();
			mesh.name = terrain.name + "_lightbake";
			cache[terrain] = mesh;

			return mesh;

			//Original
			//var request = ConvertAsync(terrain);
			//request.WaitForCompletion();
			//return request.GetMesh();
		}

		static public Mesh Convert(int heightmapWidth, int heightmapHeight, short[] heightmapData, Vector3 heightmapScale, int holeWidth, int holeHeight, byte[] holedata)
		{
			var request = ConvertAsync(heightmapWidth, heightmapHeight, heightmapData, heightmapScale, holeWidth, holeHeight, holedata);
			request.WaitForCompletion();
			return request.GetMesh();
		}
	}

    internal struct AsyncTerrainToMeshRequest
    {
        internal AsyncTerrainToMeshRequest(ComputeTerrainMeshJob job, JobHandle jobHandle)
        {
            m_Job = job;
            m_JobHandle = jobHandle;
        }

        public bool done { get { return m_JobHandle.IsCompleted; } }
        public Mesh GetMesh()
        {
            if (!done)
                return null;

            Mesh mesh = new Mesh();
            mesh.indexFormat = IndexFormat.UInt32;
            mesh.SetVertices(m_Job.positions);
            mesh.SetUVs(0, m_Job.uvs);
            mesh.SetNormals(m_Job.normals);
            mesh.SetIndices(TriangleIndicesWithoutHoles().ToArray(), MeshTopology.Triangles, 0);

            m_Job.DisposeArrays();

            return mesh;
        }

        public void WaitForCompletion()
        {
            m_JobHandle.Complete();
        }

        List<int> TriangleIndicesWithoutHoles()
        {
            var trianglesWithoutHoles = new List<int>((m_Job.width - 1) * (m_Job.height - 1) * 6);
            for (int i = 0; i < m_Job.indices.Length; i += 3)
            {
                int i1 = m_Job.indices[i];
                int i2 = m_Job.indices[i + 1];
                int i3 = m_Job.indices[i + 2];

                if (i1 != 0 && i2 != 0 && i3 != 0)
                {
                    trianglesWithoutHoles.Add(i1);
                    trianglesWithoutHoles.Add(i2);
                    trianglesWithoutHoles.Add(i3);
                }
            }

            if (trianglesWithoutHoles.Count == 0)
            {
                trianglesWithoutHoles.Add(0);
                trianglesWithoutHoles.Add(0);
                trianglesWithoutHoles.Add(0);
            }

            return trianglesWithoutHoles;
        }

        JobHandle m_JobHandle;
        ComputeTerrainMeshJob m_Job;
    }

    [BurstCompile]
    internal struct ComputeTerrainMeshJob : IJobParallelFor
    {
        [ReadOnly]
        public NativeArray<float> heightmap;

        [ReadOnly]
        public NativeArray<bool> holes;

        public int width;
        public int height;
        public float3 heightmapScale;

        public NativeArray<float3> positions;
        public NativeArray<float2> uvs;
        public NativeArray<float3> normals;

        [NativeDisableParallelForRestriction]
        public NativeArray<int> indices;

        public void DisposeArrays()
        {
            heightmap.Dispose();
            holes.Dispose();
            positions.Dispose();
            uvs.Dispose();
            normals.Dispose();
            indices.Dispose();
        }

        public void Execute(int index)
        {
            int vertexIndex = index;
            int x = vertexIndex % width;
            int y = vertexIndex / height;

            float3 v = new float3(x, heightmap[y*width +x], y);

            positions[vertexIndex] = v * heightmapScale;
            uvs[vertexIndex] = v.xz / new float2(width, height);
            normals[vertexIndex] = CalculateTerrainNormal(heightmap, x, y, width, height, heightmapScale);

            if (x < width - 1 && y < height - 1)
            {
                int i1 = y * width + x;
                int i2 = i1 + 1;
                int i3 = i1 + width;
                int i4 = i3 + 1;

                int faceIndex = x + y * (width - 1);

                if (!holes[faceIndex])
                {
                    i1 = i2 = i3 = i4 = 0;
                }

                indices[6* faceIndex + 0] = i1;
                indices[6* faceIndex + 1] = i4;
                indices[6* faceIndex + 2] = i2;

                indices[6* faceIndex + 3] = i1;
                indices[6* faceIndex + 4] = i3;
                indices[6* faceIndex + 5] = i4;
            }
        }

        static float3 CalculateTerrainNormal(NativeArray<float> heightmap, int x, int y, int width, int height, float3 scale)
        {
            float dY, dX;

            dX = SampleHeight(x - 1, y - 1, width, height, heightmap, scale.y) * -1.0F;
            dX += SampleHeight(x - 1, y, width, height, heightmap, scale.y) * -2.0F;
            dX += SampleHeight(x - 1, y + 1, width, height, heightmap, scale.y) * -1.0F;
            dX += SampleHeight(x + 1, y - 1, width, height, heightmap, scale.y) * 1.0F;
            dX += SampleHeight(x + 1, y, width, height, heightmap, scale.y) * 2.0F;
            dX += SampleHeight(x + 1, y + 1, width, height, heightmap, scale.y) * 1.0F;

            dX /= scale.x;

            dY = SampleHeight(x - 1, y - 1, width, height, heightmap, scale.y) * -1.0F;
            dY += SampleHeight(x, y - 1, width, height, heightmap, scale.y) * -2.0F;
            dY += SampleHeight(x + 1, y - 1, width, height, heightmap, scale.y) * -1.0F;
            dY += SampleHeight(x - 1, y + 1, width, height, heightmap, scale.y) * 1.0F;
            dY += SampleHeight(x, y + 1, width, height, heightmap, scale.y) * 2.0F;
            dY += SampleHeight(x + 1, y + 1, width, height, heightmap, scale.y) * 1.0F;
            dY /= scale.z;

            // Cross Product of components of gradient reduces to
            return math.normalize(new float3(-dX, 8, -dY));
        }
        static float SampleHeight(int x, int y, int width, int height, NativeArray<float> heightmap, float scale)
        {
            x = math.clamp(x, 0, width - 1);
            y = math.clamp(y, 0, height - 1);
            return heightmap[x + y* width] * scale;
        }

    }

}