package net.halbear.Terrain4J.EngineCore.Vulkan.GUI;

import imgui.ImDrawData;
import imgui.ImGui;
import imgui.ImGuiIO;
import imgui.ImVec4;
import imgui.glfw.ImGuiImplGlfw;
import imgui.type.ImInt;
import net.halbear.Terrain4J.EngineCore.Input.KeyboardInput;
import net.halbear.Terrain4J.EngineCore.Logic.EngineConfig;
import net.halbear.Terrain4J.EngineCore.Logic.EngineInstance;
import net.halbear.Terrain4J.EngineCore.Logic.Rendering.VulkanContext;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Pipeline.DefaultPipeline;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Pipeline.Images.*;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Pipeline.Pipeline;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Pipeline.PipelineBuildInfo;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Pipeline.PushConstantsRange;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.PostProcessing.PostProcess;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Shader.*;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.CommandBuffer;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.CommandPool;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.Device;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.Queues.Queue;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanBuffer;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanUtils;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.system.MemoryUtil;
import org.lwjgl.util.shaderc.Shaderc;
import org.lwjgl.vulkan.*;
import org.tinylog.Logger;

import java.nio.ByteBuffer;
import java.nio.FloatBuffer;
import java.nio.LongBuffer;
import java.util.*;

import static org.lwjgl.util.vma.Vma.VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT;
import static org.lwjgl.util.vma.Vma.VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;
import static org.lwjgl.vulkan.VK13.*;

public class GuiRenderer {
    private static final String DESCRIPTOR_ID_TEXTURE = "GUI_DESC_ID_TEXT";
    private static final String GUI_FRAGMENT_SHADER_GLSL = "resources/EngineResources/GuiShaders/gui_frag.glsl";
    private static final String GUI_FRAGMENT_SHADER_SPV = GUI_FRAGMENT_SHADER_GLSL + ".spv";
    private static final String GUI_VERTEX_SHADER_GLSL = "resources/EngineResources/GuiShaders/gui_vertex.glsl";
    private static final String GUI_VERTEX_SHADER_SPV = GUI_VERTEX_SHADER_GLSL + ".spv";

    private final VulkanBuffer[] IndexBuffers;
    private final VulkanBuffer[] VertexBuffers;
    private final Texture FontsTexture;
    private final TextureSampler FontsTextureSampler;
    private final Map<Long,Long> GuiTexturesMap;
    private final Pipeline pipeline;
    private final DescriptorSetLayout TextDescriptorSetLayout;
    private VkRenderingAttachmentInfo.Buffer AttachmentInfoColour;
    private VkRenderingInfo RenderingInfo;

    public long GetVkDescriptorImage(long ID){
        return GuiTexturesMap.get(ID);
    }

    public GuiRenderer(EngineInstance engineInstance, VulkanContext VkCtx, Queue queue, Attachment dstAttachment){
        AttachmentInfoColour = CreateColourAttachmentInfo(dstAttachment);
        RenderingInfo = CreateRenderInfo(dstAttachment,AttachmentInfoColour);
        ShaderModule[] shadermodules = CreateShaderModules(VkCtx);
        TextDescriptorSetLayout = new DescriptorSetLayout(VkCtx,new DescriptorSetLayout.LayoutInformation(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,0,1,VK_SHADER_STAGE_FRAGMENT_BIT));
        pipeline = CreatePipeline(VkCtx, shadermodules, new DescriptorSetLayout[]{TextDescriptorSetLayout});
        Arrays.asList(shadermodules).forEach(s->s.CleanUp(VkCtx));

        VertexBuffers = new VulkanBuffer[VulkanUtils.MAX_IN_FLIGHT];
        IndexBuffers = new VulkanBuffer[VulkanUtils.MAX_IN_FLIGHT];

        FontsTexture = InitGUI(VkCtx,queue);
        var textureSamplerInfo = new TextureSamplerInfo(VK_SAMPLER_ADDRESS_MODE_REPEAT, VK_BORDER_COLOR_INT_OPAQUE_BLACK,1,true);
        FontsTextureSampler = new TextureSampler(VkCtx,textureSamplerInfo);
        Device device = VkCtx.GetDevice();
        DescriptorAllocator descriptorAllocator = VkCtx.GetDescriptorAllocator();
        DescriptorSet descriptorSet = descriptorAllocator.AddDescriptorSets(device,DESCRIPTOR_ID_TEXTURE,1,TextDescriptorSetLayout)[0];
        descriptorSet.SetImage(device,FontsTexture.GetImageView(),FontsTextureSampler,TextDescriptorSetLayout.GetLayoutInfo().Binding());
        ImGui.getIO().getFonts().setTexID(descriptorSet.GetVkDescriptorSet());

        KeyboardInput keyboardInput=engineInstance.window().getKeyboardInput();
        keyboardInput.setCharCallBack(new GuiUtils.CharacterCallBack());
        keyboardInput.AddKeyCallBack(new GuiUtils.KeyCallBack());

        GuiTexturesMap = new HashMap<>();
    }

