package net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.PostProcessing;

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.Images.Image;
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.Structure.SwapChain.SwapChain;
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 java.nio.FloatBuffer;
import java.nio.LongBuffer;
import java.util.Arrays;

import static org.lwjgl.vulkan.VK10.*;
import static org.lwjgl.vulkan.VK13.*;

public class PostProcess {
    public static final int COLOUR_FORMAT = VK_FORMAT_R16G16B16A16_SFLOAT;
    private static final String DESCRIPTOR_ID_ATTACHMENT = "POST_DESC_ID_ATT";
    private static final String DESCRIPTOR_ID_SCREEN_SIZE = "POST_DESC_ID_SCREEN_SIZE";
    private static final String MULTI_PASS_FRAGMENT_SHADER_FILE_GLSL = "resources/EngineResources/shaders/multi-sampled_post_process_frag.glsl";
    private static final String MULTI_PASS_FRAGMENT_SHADER_FILE_SPV = MULTI_PASS_FRAGMENT_SHADER_FILE_GLSL + ".spv";
    private static final String SINGLE_PASS_FRAGMENT_SHADER_FILE_GLSL = "resources/EngineResources/shaders/post_process_frag.glsl";
    private static final String SINGLE_PASS_FRAGMENT_SHADER_FILE_SPV = SINGLE_PASS_FRAGMENT_SHADER_FILE_GLSL + ".spv";
    private static final String VERTEX_SHADER_FILE_GLSL = "resources/EngineResources/shaders/post_process_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 DescriptorSetLayout FragmentUniformDescriptorSetLayout;
    private final Pipeline pipeline;
    private final VulkanBuffer ScreenSizeBuffer;
    private final SpecializationConstants specConstants;
    private final TextureSampler textureSampler;
    private Attachment ColourAttachment;
    private VkRenderingAttachmentInfo.Buffer ColourAttachmentInfo;
    private VkRenderingInfo RenderingInfo;

