using MoonWorks.Graphics;
using MoonWorks;
using ImGuiNET;
using System.Runtime.InteropServices;
using System.IO;
using MoonWorks.Input;
using MoonWorks.Math;
using System.Collections.Generic;

namespace MoonWorksImGuiExample
{
	class MoonWorksImGuiExampleGame : Game
	{
		private string ContentPath = Path.Combine(System.AppDomain.CurrentDomain.BaseDirectory, "Content");

		private ImFontPtr DefaultFontPtr;
		private int TextureID;
		private System.IntPtr? FontTexturePtr;
		private Dictionary<System.IntPtr, Texture> LoadedTextures = new Dictionary<System.IntPtr, Texture>();
		private Dictionary<Texture, System.IntPtr> TextureHandles = new Dictionary<Texture, System.IntPtr>();

		private uint VertexCount = 0;
		private uint IndexCount = 0;
		private Buffer ImGuiVertexBuffer = null;
		private Buffer ImGuiIndexBuffer = null;
		private GraphicsPipeline ImGuiPipeline;
		private ShaderModule ImGuiVertexShader;
		private ShaderModule ImGuiFragmentShader;
		private Sampler ImGuiSampler;

		public MoonWorksImGuiExampleGame(
			WindowCreateInfo windowCreateInfo,
			PresentMode presentMode,
			bool debugMode
		) : base(windowCreateInfo, presentMode, 60, debugMode)
		{
			var context = ImGui.CreateContext();
			ImGui.SetCurrentContext(context);

			var io = ImGui.GetIO();

			io.Fonts.AddFontDefault();

			RebuildFontAtlas();

			io.DisplaySize = new System.Numerics.Vector2(
				windowCreateInfo.WindowWidth,
				windowCreateInfo.WindowHeight
			);
			io.DisplayFramebufferScale = System.Numerics.Vector2.One;

			ImGuiVertexShader = new ShaderModule(GraphicsDevice, Path.Combine(ContentPath, "ImGuiVert.spv"));
			ImGuiFragmentShader = new ShaderModule(GraphicsDevice, Path.Combine(ContentPath, "ImGuiFrag.spv"));
			ImGuiSampler = new Sampler(GraphicsDevice, SamplerCreateInfo.PointClamp);

			ImGuiPipeline = new GraphicsPipeline(
				GraphicsDevice,
				new GraphicsPipelineCreateInfo
				{
					AttachmentInfo = new GraphicsPipelineAttachmentInfo(
						new ColorAttachmentDescription(
							GraphicsDevice.GetSwapchainFormat(Window),
							ColorAttachmentBlendState.NonPremultiplied
						)
					),
					DepthStencilState = DepthStencilState.Disable,
					VertexShaderInfo = GraphicsShaderInfo.Create<Matrix4x4>(ImGuiVertexShader, "main", 0),
					FragmentShaderInfo = GraphicsShaderInfo.Create(ImGuiFragmentShader, "main", 1),
					VertexInputState = new VertexInputState(
						VertexBinding.Create<VertexPositionTextureColor>(),
						VertexAttribute.Create<VertexPositionTextureColor>("Position", 0),
						VertexAttribute.Create<VertexPositionTextureColor>("Texture", 1),
						VertexAttribute.Create<VertexPositionTextureColor>("Color", 2)
					),
					PrimitiveType = PrimitiveType.TriangleList,
					RasterizerState = RasterizerState.CW_CullNone,
					ViewportState = new ViewportState((int) windowCreateInfo.WindowWidth, (int) windowCreateInfo.WindowHeight),
					MultisampleState = MultisampleState.None
				}
			);
		}

		protected override void Update(System.TimeSpan dt)
		{
			// Insert your game update logic here.
		}

		private System.IntPtr BindTexture(Texture texture)
		{
			var id = new System.IntPtr(TextureID);
			LoadedTextures.Add(id, texture);
			TextureHandles.Add(texture, id);
			return id;
		}

		private void UnbindTexture(System.IntPtr textureID)
		{
			var texture = LoadedTextures[textureID];
			LoadedTextures.Remove(textureID);
			TextureHandles.Remove(texture);
		}

		private void RebuildFontAtlas()
		{
			var commandBuffer = GraphicsDevice.AcquireCommandBuffer();

			var io = ImGui.GetIO();

			io.Fonts.GetTexDataAsRGBA32(
				out System.IntPtr pixelData,
				out int width,
				out int height,
				out int bytesPerPixel
			);

			var texture = Texture.CreateTexture2D(
				GraphicsDevice,
				(uint) width,
				(uint) height,
				TextureFormat.R8G8B8A8,
				TextureUsageFlags.Sampler
			);

			commandBuffer.SetTextureData(texture, pixelData, (uint) (width * height * bytesPerPixel));

			GraphicsDevice.Submit(commandBuffer);

			if (FontTexturePtr.HasValue)
			{
				UnbindTexture(FontTexturePtr.Value);
			}

			FontTexturePtr = BindTexture(texture);

			io.Fonts.SetTexID(FontTexturePtr.Value);
			io.Fonts.ClearTexData();
		}