    private static VkRenderingAttachmentInfo.Buffer CreateColourAttachmentInfo(Attachment dstAttachment){
        return VkRenderingAttachmentInfo.calloc(1)
                .sType$Default()
                .imageView(dstAttachment.GetVkImageView().GetVulkanImageView())
                .imageLayout(VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL)
                .loadOp(VK_ATTACHMENT_LOAD_OP_LOAD)
                .storeOp(VK_ATTACHMENT_STORE_OP_STORE);
    }

    private static Pipeline CreatePipeline(VulkanContext VkCtx, ShaderModule[] shaderModules, DescriptorSetLayout[] descriptorSetLayouts){
        var VertexBufferStructure = new GuiVertexBufferStruct();
        var BuildInfo = new PipelineBuildInfo(shaderModules, VertexBufferStructure.GetVertexInput(),
                PostProcess.COLOUR_FORMAT)
                .SetPushConstantRanges(
                        new PushConstantsRange[]{
                            new PushConstantsRange(VK_SHADER_STAGE_VERTEX_BIT,0,VulkanUtils.VEC2_SIZE)
                        })
                .SetDescriptorSetLayouts(descriptorSetLayouts)
                .BlendingIsUsed(true);
        var pipeline = new DefaultPipeline(VkCtx, BuildInfo);
        VertexBufferStructure.CleanUp();
        return pipeline;
    }

    private static VkRenderingInfo CreateRenderInfo(Attachment ColourAttachment, VkRenderingAttachmentInfo.Buffer ColourAttachmentInfo){
        VkRenderingInfo renderingInfo;
        try(var MemStack = MemoryStack.stackPush()){
            Image image = ColourAttachment.GetVkImage();
            VkExtent2D extent2D = VkExtent2D.calloc(MemStack).width(image.GetWidth()).height(image.GetHeight());
            var RenderArea = VkRect2D.calloc(MemStack).extent(extent2D);
            renderingInfo = VkRenderingInfo.calloc()
                    .sType$Default()
                    .renderArea(RenderArea)
                    .layerCount(1)
                    .pColorAttachments(ColourAttachmentInfo);
        }
        return renderingInfo;
    }

    private static ShaderModule[] CreateShaderModules(VulkanContext VkCtx){
        if(EngineConfig.getInstance().RecompileShaders()){
            ShaderCompiler.CompileGLSLShaderOnChange(GUI_VERTEX_SHADER_GLSL, Shaderc.shaderc_glsl_vertex_shader);
            ShaderCompiler.CompileGLSLShaderOnChange(GUI_FRAGMENT_SHADER_GLSL, Shaderc.shaderc_glsl_fragment_shader);
        }
        return new ShaderModule[]{
                new ShaderModule(VkCtx, VK_SHADER_STAGE_VERTEX_BIT, GUI_VERTEX_SHADER_SPV, null),
                new ShaderModule(VkCtx, VK_SHADER_STAGE_FRAGMENT_BIT, GUI_FRAGMENT_SHADER_SPV, null)
        };
    }

    public void LoadTextures(VulkanContext VkCtx, List<GuiTexture> GuiTextures, TextureCache textureCache){
        if(GuiTextures == null){
            return;
        }
        DescriptorAllocator descriptorAllocator = VkCtx.GetDescriptorAllocator();
        int TextureCount = GuiTextures.size();
        Device device = VkCtx.GetDevice();
        for(int i = 0; i < TextureCount; i++){
            var guiTexture = GuiTextures.get(i);
            String descriptorID = guiTexture.TexturePath();
            DescriptorSet descriptorSet = descriptorAllocator.AddDescriptorSets(device, descriptorID, 1, TextDescriptorSetLayout)[0];
            Texture texture = textureCache.GetTexture(guiTexture.TexturePath());
            descriptorSet.SetImage(device,texture.GetImageView(),FontsTextureSampler,TextDescriptorSetLayout.GetLayoutInfo().Binding());
            //Logger.debug("TEXTURE ID -> [{}] DESCRIPTOR SET -> [{}] TEXTURE -> [{}] IMAGE PATH -> [{}] TEXTURE SAMPLER ->[{}]",guiTexture.ID(),descriptorSet.GetVkDescriptorSet(),texture,guiTexture.TexturePath(),FontsTextureSampler);
            GuiTexturesMap.put(guiTexture.ID(),descriptorSet.GetVkDescriptorSet());
        }
    }

