package net.halbear.Terrain4J.EngineCore.Vulkan.Structure.SwapChain;

import net.halbear.Terrain4J.EngineCore.Logic.EngineConfig;
import net.halbear.Terrain4J.EngineCore.Logic.Rendering.VulkanContext;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.EmptyVertexBufferStruct;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Pipeline.DefaultPipeline;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Pipeline.Images.Attachment;
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.Shader.*;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.CommandBuffer;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.Device;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanUtils;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.util.shaderc.Shaderc;
import org.lwjgl.vulkan.*;

import java.nio.LongBuffer;
import java.util.Arrays;

import static org.lwjgl.vulkan.KHRSwapchain.VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
import static org.lwjgl.vulkan.KHRSynchronization2.VK_IMAGE_LAYOUT_ATTACHMENT_OPTIMAL_KHR;
import static org.lwjgl.vulkan.VK10.*;
import static org.lwjgl.vulkan.VK13.*;

public class SwapChainRender {

    private static final String DESCRIPTOR_ID_ATTACHMENT = "FWD_DESC_ID_ATT";
    private static final String FRAGMENT_SHADER_FILE_GLSL = "resources/EngineResources/shaders/swap_frag.glsl";
    private static final String FRAGMENT_SHADER_FILE_SPV = FRAGMENT_SHADER_FILE_GLSL + ".spv";
    private static final String VERTEX_SHADER_FILE_GLSL = "resources/EngineResources/shaders/swap_vtx.glsl";
    private static final String VERTEX_SHADER_FILE_SPV = VERTEX_SHADER_FILE_GLSL + ".spv";
    private final DescriptorSetLayout AttachmentDescriptorSetLayout;
    private final VkClearValue ClearValueColour;
    private final Pipeline pipeline;
    private final TextureSampler textureSampler;
    private VkRenderingAttachmentInfo.Buffer[] ColourAttachmentsInfo;
    private VkRenderingInfo[] RenderInfo;

    public SwapChainRender(VulkanContext VkCtx, Attachment SrcAttachment){
        ClearValueColour = VkClearValue.calloc();
        ClearValueColour.color(c->c.float32(0,0.0f).float32(1,0.0f).float32(2,0.0f).float32(3,0.0f));

        ColourAttachmentsInfo = CreateColourAttachmentsInfo(VkCtx,ClearValueColour);
        RenderInfo = CreateRenderInfo(VkCtx,ColourAttachmentsInfo);

        var TextureSamplerInfo = new TextureSamplerInfo(VK_SAMPLER_ADDRESS_MODE_REPEAT,VK_BORDER_COLOR_INT_OPAQUE_BLACK,1,true);
        textureSampler = new TextureSampler(VkCtx,TextureSamplerInfo);

        var LayoutInfo = new DescriptorSetLayout.LayoutInformation(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,0,1,VK_SHADER_STAGE_FRAGMENT_BIT);
        AttachmentDescriptorSetLayout = new DescriptorSetLayout(VkCtx,LayoutInfo);
        CreateAttachmentDescriptorSet(VkCtx,AttachmentDescriptorSetLayout,SrcAttachment,textureSampler);

        ShaderModule[] shaderModules = CreateShaderModules(VkCtx);

        pipeline = CreatePipeline(VkCtx,shaderModules,new DescriptorSetLayout[]{AttachmentDescriptorSetLayout});
        Arrays.asList(shaderModules).forEach(s->s.CleanUp(VkCtx));
    }

    private static void CreateAttachmentDescriptorSet(VulkanContext VkCtx, DescriptorSetLayout descriptorSetLayout, Attachment attachment, TextureSampler sampler){
        DescriptorAllocator descriptorAllocator = VkCtx.GetDescriptorAllocator();
        Device device = VkCtx.GetDevice();
        DescriptorSet descriptorSet = descriptorAllocator.AddDescriptorSets(device,DESCRIPTOR_ID_ATTACHMENT,1,descriptorSetLayout)[0];
        descriptorSet.SetImage(device,attachment.GetVkImageView(),sampler,0);
    }

    private static VkRenderingAttachmentInfo.Buffer[] CreateColourAttachmentsInfo(VulkanContext VkCtx, VkClearValue ClearView){
        SwapChain swapChain = VkCtx.GetSwapChain();
        int ImageCount = swapChain.GetImageCount();
        var Result = new VkRenderingAttachmentInfo.Buffer[ImageCount];

        for(int i = 0; i < ImageCount; i++){
            var Attachments = VkRenderingAttachmentInfo.calloc(1);
            Attachments.get(0)
                    .sType$Default()
                    .imageView(swapChain.GetVulkanImageView(i).GetVulkanImageView())
                    .imageLayout(VK_IMAGE_LAYOUT_ATTACHMENT_OPTIMAL_KHR)
                    .loadOp(VK_ATTACHMENT_LOAD_OP_CLEAR)
                    .storeOp(VK_ATTACHMENT_STORE_OP_STORE)
                    .clearValue(ClearView);
            Result[i] = Attachments;
        }
        return Result;
    }
    private static Pipeline CreatePipeline(VulkanContext VkCtx, ShaderModule[] shaderModules, DescriptorSetLayout[] descriptorSetLayouts){
        var VertexBufferStruct = new EmptyVertexBufferStruct();
        var BuildInfo = new PipelineBuildInfo(shaderModules, VertexBufferStruct.GetVertexInput(),VkCtx.GetSurface().GetSurfaceFormat().ImageFormat()).SetDescriptorSetLayouts(descriptorSetLayouts).BlendingIsUsed(true);
        var pipeline = new DefaultPipeline(VkCtx,BuildInfo);
        VertexBufferStruct.CleanUp();
        return pipeline;
    }

