using System.IO;
using System.Runtime.InteropServices;
using MoonWorks.Graphics;
using MoonWorks.Math;

namespace MoonWorksComputeSpriteBatch
{
	[StructLayout(LayoutKind.Sequential)]
	public struct SpriteBatchUniforms
	{
		public uint VertexCount { get; set; }
	}

	public class SpriteBatch
	{
		public const int MAX_SPRITES = 65536;
		private const int MAX_VERTICES = MAX_SPRITES * 4;
		private const int MAX_INDICES = MAX_SPRITES * 6;
		private const int MAX_MATRICES = MAX_SPRITES;

		private static ComputePipeline ComputePipeline = null;

		private Buffer VertexBuffer { get; }
		private Buffer IndexBuffer { get; }
		private Buffer TransformBuffer { get; }

		private readonly VertexPositionTexcoord[] Vertices;
		private static readonly short[] Indices = GenerateIndexArray();
		private readonly Matrix4x4[] Transforms;

		private Texture CurrentTexture { get; set; }
		private Sampler CurrentSampler { get; set; }
		private uint VertexCount { get; set; }
		private uint TransformCount { get; set; }

		public SpriteBatch(GraphicsDevice graphicsDevice)
		{
			if (ComputePipeline == null)
			{
				var computeShaderModule = new ShaderModule(graphicsDevice, Path.Combine(System.Environment.CurrentDirectory, "Content", "spritebatch.comp.spv"));

				var computeShaderState = new ShaderStageState
				{
					ShaderModule = computeShaderModule,
					EntryPointName = "main",
					UniformBufferSize = (uint)Marshal.SizeOf<SpriteBatchUniforms>()
				};

				ComputePipeline = new ComputePipeline(graphicsDevice, computeShaderState, 2, 0);
			}

			Vertices = new VertexPositionTexcoord[MAX_VERTICES];
			VertexBuffer = new Buffer(
				graphicsDevice,
				BufferUsageFlags.Vertex | BufferUsageFlags.Compute,
				(uint)(MAX_VERTICES * Marshal.SizeOf<VertexPositionTexcoord>())
			);

			IndexBuffer = new Buffer(
				graphicsDevice,
				BufferUsageFlags.Index | BufferUsageFlags.Compute,
				MAX_INDICES * sizeof(short)
			);

			Transforms = new Matrix4x4[MAX_MATRICES];
			TransformBuffer = new Buffer(
				graphicsDevice,
				BufferUsageFlags.Compute,
				(uint)(MAX_MATRICES * Marshal.SizeOf<Matrix4x4>())
			);

			var commandBuffer = graphicsDevice.AcquireCommandBuffer();
			commandBuffer.SetBufferData(IndexBuffer, Indices);
			graphicsDevice.Submit(commandBuffer);
		}

		public void Start(Texture texture, Sampler sampler)
		{
			TransformCount = 0;
			VertexCount = 0;
			CurrentTexture = texture;
			CurrentSampler = sampler;
		}

