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,
            RenderPass renderPass,
            Framebuffer framebuffer,
            Rect renderArea,
            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(renderPass, framebuffer, renderArea, Vector4.Zero);
            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;
        }
    }
}