    private static VkRenderingInfo[] CreateRenderInfo(VulkanContext VkCtx, VkRenderingAttachmentInfo.Buffer[] ColourAttachments){
        SwapChain swapChain = VkCtx.GetSwapChain();
        int ImageCount = swapChain.GetImageCount();
        var Result = new VkRenderingInfo[ImageCount];

        try(var MemStack = MemoryStack.stackPush()){
            VkExtent2D extent2D = swapChain.GetSwapChainExtent();
            var RenderArea = VkRect2D.calloc(MemStack).extent(extent2D);

            for(int i = 0; i < ImageCount; i++){
                var renderingInfo = VkRenderingInfo.calloc()
                        .sType$Default()
                        .renderArea(RenderArea)
                        .layerCount(1)
                        .pColorAttachments(ColourAttachments[i]);
                Result[i] = renderingInfo;
            }
        }
        return Result;
    }

    private static ShaderModule[] CreateShaderModules(VulkanContext VkCtx){
        if(EngineConfig.getInstance().RecompileShaders()){
            ShaderCompiler.CompileGLSLShaderOnChange(VERTEX_SHADER_FILE_GLSL, Shaderc.shaderc_glsl_vertex_shader);
            ShaderCompiler.CompileGLSLShaderOnChange(FRAGMENT_SHADER_FILE_GLSL, Shaderc.shaderc_glsl_fragment_shader);
        }
        return new ShaderModule[]{
                new ShaderModule(VkCtx, VK_SHADER_STAGE_VERTEX_BIT, VERTEX_SHADER_FILE_SPV, null),
                new ShaderModule(VkCtx, VK_SHADER_STAGE_FRAGMENT_BIT, FRAGMENT_SHADER_FILE_SPV, null)
        };
    }

    public void Render(VulkanContext VkCtx, CommandBuffer commandBuffer, Attachment SrcAttachment, int ImageIndex){
        try(var MemStack = MemoryStack.stackPush()){
            SwapChain swapChain = VkCtx.GetSwapChain();

            long SwapChainImage = swapChain.GetVulkanImageView(ImageIndex).GetVulkanImage();
            VkCommandBuffer CommandHandle = commandBuffer.GetVulkanCommandBuffer();

            VulkanUtils.ImageBarrier(MemStack,CommandHandle,SwapChainImage,
                    VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL, VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
                    VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT, VK_ACCESS_2_NONE,VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,VK_IMAGE_ASPECT_COLOR_BIT);
            VulkanUtils.ImageBarrier(MemStack,CommandHandle,SrcAttachment.GetVkImage().getVulkanImage(),
                    VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
                    VK_PIPELINE_STAGE_2_FRAGMENT_SHADER_BIT,VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,VK_ACCESS_2_SHADER_READ_BIT,VK_IMAGE_ASPECT_COLOR_BIT);

            vkCmdBeginRendering(CommandHandle, RenderInfo[ImageIndex]);
            vkCmdBindPipeline(CommandHandle,VK_PIPELINE_BIND_POINT_GRAPHICS,pipeline.GetVulkanPipeline());

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

            var Scissor = VkRect2D.calloc(1,MemStack)
                    .extent(it->it.width(Width).height(Height))
                    .offset(it->it.x(0).y(0));
            vkCmdSetScissor(CommandHandle,0,Scissor);

            DescriptorAllocator descriptorAllocator = VkCtx.GetDescriptorAllocator();
            LongBuffer DescriptorSets = MemStack.mallocLong(1).put(0,descriptorAllocator.GetDescriptorSet(DESCRIPTOR_ID_ATTACHMENT).GetVkDescriptorSet());
            vkCmdBindDescriptorSets(CommandHandle,VK_PIPELINE_BIND_POINT_GRAPHICS,pipeline.GetVulkanPipelineLayout(),0,DescriptorSets,null);

            vkCmdDraw(CommandHandle,3,1,0,0);

            vkCmdEndRendering(CommandHandle);

            VulkanUtils.ImageBarrier(MemStack,CommandHandle,SwapChainImage,VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
                    VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT, VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT, VK_ACCESS_2_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,
                    VK_PIPELINE_STAGE_2_NONE, VK_IMAGE_ASPECT_COLOR_BIT);
        }
    }

    public void Resize(VulkanContext VkCtx, Attachment SrcAttachment){
        Arrays.asList(RenderInfo).forEach(VkRenderingInfo::free);
        Arrays.asList(ColourAttachmentsInfo).forEach(VkRenderingAttachmentInfo.Buffer::free);
        ColourAttachmentsInfo = CreateColourAttachmentsInfo(VkCtx,ClearValueColour);
        RenderInfo = CreateRenderInfo(VkCtx,ColourAttachmentsInfo);

        DescriptorAllocator descriptorAllocator = VkCtx.GetDescriptorAllocator();
        DescriptorSet descriptorSet = descriptorAllocator.GetDescriptorSet(DESCRIPTOR_ID_ATTACHMENT);
        descriptorSet.SetImage(VkCtx.GetDevice(),SrcAttachment.GetVkImageView(),textureSampler,0);
    }

    public void CleanUp(VulkanContext VkCtx){
        textureSampler.CleanUp(VkCtx);
        AttachmentDescriptorSetLayout.CleanUp(VkCtx);
        pipeline.CleanUp(VkCtx);
        Arrays.asList(RenderInfo).forEach(VkRenderingInfo::free);
        Arrays.asList(ColourAttachmentsInfo).forEach(VkRenderingAttachmentInfo.Buffer::free);
        ClearValueColour.free();
    }
}