    private static Texture InitGUI(VulkanContext VkCtx, Queue queue){
        ImGui.createContext();
        ImGuiIO imGuiIO = ImGui.getIO();
        imGuiIO.setIniFilename(null);
        VkExtent2D SwapChainExtent = VkCtx.GetSwapChain().GetSwapChainExtent();
        imGuiIO.setDisplaySize(SwapChainExtent.width(), SwapChainExtent.height());
        imGuiIO.setDisplayFramebufferScale(1.0f,1.0f);

        ImInt textureWidth = new ImInt();
        ImInt textureHeight = new ImInt();
        ByteBuffer buffer = imGuiIO.getFonts().getTexDataAsRGBA32(textureWidth,textureHeight);
        ImageSrc imageSrc = new ImageSrc(buffer,textureWidth.get(),textureHeight.get(),4);
        Texture FontsTexture = new Texture(VkCtx, "GUI_TEXTURE",imageSrc,VK_FORMAT_R8G8B8A8_SRGB);

        var commandPool = new CommandPool(VkCtx, queue.GetQueueFamilyIndex(), false);
        var commandBuffer = new CommandBuffer(VkCtx, commandPool, true, true);
        commandBuffer.BeginRecording();
        FontsTexture.RecordTextureTransition(commandBuffer);
        commandBuffer.EndRecording();
        commandBuffer.SubmitAndWait(VkCtx, queue);
        commandBuffer.cleanup(VkCtx,commandPool);
        commandPool.cleanup(VkCtx);
        return FontsTexture;
    }

    public void Render(VulkanContext VkCtx, CommandBuffer commandBuffer, int CurrentFrame, Attachment dstAttachment){
        try(var MemStack = MemoryStack.stackPush()){
            UpdateBuffers(VkCtx, CurrentFrame);
            if(VertexBuffers[CurrentFrame] == null){
                return;
            }
            Image dstImage = dstAttachment.GetVkImage();
            int Width = dstImage.GetWidth();
            int Height = dstImage.GetHeight();

            VkCommandBuffer CommandHandle = commandBuffer.GetVulkanCommandBuffer();

            vkCmdBeginRendering(CommandHandle, RenderingInfo);
            vkCmdBindPipeline(CommandHandle, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline.GetVulkanPipeline());

            var Viewport = VkViewport.calloc(1,MemStack)
                    .x(0)
                    .y(Height)
                    .height(-Height)
                    .width(Width)
                    .minDepth(0.0f)
                    .maxDepth(1.0f);
            vkCmdSetViewport(CommandHandle,0,Viewport);

            LongBuffer VertexBuff = MemStack.mallocLong(1);
            VertexBuff.put(0,VertexBuffers[CurrentFrame].GetBuffer());
            LongBuffer Offsets = MemStack.mallocLong(1);
            Offsets.put(0,0L);
            vkCmdBindVertexBuffers(CommandHandle,0,VertexBuff,Offsets);
            vkCmdBindIndexBuffer(CommandHandle,IndexBuffers[CurrentFrame].GetBuffer(),0,VK_INDEX_TYPE_UINT16);

            ImGuiIO imGuiIO = ImGui.getIO();
            FloatBuffer PushConstantBuffer = MemStack.mallocFloat(2);
            PushConstantBuffer.put(0,2.0f/imGuiIO.getDisplaySizeX());
            PushConstantBuffer.put(1,-2.0f/imGuiIO.getDisplaySizeY());
            vkCmdPushConstants(CommandHandle,pipeline.GetVulkanPipelineLayout(),VK_SHADER_STAGE_VERTEX_BIT,0,PushConstantBuffer);

            LongBuffer DescriptorSets = MemStack.mallocLong(1);

            ImVec4 imVec4 = new ImVec4();
            VkRect2D.Buffer rect = VkRect2D.calloc(1,MemStack);
            ImDrawData imDrawData = ImGui.getDrawData();
            int CmdListCount = imDrawData.getCmdListsCount();
            int OffsetIndex = 0;
            int OffsetVertex = 0;
            for(int i = 0; i < CmdListCount; i++){
                int CommandBufferSize = imDrawData.getCmdListCmdBufferSize(i);
                for(int j = 0; j < CommandBufferSize; j++){
                    long TextureID = imDrawData.getCmdListCmdBufferTextureId(i,j);
                    Long TextureDescriptorSet = GuiTexturesMap.get(TextureID);
                    if(TextureDescriptorSet == null){
                        TextureDescriptorSet = TextureID;
                    }
                    DescriptorSets.put(0,TextureDescriptorSet);
                    vkCmdBindDescriptorSets(CommandHandle,VK_PIPELINE_BIND_POINT_GRAPHICS,pipeline.GetVulkanPipelineLayout(),0,DescriptorSets,null);

                    imDrawData.getCmdListCmdBufferClipRect(imVec4,i,j);
                    rect.offset(it->it.x((int)Math.max(imVec4.x,0)).y((int)Math.max(imVec4.y,1)));
                    rect.extent(it->it.width((int)(imVec4.z - imVec4.x)).height((int)(imVec4.w - imVec4.y)));
                    vkCmdSetScissor(CommandHandle,0,rect);
                    int ElementCount = imDrawData.getCmdListCmdBufferElemCount(i,j);
                    vkCmdDrawIndexed(CommandHandle,ElementCount,1,
                            OffsetIndex + imDrawData.getCmdListCmdBufferIdxOffset(i,j),
                            OffsetVertex + imDrawData.getCmdListCmdBufferVtxOffset(i,j),0);
                }
                OffsetIndex += imDrawData.getCmdListIdxBufferSize(i);
                OffsetVertex += imDrawData.getCmdListVtxBufferSize(i);
            }
            vkCmdEndRendering(CommandHandle);
        }
    }