		public void Add(Sprite sprite, Matrix4x4 transform, Color color)
		{
			Vertices[VertexCount].position.X = 0;
			Vertices[VertexCount].position.Y = 0;
			Vertices[VertexCount].texcoord.X = sprite.Texcoord.X;
			Vertices[VertexCount].texcoord.Y = sprite.Texcoord.Y;
			Vertices[VertexCount].color.X = color.R / 255f;
			Vertices[VertexCount].color.Y = color.G / 255f;
			Vertices[VertexCount].color.Z = color.B / 255f;
			Vertices[VertexCount].color.W = color.A / 255f;

			Vertices[VertexCount + 1].position.X = sprite.Width;
			Vertices[VertexCount + 1].position.Y = 0;
			Vertices[VertexCount + 1].texcoord.X = sprite.Texcoord.X + sprite.Texcoord.W;
			Vertices[VertexCount + 1].texcoord.Y = sprite.Texcoord.Y;
			Vertices[VertexCount + 1].color.X = color.R / 255f;
			Vertices[VertexCount + 1].color.Y = color.G / 255f;
			Vertices[VertexCount + 1].color.Z = color.B / 255f;
			Vertices[VertexCount + 1].color.W = color.A / 255f;

			Vertices[VertexCount + 2].position.X = 0;
			Vertices[VertexCount + 2].position.Y = sprite.Height;
			Vertices[VertexCount + 2].texcoord.X = sprite.Texcoord.X;
			Vertices[VertexCount + 2].texcoord.Y = sprite.Texcoord.Y + sprite.Texcoord.H;
			Vertices[VertexCount + 2].color.X = color.R / 255f;
			Vertices[VertexCount + 2].color.Y = color.G / 255f;
			Vertices[VertexCount + 2].color.Z = color.B / 255f;
			Vertices[VertexCount + 2].color.W = color.A / 255f;

			Vertices[VertexCount + 3].position.X = sprite.Width;
			Vertices[VertexCount + 3].position.Y = sprite.Height;
			Vertices[VertexCount + 3].texcoord.X = sprite.Texcoord.X + sprite.Texcoord.W;
			Vertices[VertexCount + 3].texcoord.Y = sprite.Texcoord.Y + sprite.Texcoord.H;
			Vertices[VertexCount + 3].color.X = color.R / 255f;
			Vertices[VertexCount + 3].color.Y = color.G / 255f;
			Vertices[VertexCount + 3].color.Z = color.B / 255f;
			Vertices[VertexCount + 3].color.W = color.A / 255f;

			VertexCount += 4;

			Transforms[TransformCount] = transform;
			TransformCount += 1;
		}

		public void Flush(
			CommandBuffer commandBuffer,
			ColorAttachmentInfo colorAttachmentInfo,
			GraphicsPipeline graphicsPipeline,
			CameraUniforms cameraUniforms
		)
		{
			if (VertexCount == 0)
			{
				return;
			}

			commandBuffer.SetBufferData(VertexBuffer, Vertices, 0, 0, VertexCount);
			commandBuffer.SetBufferData(TransformBuffer, Transforms, 0, 0, TransformCount);

			commandBuffer.BindComputePipeline(ComputePipeline);
			commandBuffer.BindComputeBuffers(VertexBuffer, TransformBuffer);
			var offset = commandBuffer.PushComputeShaderUniforms(new SpriteBatchUniforms
			{
				VertexCount = VertexCount
			});
			commandBuffer.DispatchCompute(System.Math.Max(1, VertexCount / 256), 1, 1, offset);

			commandBuffer.BeginRenderPass(colorAttachmentInfo);
			commandBuffer.BindGraphicsPipeline(graphicsPipeline);
			commandBuffer.BindVertexBuffers(VertexBuffer);
			commandBuffer.BindIndexBuffer(IndexBuffer, IndexElementSize.Sixteen);
			commandBuffer.BindFragmentSamplers(new TextureSamplerBinding { Texture = CurrentTexture, Sampler = CurrentSampler });
			offset = commandBuffer.PushVertexShaderUniforms(cameraUniforms);

			commandBuffer.DrawIndexedPrimitives(
				0,
				0,
				VertexCount / 2,
				offset,
				0
			);

			commandBuffer.EndRenderPass();

			VertexCount = 0;
			TransformCount = 0;
		}

		private static short[] GenerateIndexArray()
		{
			var result = new short[MAX_INDICES];
			for (int i = 0, j = 0; i < MAX_INDICES; i += 6, j += 4)
			{
				result[i] = (short)j;
				result[i + 1] = (short)(j + 1);
				result[i + 2] = (short)(j + 2);
				result[i + 3] = (short)(j + 2);
				result[i + 4] = (short)(j + 1);
				result[i + 5] = (short)(j + 3);
			}
			return result;
		}
	}
}