		private void UpdateImGuiBuffers(ImDrawDataPtr drawDataPtr)
		{
			if (drawDataPtr.TotalVtxCount == 0) { return; }

			var commandBuffer = GraphicsDevice.AcquireCommandBuffer();

			var vertexSize = Marshal.SizeOf<VertexPositionTextureColor>();

			if (drawDataPtr.TotalVtxCount > VertexCount)
			{
				ImGuiVertexBuffer?.Dispose();

				VertexCount = (uint) (drawDataPtr.TotalVtxCount * 1.5f);
				ImGuiVertexBuffer = Buffer.Create<VertexPositionTextureColor>(
					GraphicsDevice,
					BufferUsageFlags.Vertex,
					VertexCount
				);
			}

			if (drawDataPtr.TotalIdxCount > IndexCount)
			{
				ImGuiIndexBuffer?.Dispose();

				IndexCount = (uint) (drawDataPtr.TotalIdxCount * 1.5f);
				ImGuiIndexBuffer = Buffer.Create<ushort>(
					GraphicsDevice,
					BufferUsageFlags.Index,
					IndexCount
				);
			}

			uint vertexOffset = 0;
			uint indexOffset = 0;

			for (var n = 0; n < drawDataPtr.CmdListsCount; n += 1)
			{
				var cmdList = drawDataPtr.CmdListsRange[n];

				commandBuffer.SetBufferData(
					ImGuiVertexBuffer,
					cmdList.VtxBuffer.Data,
					vertexOffset,
					(uint) (cmdList.VtxBuffer.Size * vertexSize)
				);

				commandBuffer.SetBufferData(
					ImGuiIndexBuffer,
					cmdList.IdxBuffer.Data,
					indexOffset,
					(uint) cmdList.IdxBuffer.Size * sizeof(ushort)
				);

				vertexOffset += (uint) cmdList.VtxBuffer.Size;
				indexOffset += (uint) cmdList.IdxBuffer.Size;
			}

			GraphicsDevice.Submit(commandBuffer);
		}

		protected override void Draw(System.TimeSpan dt, double alpha)
		{
			// Replace this with your own drawing code.

			var io = ImGui.GetIO();

			io.KeyShift = Inputs.Keyboard.IsDown(Keycode.LeftShift) || Inputs.Keyboard.IsDown(Keycode.RightShift);
			io.KeyCtrl = Inputs.Keyboard.IsDown(Keycode.LeftControl) || Inputs.Keyboard.IsDown(Keycode.RightControl);
			io.KeyAlt = Inputs.Keyboard.IsDown(Keycode.LeftAlt) || Inputs.Keyboard.IsDown(Keycode.RightAlt);
			io.KeySuper = Inputs.Keyboard.IsDown(Keycode.LeftMeta) || Inputs.Keyboard.IsDown(Keycode.RightMeta);
			io.KeyMap[(int) ImGuiKey.Backspace] = (int) Keycode.Backspace;

			io.MousePos = new System.Numerics.Vector2(Inputs.Mouse.X, Inputs.Mouse.Y);
			io.MouseDown[0] = Inputs.Mouse.LeftButton.IsDown;
			io.MouseDown[1] = Inputs.Mouse.RightButton.IsDown;
			io.MouseDown[2] = Inputs.Mouse.MiddleButton.IsDown;
			io.MouseWheel = Inputs.Mouse.Wheel > 0 ? 1 : Inputs.Mouse.Wheel < 0 ? -1 : 0;

			ImGui.NewFrame();

			ImGui.SetNextWindowSize(new System.Numerics.Vector2(400, 200));
			ImGui.Begin("Help");
			ImGui.Columns(2);

			ImGui.Text("F7");
			ImGui.NextColumn();
			ImGui.Text("Pause Action");
			ImGui.NextColumn();

			ImGui.Text("Ctrl-E");
			ImGui.NextColumn();
			ImGui.Text("Entity Search");
			ImGui.NextColumn();

			ImGui.Text("Left Click");
			ImGui.NextColumn();
			ImGui.Text("Inspect Entity");
			ImGui.NextColumn();

			ImGui.Text("Ctrl-H");
			ImGui.NextColumn();
			ImGui.Text("Help");
			ImGui.NextColumn();

			ImGui.End();

			ImGui.Render();

			var drawDataPtr = ImGui.GetDrawData();
			UpdateImGuiBuffers(drawDataPtr);

			if (ImGuiVertexBuffer == null) { return; }

			var commandBuffer = GraphicsDevice.AcquireCommandBuffer();
			var swapchainTexture = commandBuffer.AcquireSwapchainTexture(Window);

			if (swapchainTexture != null)
			{
				commandBuffer.BeginRenderPass(
					new ColorAttachmentInfo(swapchainTexture, Color.Transparent)
				);

				commandBuffer.BindGraphicsPipeline(ImGuiPipeline);

				var vertexUniformOffset = commandBuffer.PushVertexShaderUniforms(
					Matrix4x4.CreateOrthographicOffCenter(0, io.DisplaySize.X, io.DisplaySize.Y, 0, -1, 1)
				);

				commandBuffer.BindVertexBuffers(ImGuiVertexBuffer);
				commandBuffer.BindIndexBuffer(ImGuiIndexBuffer, IndexElementSize.Sixteen);

				uint vertexOffset = 0;
				uint indexOffset = 0;

				for (int n = 0; n < drawDataPtr.CmdListsCount; n += 1)
				{
					var cmdList = drawDataPtr.CmdListsRange[n];

					for (int cmdIndex = 0; cmdIndex < cmdList.CmdBuffer.Size; cmdIndex += 1)
					{
						var drawCmd = cmdList.CmdBuffer[cmdIndex];

						commandBuffer.BindFragmentSamplers(
							new TextureSamplerBinding(LoadedTextures[drawCmd.TextureId], ImGuiSampler)
						);

						commandBuffer.DrawIndexedPrimitives(
							vertexOffset,
							indexOffset,
							drawCmd.ElemCount / 3,
							vertexUniformOffset,
							0
						);

						indexOffset += drawCmd.ElemCount;
					}

					vertexOffset += (uint) cmdList.VtxBuffer.Size;
				}

				commandBuffer.EndRenderPass();
			}

			GraphicsDevice.Submit(commandBuffer);
		}

		protected override void OnDestroy()
		{

		}
	}
}