    public PostProcess(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));

        ColourAttachment = CreateColourAttachment(VkCtx);
        ColourAttachmentInfo = CreateColourAttachmentInfo(ColourAttachment,ClearValueColour);
        RenderingInfo = CreateRenderInfo(ColourAttachment,ColourAttachmentInfo);

        var TextureSamplerInfo = new TextureSamplerInfo(VK_SAMPLER_ADDRESS_MODE_REPEAT,VK_BORDER_COLOR_FLOAT_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);

        LayoutInfo = new DescriptorSetLayout.LayoutInformation(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 0,1,VK_SHADER_STAGE_FRAGMENT_BIT);
        FragmentUniformDescriptorSetLayout = new DescriptorSetLayout(VkCtx, LayoutInfo);
        ScreenSizeBuffer = VulkanUtils.CreateHostVisibleBuffer(VkCtx, VulkanUtils.VEC2_SIZE, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,DESCRIPTOR_ID_SCREEN_SIZE,FragmentUniformDescriptorSetLayout);
        SetScreenSizeBuffer(VkCtx);

        specConstants = new SpecializationConstants();
        ShaderModule[] shaderModules = CreateShaderModules(VkCtx, specConstants);

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

    private static Attachment CreateColourAttachment(VulkanContext VkCtx){
        SwapChain swapChain = VkCtx.GetSwapChain();
        VkExtent2D SwapChainExtent = swapChain.GetSwapChainExtent();
        return new Attachment(VkCtx,SwapChainExtent.width(),SwapChainExtent.height(),COLOUR_FORMAT,VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT);
    }

    private static VkRenderingAttachmentInfo.Buffer CreateColourAttachmentInfo(Attachment SrcAttachment, VkClearValue ClearValue){
        return VkRenderingAttachmentInfo.calloc(1)
                .sType$Default()
                .imageView(SrcAttachment.GetVkImageView().GetVulkanImageView())
                .imageLayout(VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL)
                .loadOp(VK_ATTACHMENT_LOAD_OP_CLEAR)
                .storeOp(VK_ATTACHMENT_STORE_OP_STORE)
                .clearValue(ClearValue);
    }

    private static Pipeline CreatePipeline(VulkanContext VkCtx, ShaderModule[] shaderModules, DescriptorSetLayout[] descriptorSetLayouts){
        var VertexBufferStruct = new EmptyVertexBufferStruct();
        var BuildInfo = new PipelineBuildInfo(shaderModules,VertexBufferStruct.GetVertexInput(),COLOUR_FORMAT).SetDescriptorSetLayouts(descriptorSetLayouts).BlendingIsUsed(true);
        var PipeLine = new DefaultPipeline(VkCtx,BuildInfo);
        VertexBufferStruct.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 void CreateAttachmentDescriptorSet(VulkanContext VkCtx, DescriptorSetLayout descriptorSetLayout, Attachment attachment, TextureSampler textureSampler){
        DescriptorAllocator descriptorAllocator = VkCtx.GetDescriptorAllocator();
        Device device = VkCtx.GetDevice();
        DescriptorSet descriptorSet = descriptorAllocator.AddDescriptorSets(device,DESCRIPTOR_ID_ATTACHMENT,1 , descriptorSetLayout)[0];
        descriptorSet.SetImage(device,attachment.GetVkImageView(),textureSampler,0);
    }

    private static ShaderModule[] CreateShaderModules(VulkanContext VkCtx, SpecializationConstants specConstants){
        String ShaderPath = SINGLE_PASS_FRAGMENT_SHADER_FILE_GLSL;
        if(EngineConfig.getInstance().RecompileShaders()){
            ShaderCompiler.CompileGLSLShaderOnChange(VERTEX_SHADER_FILE_GLSL, Shaderc.shaderc_glsl_vertex_shader);
            if(EngineConfig.getInstance().RenderAA() > 1) ShaderPath = MULTI_PASS_FRAGMENT_SHADER_FILE_GLSL;
            ShaderCompiler.CompileGLSLShaderOnChange(ShaderPath, Shaderc.shaderc_glsl_fragment_shader);
        }
        ShaderPath = SINGLE_PASS_FRAGMENT_SHADER_FILE_SPV;
        if(EngineConfig.getInstance().RenderAA() > 1) ShaderPath = MULTI_PASS_FRAGMENT_SHADER_FILE_SPV;
        return new ShaderModule[]{
                new ShaderModule(VkCtx,VK_SHADER_STAGE_VERTEX_BIT,VERTEX_SHADER_FILE_SPV,null),
                new ShaderModule(VkCtx,VK_SHADER_STAGE_FRAGMENT_BIT,ShaderPath,specConstants.GetSpecializationInfo())
        };
    }

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

            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);
            VulkanUtils.ImageBarrier(MemStack,CommandHandle,ColourAttachment.GetVkImage().getVulkanImage(),
                    VK_IMAGE_LAYOUT_UNDEFINED,VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,VK_PIPELINE_STAGE_2_TOP_OF_PIPE_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);
            vkCmdBeginRendering(CommandHandle,RenderingInfo);
            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(2)
                    .put(0,descriptorAllocator.GetDescriptorSet(DESCRIPTOR_ID_ATTACHMENT).GetVkDescriptorSet())
                    .put(1,descriptorAllocator.GetDescriptorSet(DESCRIPTOR_ID_SCREEN_SIZE).GetVkDescriptorSet());
            vkCmdBindDescriptorSets(CommandHandle,VK_PIPELINE_BIND_POINT_GRAPHICS,pipeline.GetVulkanPipelineLayout(),0,DescriptorSets,null);
            vkCmdDraw(CommandHandle,3,1,0,0);
            vkCmdEndRendering(CommandHandle);
        }
    }

    public void Resize(VulkanContext VkCtx, Attachment SrcAttachment){
        RenderingInfo.free();;
        ColourAttachment.CleanUp(VkCtx);
        ColourAttachmentInfo.free();
        ColourAttachment = CreateColourAttachment(VkCtx);
        ColourAttachmentInfo = CreateColourAttachmentInfo(ColourAttachment,ClearValueColour);
        RenderingInfo = CreateRenderInfo(ColourAttachment,ColourAttachmentInfo);

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

        SetScreenSizeBuffer(VkCtx);
    }

    private void SetScreenSizeBuffer(VulkanContext VkCtx){
        long MappedMemory = ScreenSizeBuffer.MapMemory(VkCtx);
        FloatBuffer dataBuffer = MemoryUtil.memFloatBuffer(MappedMemory,(int)ScreenSizeBuffer.GetRequestedSize());
        VkExtent2D SwapChainExtent = VkCtx.GetSwapChain().GetSwapChainExtent();
        dataBuffer.put(0,SwapChainExtent.width());
        dataBuffer.put(1,SwapChainExtent.height());
        ScreenSizeBuffer.UnMapMemory(VkCtx);
    }

    public Attachment GetAttachment(){return ColourAttachment;}

    public void CleanUp(VulkanContext VkCtx){
        ClearValueColour.free();
        ColourAttachment.CleanUp(VkCtx);
        textureSampler.CleanUp(VkCtx);
        AttachmentDescriptorSetLayout.CleanUp(VkCtx);
        FragmentUniformDescriptorSetLayout.CleanUp(VkCtx);
        pipeline.CleanUp(VkCtx);
        RenderingInfo.free();
        ColourAttachmentInfo.free();
        ScreenSizeBuffer.cleanup(VkCtx);
        specConstants.CleanUp();
    }
}