    private void UpdateBuffers(VulkanContext VkCtx, int Index){
        ImDrawData imDrawData = ImGui.getDrawData();

        if(imDrawData.ptr == 0) return;
        int VertexBufferSize = imDrawData.getTotalVtxCount() * GuiVertexBufferStruct.VERTEX_SIZE;
        int IndexBufferSize = imDrawData.getTotalIdxCount() * VulkanUtils.SHORT_LENGTH;
        if(VertexBufferSize == 0 || IndexBufferSize == 0) return;
        var VertexBuff = VertexBuffers[Index];
        if(VertexBuff == null || VertexBufferSize > VertexBuff.GetRequestedSize()){
            if(VertexBuff != null){
                VertexBuff.cleanup(VkCtx);
            }
            VertexBuff = new VulkanBuffer(VkCtx, VertexBufferSize, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE,
                    VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT);
            VertexBuffers[Index] = VertexBuff;
        }

        var IndiciesBuffer = IndexBuffers[Index];
        if(IndiciesBuffer == null || IndexBufferSize > IndiciesBuffer.GetRequestedSize()){
            if(IndiciesBuffer != null){
                IndiciesBuffer.cleanup(VkCtx);
            }
            IndiciesBuffer = new VulkanBuffer(VkCtx,IndexBufferSize,VK_BUFFER_USAGE_INDEX_BUFFER_BIT,VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE,
                    VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT);
            IndexBuffers[Index] = IndiciesBuffer;
        }

        ByteBuffer dstVertexBuffer = MemoryUtil.memByteBuffer(VertexBuff.MapMemory(VkCtx),VertexBufferSize);
        ByteBuffer dstIndexBuffer = MemoryUtil.memByteBuffer(IndiciesBuffer.MapMemory(VkCtx),IndexBufferSize);

        int CommandListCount = imDrawData.getCmdListsCount();
        for(int i = 0; i < CommandListCount; i++){
            ByteBuffer ImGuiVertexBuffer = imDrawData.getCmdListVtxBufferData(i);
            dstVertexBuffer.put(ImGuiVertexBuffer);
            ByteBuffer ImGuiIndicesBuffer = imDrawData.getCmdListIdxBufferData(i);
            dstIndexBuffer.put(ImGuiIndicesBuffer);
        }
        VertexBuff.Flush(VkCtx);
        IndiciesBuffer.Flush(VkCtx);
        VertexBuff.UnMapMemory(VkCtx);
        IndiciesBuffer.UnMapMemory(VkCtx);
    }

    public void Resize(VulkanContext VkCtx, Attachment dstAttachment){
        ImGuiIO imGuiIO = ImGui.getIO();
        VkExtent2D SwapChainExtent = VkCtx.GetSwapChain().GetSwapChainExtent();
        imGuiIO.setDisplaySize(SwapChainExtent.width(),SwapChainExtent.height());

        RenderingInfo.free();
        AttachmentInfoColour.free();
        AttachmentInfoColour = CreateColourAttachmentInfo(dstAttachment);
        RenderingInfo = CreateRenderInfo(dstAttachment,AttachmentInfoColour);
    }

    public void CleanUp(VulkanContext VkCtx){
        FontsTextureSampler.CleanUp(VkCtx);
        FontsTexture.CleanUp(VkCtx);
        TextDescriptorSetLayout.CleanUp(VkCtx);
        pipeline.CleanUp(VkCtx);
        Arrays.stream(VertexBuffers).filter(Objects::nonNull).forEach(b->b.cleanup(VkCtx));
        Arrays.stream(IndexBuffers).filter(Objects::nonNull).forEach(b->b.cleanup(VkCtx));
        RenderingInfo.free();
        AttachmentInfoColour.free();
    }
}
