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

import net.halbear.Terrain4J.EngineCore.Logic.Rendering.VulkanContext;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Shader.DescriptorSetLayout;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Shader.ShaderModule;
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.vulkan.*;
import org.tinylog.Logger;

import java.nio.ByteBuffer;
import java.nio.IntBuffer;
import java.nio.LongBuffer;

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

public class Pipeline {
    private final long VulkanPipeline;
    private final long VulkanPipelineLayout;

    public Pipeline(VulkanContext VkCtx, PipelineBuildInfo BuildInfo){
        Logger.debug("Creating Pipeline");
        Device device = VkCtx.GetDevice();
        try(var MemStack = MemoryStack.stackPush()){
            LongBuffer longPtr = MemStack.mallocLong(1);
            ByteBuffer main = MemStack.UTF8("main");
            ShaderModule[] ShaderModules = BuildInfo.GetShaderModules();
            int ModuleCount = ShaderModules.length;
            var ShaderStages = VkPipelineShaderStageCreateInfo.calloc(ModuleCount, MemStack);
            for(int i = 0; i < ModuleCount; i++){
                ShaderModule shaderModule = ShaderModules[i];
                ShaderStages.get(i)
                        .sType$Default()
                        .stage(shaderModule.GetShaderStage())
                        .module(shaderModule.GetHandle())
                        .pName(main);
            }

            var AssemblyStateCreateInfo = VkPipelineInputAssemblyStateCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .topology(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST);
            var ViewportCreateStateInfo = VkPipelineViewportStateCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .viewportCount(1)
                    .scissorCount(1);
            var RasterizationStateCreateInfo = VkPipelineRasterizationStateCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .polygonMode(VK_POLYGON_MODE_FILL)
                    .cullMode(VK_CULL_MODE_NONE)
                    .frontFace(VK_FRONT_FACE_CLOCKWISE)
                    .lineWidth(1.0f);
            var MultisampleStateCreateInfo = VkPipelineMultisampleStateCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .rasterizationSamples(VK_SAMPLE_COUNT_1_BIT);

            VkPipelineDepthStencilStateCreateInfo DepthStencil = null;
            if(BuildInfo.GetDepthFormat() != VK_FORMAT_UNDEFINED){
                DepthStencil = VkPipelineDepthStencilStateCreateInfo.calloc(MemStack)
                        .sType$Default()
                        .depthTestEnable(true)
                        .depthWriteEnable(true)
                        .depthCompareOp(VK_COMPARE_OP_LESS_OR_EQUAL)
                        .depthBoundsTestEnable(false)
                        .stencilTestEnable(false);
            }


            VkPipelineDynamicStateCreateInfo VulkanPipelineDynamicStateCreateInfo = VkPipelineDynamicStateCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .pDynamicStates(MemStack.ints(
                            VK_DYNAMIC_STATE_VIEWPORT,
                            VK_DYNAMIC_STATE_SCISSOR
                    ));


            VkPushConstantRange.Buffer VkPushConstRangeBuffer = null;
            PushConstantsRange[] PushConstRanges = BuildInfo.GetPushConstantRanges();
            int PushConstCount = PushConstRanges.length != 0 ? PushConstRanges.length : 0;
            if(PushConstCount > 0){
                VkPushConstRangeBuffer = VkPushConstantRange.calloc(PushConstCount,MemStack);
                for(int i = 0; i < PushConstCount; i++){
                    PushConstantsRange pushConstantsRange = PushConstRanges[i];
                    VkPushConstRangeBuffer.get(i)
                            .stageFlags(pushConstantsRange.Stage())
                            .offset(pushConstantsRange.Offset())
                            .size(pushConstantsRange.Size());
                }
            }
            DescriptorSetLayout[] descriptorSetLayouts = BuildInfo.GetDescriptorSetLayouts();
            int LayoutCount = descriptorSetLayouts != null ? descriptorSetLayouts.length : 0;
            LongBuffer ppLayout = MemStack.mallocLong(LayoutCount);
            for(int i = 0; i < LayoutCount; i++){
                ppLayout.put(i, descriptorSetLayouts[i].GetVkDescriptorLayout());
            }
            var BlendAttributeState = VkPipelineColorBlendAttachmentState.calloc(1,MemStack)
                    .colorWriteMask(VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT)
                    .blendEnable(BuildInfo.BlendingUsed());
            if(BuildInfo.BlendingUsed()){
                BlendAttributeState.get(0).colorBlendOp(VK_BLEND_OP_ADD)
                        .alphaBlendOp(VK_BLEND_OP_ADD)
                        .srcColorBlendFactor(VK_BLEND_FACTOR_SRC_ALPHA)
                        .dstColorBlendFactor(VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA)
                        .srcAlphaBlendFactor(VK_BLEND_FACTOR_SRC_ALPHA)
                        .dstAlphaBlendFactor(VK_BLEND_FACTOR_ZERO);
            }


            var ColourBlendState = VkPipelineColorBlendStateCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .pAttachments(BlendAttributeState);
            IntBuffer ColourFormats = MemStack.mallocInt(1);
            ColourFormats.put(0,BuildInfo.GetColourFormat());

            var RendererCreateInfo = VkPipelineRenderingCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .colorAttachmentCount(1)
                    .pColorAttachmentFormats(ColourFormats);

            if(DepthStencil != null){RendererCreateInfo.depthAttachmentFormat(BuildInfo.GetDepthFormat());}

            var PipelineLayoutCreateInfoPtr = VkPipelineLayoutCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .pSetLayouts(ppLayout)
                    .pPushConstantRanges(VkPushConstRangeBuffer);

            VulkanUtils.vkCheck(vkCreatePipelineLayout(device.FetchVulkanDevice(), PipelineLayoutCreateInfoPtr,null,longPtr)
            ,"Unable to create new pipeline layout");
            VulkanPipelineLayout = longPtr.get(0);

            var PipelineCreateInfo = VkGraphicsPipelineCreateInfo.calloc(1,MemStack)
                    .sType$Default()
                    .renderPass(VK_NULL_HANDLE)
                    .pStages(ShaderStages)
                    .pVertexInputState(BuildInfo.GetVertexInputStateCreateInfo())
                    .pInputAssemblyState(AssemblyStateCreateInfo)
                    .pViewportState(ViewportCreateStateInfo)
                    .pRasterizationState(RasterizationStateCreateInfo)
                    .pColorBlendState(ColourBlendState)
                    .pMultisampleState(MultisampleStateCreateInfo)
                    .pDynamicState(VulkanPipelineDynamicStateCreateInfo)
                    .layout(VulkanPipelineLayout)
                    .pNext(RendererCreateInfo);

            if(DepthStencil != null){PipelineCreateInfo.pDepthStencilState(DepthStencil);}

            VulkanUtils.vkCheck(vkCreateGraphicsPipelines(device.FetchVulkanDevice(),
                    VkCtx.GetVkPipelineCache().GetVkPipelineCache(), PipelineCreateInfo,
                    null, longPtr),"Could not create new pipeline");
            VulkanPipeline = longPtr.get(0);
        }
    }
    public long GetVulkanPipeline(){return VulkanPipeline;}
    public long GetVulkanPipelineLayout(){return VulkanPipelineLayout;}
    public void CleanUp(VulkanContext VkCtx){
        Logger.debug("destroying Pipeline");
        VkDevice vkDevice = VkCtx.GetDevice().FetchVulkanDevice();
        vkDestroyPipelineLayout(vkDevice,VulkanPipelineLayout,null);
        vkDestroyPipeline(vkDevice,VulkanPipeline,null);
    }

}
